graph.cc 20.9 KB
Newer Older
1
2
3
4
5
6
/*!
 *  Copyright (c) 2018 by Contributors
 * \file graph/graph.cc
 * \brief DGL graph index implementation
 */
#include <dgl/graph.h>
Minjie Wang's avatar
impl  
Minjie Wang committed
7
#include <algorithm>
Minjie Wang's avatar
Minjie Wang committed
8
#include <unordered_map>
9
10
#include <set>
#include <functional>
GaiYu0's avatar
GaiYu0 committed
11
12
#include <tuple>
#include "../c_api_common.h"
Minjie Wang's avatar
Minjie Wang committed
13

Minjie Wang's avatar
impl  
Minjie Wang committed
14
namespace dgl {
15
16
17
18
19
20
21
22

Graph::Graph(IdArray src_ids, IdArray dst_ids, IdArray edge_ids, size_t num_nodes,
    bool multigraph): is_multigraph_(multigraph) {
  CHECK(IsValidIdArray(src_ids));
  CHECK(IsValidIdArray(dst_ids));
  CHECK(IsValidIdArray(edge_ids));
  this->AddVertices(num_nodes);
  num_edges_ = src_ids->shape[0];
23
24
25
26
  CHECK(static_cast<int64_t>(num_edges_) == dst_ids->shape[0])
    << "vectors in COO must have the same length";
  CHECK(static_cast<int64_t>(num_edges_) == edge_ids->shape[0])
    << "vectors in COO must have the same length";
27
28
29
30
31
  const dgl_id_t *src_data = static_cast<dgl_id_t*>(src_ids->data);
  const dgl_id_t *dst_data = static_cast<dgl_id_t*>(dst_ids->data);
  const dgl_id_t *edge_data = static_cast<dgl_id_t*>(edge_ids->data);
  all_edges_src_.reserve(num_edges_);
  all_edges_dst_.reserve(num_edges_);
32
  for (uint64_t i = 0; i < num_edges_; i++) {
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
    auto src = src_data[i];
    auto dst = dst_data[i];
    auto eid = edge_data[i];
    CHECK(HasVertex(src) && HasVertex(dst))
      << "Invalid vertices: src=" << src << " dst=" << dst;

    adjlist_[src].succ.push_back(dst);
    adjlist_[src].edge_id.push_back(eid);
    reverse_adjlist_[dst].succ.push_back(src);
    reverse_adjlist_[dst].edge_id.push_back(eid);

    all_edges_src_.push_back(src);
    all_edges_dst_.push_back(dst);
  }
}

Minjie Wang's avatar
impl  
Minjie Wang committed
49
50
51
void Graph::AddVertices(uint64_t num_vertices) {
  CHECK(!read_only_) << "Graph is read-only. Mutations are not allowed.";
  adjlist_.resize(adjlist_.size() + num_vertices);
Minjie Wang's avatar
Minjie Wang committed
52
  reverse_adjlist_.resize(reverse_adjlist_.size() + num_vertices);
Minjie Wang's avatar
impl  
Minjie Wang committed
53
54
55
56
57
}

void Graph::AddEdge(dgl_id_t src, dgl_id_t dst) {
  CHECK(!read_only_) << "Graph is read-only. Mutations are not allowed.";
  CHECK(HasVertex(src) && HasVertex(dst))
58
    << "Invalid vertices: src=" << src << " dst=" << dst;
59

Minjie Wang's avatar
impl  
Minjie Wang committed
60
  dgl_id_t eid = num_edges_++;
61

Minjie Wang's avatar
impl  
Minjie Wang committed
62
  adjlist_[src].succ.push_back(dst);
Minjie Wang's avatar
Minjie Wang committed
63
64
65
  adjlist_[src].edge_id.push_back(eid);
  reverse_adjlist_[dst].succ.push_back(src);
  reverse_adjlist_[dst].edge_id.push_back(eid);
66

Minjie Wang's avatar
Minjie Wang committed
67
68
  all_edges_src_.push_back(src);
  all_edges_dst_.push_back(dst);
Minjie Wang's avatar
impl  
Minjie Wang committed
69
70
71
72
73
74
75
}

void Graph::AddEdges(IdArray src_ids, IdArray dst_ids) {
  CHECK(!read_only_) << "Graph is read-only. Mutations are not allowed.";
  CHECK(IsValidIdArray(src_ids)) << "Invalid src id array.";
  CHECK(IsValidIdArray(dst_ids)) << "Invalid dst id array.";
  const auto srclen = src_ids->shape[0];
Minjie Wang's avatar
Minjie Wang committed
76
  const auto dstlen = dst_ids->shape[0];
Minjie Wang's avatar
impl  
Minjie Wang committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
  const int64_t* src_data = static_cast<int64_t*>(src_ids->data);
  const int64_t* dst_data = static_cast<int64_t*>(dst_ids->data);
  if (srclen == 1) {
    // one-many
    for (int64_t i = 0; i < dstlen; ++i) {
      AddEdge(src_data[0], dst_data[i]);
    }
  } else if (dstlen == 1) {
    // many-one
    for (int64_t i = 0; i < srclen; ++i) {
      AddEdge(src_data[i], dst_data[0]);
    }
  } else {
    // many-many
    CHECK(srclen == dstlen) << "Invalid src and dst id array.";
    for (int64_t i = 0; i < srclen; ++i) {
      AddEdge(src_data[i], dst_data[i]);
    }
  }
}

BoolArray Graph::HasVertices(IdArray vids) const {
  CHECK(IsValidIdArray(vids)) << "Invalid vertex id array.";
  const auto len = vids->shape[0];
  BoolArray rst = BoolArray::Empty({len}, vids->dtype, vids->ctx);
  const int64_t* vid_data = static_cast<int64_t*>(vids->data);
  int64_t* rst_data = static_cast<int64_t*>(rst->data);
Minjie Wang's avatar
Minjie Wang committed
104
  const int64_t nverts = NumVertices();
Minjie Wang's avatar
impl  
Minjie Wang committed
105
106
107
108
109
110
111
  for (int64_t i = 0; i < len; ++i) {
    rst_data[i] = (vid_data[i] < nverts)? 1 : 0;
  }
  return rst;
}

// O(E)
112
bool Graph::HasEdgeBetween(dgl_id_t src, dgl_id_t dst) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
113
114
115
116
117
  if (!HasVertex(src) || !HasVertex(dst)) return false;
  const auto& succ = adjlist_[src].succ;
  return std::find(succ.begin(), succ.end(), dst) != succ.end();
}

118
119
// O(E*k) pretty slow
BoolArray Graph::HasEdgesBetween(IdArray src_ids, IdArray dst_ids) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
120
121
122
  CHECK(IsValidIdArray(src_ids)) << "Invalid src id array.";
  CHECK(IsValidIdArray(dst_ids)) << "Invalid dst id array.";
  const auto srclen = src_ids->shape[0];
Minjie Wang's avatar
Minjie Wang committed
123
  const auto dstlen = dst_ids->shape[0];
Minjie Wang's avatar
impl  
Minjie Wang committed
124
125
126
127
128
129
130
131
  const auto rstlen = std::max(srclen, dstlen);
  BoolArray rst = BoolArray::Empty({rstlen}, src_ids->dtype, src_ids->ctx);
  int64_t* rst_data = static_cast<int64_t*>(rst->data);
  const int64_t* src_data = static_cast<int64_t*>(src_ids->data);
  const int64_t* dst_data = static_cast<int64_t*>(dst_ids->data);
  if (srclen == 1) {
    // one-many
    for (int64_t i = 0; i < dstlen; ++i) {
132
      rst_data[i] = HasEdgeBetween(src_data[0], dst_data[i])? 1 : 0;
Minjie Wang's avatar
impl  
Minjie Wang committed
133
134
135
136
    }
  } else if (dstlen == 1) {
    // many-one
    for (int64_t i = 0; i < srclen; ++i) {
137
      rst_data[i] = HasEdgeBetween(src_data[i], dst_data[0])? 1 : 0;
Minjie Wang's avatar
impl  
Minjie Wang committed
138
139
140
141
142
    }
  } else {
    // many-many
    CHECK(srclen == dstlen) << "Invalid src and dst id array.";
    for (int64_t i = 0; i < srclen; ++i) {
143
      rst_data[i] = HasEdgeBetween(src_data[i], dst_data[i])? 1 : 0;
Minjie Wang's avatar
impl  
Minjie Wang committed
144
145
146
147
148
149
    }
  }
  return rst;
}

// The data is copy-out; support zero-copy?
Minjie Wang's avatar
Minjie Wang committed
150
IdArray Graph::Predecessors(dgl_id_t vid, uint64_t radius) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
151
  CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
Minjie Wang's avatar
Minjie Wang committed
152
  CHECK(radius >= 1) << "invalid radius: " << radius;
153
154
155
156
157
158
  std::set<dgl_id_t> vset;

  for (auto& it : reverse_adjlist_[vid].succ)
    vset.insert(it);

  const int64_t len = vset.size();
Minjie Wang's avatar
impl  
Minjie Wang committed
159
160
  IdArray rst = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  int64_t* rst_data = static_cast<int64_t*>(rst->data);
161
162

  std::copy(vset.begin(), vset.end(), rst_data);
Minjie Wang's avatar
impl  
Minjie Wang committed
163
164
165
166
  return rst;
}

// The data is copy-out; support zero-copy?
Minjie Wang's avatar
Minjie Wang committed
167
IdArray Graph::Successors(dgl_id_t vid, uint64_t radius) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
168
  CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
Minjie Wang's avatar
Minjie Wang committed
169
  CHECK(radius >= 1) << "invalid radius: " << radius;
170
171
172
173
174
175
  std::set<dgl_id_t> vset;

  for (auto& it : adjlist_[vid].succ)
    vset.insert(it);

  const int64_t len = vset.size();
Minjie Wang's avatar
impl  
Minjie Wang committed
176
177
  IdArray rst = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  int64_t* rst_data = static_cast<int64_t*>(rst->data);
178
179

  std::copy(vset.begin(), vset.end(), rst_data);
Minjie Wang's avatar
impl  
Minjie Wang committed
180
181
182
183
  return rst;
}

// O(E)
184
185
186
IdArray Graph::EdgeId(dgl_id_t src, dgl_id_t dst) const {
  CHECK(HasVertex(src) && HasVertex(dst)) << "invalid edge: " << src << " -> " << dst;

Minjie Wang's avatar
impl  
Minjie Wang committed
187
  const auto& succ = adjlist_[src].succ;
188
189
  std::vector<dgl_id_t> edgelist;

Minjie Wang's avatar
impl  
Minjie Wang committed
190
  for (size_t i = 0; i < succ.size(); ++i) {
191
192
    if (succ[i] == dst)
      edgelist.push_back(adjlist_[src].edge_id[i]);
Minjie Wang's avatar
impl  
Minjie Wang committed
193
  }
194
195
196
197
198
199
200
201
202
203

  // FIXME: signed?  Also it seems that we are using int64_t everywhere...
  const int64_t len = edgelist.size();
  IdArray rst = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  // FIXME: signed?
  int64_t* rst_data = static_cast<int64_t*>(rst->data);

  std::copy(edgelist.begin(), edgelist.end(), rst_data);

  return rst;
Minjie Wang's avatar
impl  
Minjie Wang committed
204
205
206
}

// O(E*k) pretty slow
207
Graph::EdgeArray Graph::EdgeIds(IdArray src_ids, IdArray dst_ids) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
208
209
210
  CHECK(IsValidIdArray(src_ids)) << "Invalid src id array.";
  CHECK(IsValidIdArray(dst_ids)) << "Invalid dst id array.";
  const auto srclen = src_ids->shape[0];
Minjie Wang's avatar
Minjie Wang committed
211
  const auto dstlen = dst_ids->shape[0];
212
213
214
215
216
217
218
  int64_t i, j;

  CHECK((srclen == dstlen) || (srclen == 1) || (dstlen == 1))
    << "Invalid src and dst id array.";

  const int64_t src_stride = (srclen == 1 && dstlen != 1) ? 0 : 1;
  const int64_t dst_stride = (dstlen == 1 && srclen != 1) ? 0 : 1;
Minjie Wang's avatar
impl  
Minjie Wang committed
219
220
  const int64_t* src_data = static_cast<int64_t*>(src_ids->data);
  const int64_t* dst_data = static_cast<int64_t*>(dst_ids->data);
221
222
223
224
225

  std::vector<dgl_id_t> src, dst, eid;

  for (i = 0, j = 0; i < srclen && j < dstlen; i += src_stride, j += dst_stride) {
    const dgl_id_t src_id = src_data[i], dst_id = dst_data[j];
226
227
    CHECK(HasVertex(src_id) && HasVertex(dst_id)) <<
        "invalid edge: " << src_id << " -> " << dst_id;
228
229
230
    const auto& succ = adjlist_[src_id].succ;
    for (size_t k = 0; k < succ.size(); ++k) {
      if (succ[k] == dst_id) {
231
232
233
        src.push_back(src_id);
        dst.push_back(dst_id);
        eid.push_back(adjlist_[src_id].edge_id[k]);
234
      }
Minjie Wang's avatar
impl  
Minjie Wang committed
235
236
    }
  }
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274

  int64_t rstlen = src.size();
  IdArray rst_src = IdArray::Empty({rstlen}, src_ids->dtype, src_ids->ctx);
  IdArray rst_dst = IdArray::Empty({rstlen}, src_ids->dtype, src_ids->ctx);
  IdArray rst_eid = IdArray::Empty({rstlen}, src_ids->dtype, src_ids->ctx);
  int64_t* rst_src_data = static_cast<int64_t*>(rst_src->data);
  int64_t* rst_dst_data = static_cast<int64_t*>(rst_dst->data);
  int64_t* rst_eid_data = static_cast<int64_t*>(rst_eid->data);

  std::copy(src.begin(), src.end(), rst_src_data);
  std::copy(dst.begin(), dst.end(), rst_dst_data);
  std::copy(eid.begin(), eid.end(), rst_eid_data);

  return EdgeArray{rst_src, rst_dst, rst_eid};
}

Graph::EdgeArray Graph::FindEdges(IdArray eids) const {
  int64_t len = eids->shape[0];

  IdArray rst_src = IdArray::Empty({len}, eids->dtype, eids->ctx);
  IdArray rst_dst = IdArray::Empty({len}, eids->dtype, eids->ctx);
  IdArray rst_eid = IdArray::Empty({len}, eids->dtype, eids->ctx);
  int64_t* eid_data = static_cast<int64_t*>(eids->data);
  int64_t* rst_src_data = static_cast<int64_t*>(rst_src->data);
  int64_t* rst_dst_data = static_cast<int64_t*>(rst_dst->data);
  int64_t* rst_eid_data = static_cast<int64_t*>(rst_eid->data);

  for (uint64_t i = 0; i < (uint64_t)len; ++i) {
    dgl_id_t eid = eid_data[i];
    if (eid >= num_edges_)
      LOG(FATAL) << "invalid edge id:" << eid;

    rst_src_data[i] = all_edges_src_[eid];
    rst_dst_data[i] = all_edges_dst_[eid];
    rst_eid_data[i] = eid;
  }

  return EdgeArray{rst_src, rst_dst, rst_eid};
Minjie Wang's avatar
impl  
Minjie Wang committed
275
276
277
}

// O(E)
278
279
Graph::EdgeArray Graph::InEdges(dgl_id_t vid) const {
  CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
Minjie Wang's avatar
Minjie Wang committed
280
  const int64_t len = reverse_adjlist_[vid].succ.size();
281
282
283
284
  IdArray src = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  IdArray dst = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  IdArray eid = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  int64_t* src_data = static_cast<int64_t*>(src->data);
Minjie Wang's avatar
impl  
Minjie Wang committed
285
  int64_t* dst_data = static_cast<int64_t*>(dst->data);
286
287
  int64_t* eid_data = static_cast<int64_t*>(eid->data);
  for (int64_t i = 0; i < len; ++i) {
Minjie Wang's avatar
Minjie Wang committed
288
289
    src_data[i] = reverse_adjlist_[vid].succ[i];
    eid_data[i] = reverse_adjlist_[vid].edge_id[i];
290
291
292
  }
  std::fill(dst_data, dst_data + len, vid);
  return EdgeArray{src, dst, eid};
Minjie Wang's avatar
impl  
Minjie Wang committed
293
294
295
}

// O(E)
296
Graph::EdgeArray Graph::InEdges(IdArray vids) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
297
298
299
300
301
302
  CHECK(IsValidIdArray(vids)) << "Invalid vertex id array.";
  const auto len = vids->shape[0];
  const int64_t* vid_data = static_cast<int64_t*>(vids->data);
  int64_t rstlen = 0;
  for (int64_t i = 0; i < len; ++i) {
    CHECK(HasVertex(vid_data[i])) << "Invalid vertex: " << vid_data[i];
Minjie Wang's avatar
Minjie Wang committed
303
    rstlen += reverse_adjlist_[vid_data[i]].succ.size();
Minjie Wang's avatar
impl  
Minjie Wang committed
304
305
306
  }
  IdArray src = IdArray::Empty({rstlen}, vids->dtype, vids->ctx);
  IdArray dst = IdArray::Empty({rstlen}, vids->dtype, vids->ctx);
307
  IdArray eid = IdArray::Empty({rstlen}, vids->dtype, vids->ctx);
Minjie Wang's avatar
impl  
Minjie Wang committed
308
309
  int64_t* src_ptr = static_cast<int64_t*>(src->data);
  int64_t* dst_ptr = static_cast<int64_t*>(dst->data);
310
  int64_t* eid_ptr = static_cast<int64_t*>(eid->data);
Minjie Wang's avatar
impl  
Minjie Wang committed
311
  for (int64_t i = 0; i < len; ++i) {
Minjie Wang's avatar
Minjie Wang committed
312
313
    const auto& pred = reverse_adjlist_[vid_data[i]].succ;
    const auto& eids = reverse_adjlist_[vid_data[i]].edge_id;
Minjie Wang's avatar
impl  
Minjie Wang committed
314
315
316
    for (size_t j = 0; j < pred.size(); ++j) {
      *(src_ptr++) = pred[j];
      *(dst_ptr++) = vid_data[i];
317
      *(eid_ptr++) = eids[j];
Minjie Wang's avatar
impl  
Minjie Wang committed
318
319
    }
  }
320
  return EdgeArray{src, dst, eid};
Minjie Wang's avatar
impl  
Minjie Wang committed
321
322
323
}

// O(E)
324
325
326
327
328
329
Graph::EdgeArray Graph::OutEdges(dgl_id_t vid) const {
  CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
  const int64_t len = adjlist_[vid].succ.size();
  IdArray src = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  IdArray dst = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  IdArray eid = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
Minjie Wang's avatar
impl  
Minjie Wang committed
330
  int64_t* src_data = static_cast<int64_t*>(src->data);
331
332
333
334
  int64_t* dst_data = static_cast<int64_t*>(dst->data);
  int64_t* eid_data = static_cast<int64_t*>(eid->data);
  for (int64_t i = 0; i < len; ++i) {
    dst_data[i] = adjlist_[vid].succ[i];
Minjie Wang's avatar
Minjie Wang committed
335
    eid_data[i] = adjlist_[vid].edge_id[i];
336
337
338
  }
  std::fill(src_data, src_data + len, vid);
  return EdgeArray{src, dst, eid};
Minjie Wang's avatar
impl  
Minjie Wang committed
339
340
341
}

// O(E)
342
Graph::EdgeArray Graph::OutEdges(IdArray vids) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
343
344
345
346
347
348
349
350
351
352
  CHECK(IsValidIdArray(vids)) << "Invalid vertex id array.";
  const auto len = vids->shape[0];
  const int64_t* vid_data = static_cast<int64_t*>(vids->data);
  int64_t rstlen = 0;
  for (int64_t i = 0; i < len; ++i) {
    CHECK(HasVertex(vid_data[i])) << "Invalid vertex: " << vid_data[i];
    rstlen += adjlist_[vid_data[i]].succ.size();
  }
  IdArray src = IdArray::Empty({rstlen}, vids->dtype, vids->ctx);
  IdArray dst = IdArray::Empty({rstlen}, vids->dtype, vids->ctx);
353
  IdArray eid = IdArray::Empty({rstlen}, vids->dtype, vids->ctx);
Minjie Wang's avatar
impl  
Minjie Wang committed
354
355
  int64_t* src_ptr = static_cast<int64_t*>(src->data);
  int64_t* dst_ptr = static_cast<int64_t*>(dst->data);
356
  int64_t* eid_ptr = static_cast<int64_t*>(eid->data);
Minjie Wang's avatar
impl  
Minjie Wang committed
357
358
  for (int64_t i = 0; i < len; ++i) {
    const auto& succ = adjlist_[vid_data[i]].succ;
Minjie Wang's avatar
Minjie Wang committed
359
    const auto& eids = adjlist_[vid_data[i]].edge_id;
Minjie Wang's avatar
impl  
Minjie Wang committed
360
361
362
    for (size_t j = 0; j < succ.size(); ++j) {
      *(src_ptr++) = vid_data[i];
      *(dst_ptr++) = succ[j];
363
      *(eid_ptr++) = eids[j];
Minjie Wang's avatar
impl  
Minjie Wang committed
364
365
    }
  }
366
  return EdgeArray{src, dst, eid};
Minjie Wang's avatar
impl  
Minjie Wang committed
367
368
}

Minjie Wang's avatar
Minjie Wang committed
369
// O(E*log(E)) if sort is required; otherwise, O(E)
370
Graph::EdgeArray Graph::Edges(const std::string &order) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
371
372
373
  const int64_t len = num_edges_;
  IdArray src = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  IdArray dst = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
374
  IdArray eid = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
Minjie Wang's avatar
Minjie Wang committed
375

376
  if (order == "srcdst") {
Minjie Wang's avatar
Minjie Wang committed
377
378
379
    typedef std::tuple<int64_t, int64_t, int64_t> Tuple;
    std::vector<Tuple> tuples;
    tuples.reserve(len);
Minjie Wang's avatar
Minjie Wang committed
380
381
    for (uint64_t eid = 0; eid < num_edges_; ++eid) {
      tuples.emplace_back(all_edges_src_[eid], all_edges_dst_[eid], eid);
Minjie Wang's avatar
Minjie Wang committed
382
    }
Minjie Wang's avatar
Minjie Wang committed
383
    // sort according to src and dst ids
Minjie Wang's avatar
Minjie Wang committed
384
385
    std::sort(tuples.begin(), tuples.end(),
        [] (const Tuple& t1, const Tuple& t2) {
Minjie Wang's avatar
Minjie Wang committed
386
387
          return std::get<0>(t1) < std::get<0>(t2)
            || (std::get<0>(t1) == std::get<0>(t2) && std::get<1>(t1) < std::get<1>(t2));
Minjie Wang's avatar
Minjie Wang committed
388
        });
389

Minjie Wang's avatar
Minjie Wang committed
390
391
392
    // make return arrays
    int64_t* src_ptr = static_cast<int64_t*>(src->data);
    int64_t* dst_ptr = static_cast<int64_t*>(dst->data);
393
    int64_t* eid_ptr = static_cast<int64_t*>(eid->data);
Minjie Wang's avatar
Minjie Wang committed
394
    for (size_t i = 0; i < tuples.size(); ++i) {
Minjie Wang's avatar
Minjie Wang committed
395
396
      src_ptr[i] = std::get<0>(tuples[i]);
      dst_ptr[i] = std::get<1>(tuples[i]);
397
      eid_ptr[i] = std::get<2>(tuples[i]);
Minjie Wang's avatar
Minjie Wang committed
398
399
400
401
    }
  } else {
    int64_t* src_ptr = static_cast<int64_t*>(src->data);
    int64_t* dst_ptr = static_cast<int64_t*>(dst->data);
402
    int64_t* eid_ptr = static_cast<int64_t*>(eid->data);
Minjie Wang's avatar
Minjie Wang committed
403
404
405
406
    std::copy(all_edges_src_.begin(), all_edges_src_.end(), src_ptr);
    std::copy(all_edges_dst_.begin(), all_edges_dst_.end(), dst_ptr);
    for (uint64_t eid = 0; eid < num_edges_; ++eid) {
      eid_ptr[eid] = eid;
Minjie Wang's avatar
Minjie Wang committed
407
    }
Minjie Wang's avatar
impl  
Minjie Wang committed
408
409
  }

410
  return EdgeArray{src, dst, eid};
Minjie Wang's avatar
impl  
Minjie Wang committed
411
412
413
414
415
416
417
418
419
420
421
422
}

// O(V)
DegreeArray Graph::InDegrees(IdArray vids) const {
  CHECK(IsValidIdArray(vids)) << "Invalid vertex id array.";
  const auto len = vids->shape[0];
  const int64_t* vid_data = static_cast<int64_t*>(vids->data);
  DegreeArray rst = DegreeArray::Empty({len}, vids->dtype, vids->ctx);
  int64_t* rst_data = static_cast<int64_t*>(rst->data);
  for (int64_t i = 0; i < len; ++i) {
    const auto vid = vid_data[i];
    CHECK(HasVertex(vid)) << "Invalid vertex: " << vid;
Minjie Wang's avatar
Minjie Wang committed
423
    rst_data[i] = reverse_adjlist_[vid].succ.size();
Minjie Wang's avatar
impl  
Minjie Wang committed
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
  }
  return rst;
}

// O(V)
DegreeArray Graph::OutDegrees(IdArray vids) const {
  CHECK(IsValidIdArray(vids)) << "Invalid vertex id array.";
  const auto len = vids->shape[0];
  const int64_t* vid_data = static_cast<int64_t*>(vids->data);
  DegreeArray rst = DegreeArray::Empty({len}, vids->dtype, vids->ctx);
  int64_t* rst_data = static_cast<int64_t*>(rst->data);
  for (int64_t i = 0; i < len; ++i) {
    const auto vid = vid_data[i];
    CHECK(HasVertex(vid)) << "Invalid vertex: " << vid;
    rst_data[i] = adjlist_[vid].succ.size();
  }
  return rst;
}

Minjie Wang's avatar
Minjie Wang committed
443
444
445
446
447
448
449
450
451
452
Subgraph Graph::VertexSubgraph(IdArray vids) const {
  CHECK(IsValidIdArray(vids)) << "Invalid vertex id array.";
  const auto len = vids->shape[0];
  std::unordered_map<dgl_id_t, dgl_id_t> oldv2newv;
  std::vector<dgl_id_t> edges;
  const int64_t* vid_data = static_cast<int64_t*>(vids->data);
  for (int64_t i = 0; i < len; ++i) {
    oldv2newv[vid_data[i]] = i;
  }
  Subgraph rst;
453
  rst.graph = std::make_shared<Graph>(IsMultigraph());
Minjie Wang's avatar
Minjie Wang committed
454
  rst.induced_vertices = vids;
455
  rst.graph->AddVertices(len);
Minjie Wang's avatar
Minjie Wang committed
456
457
458
459
460
461
462
463
  for (int64_t i = 0; i < len; ++i) {
    const dgl_id_t oldvid = vid_data[i];
    const dgl_id_t newvid = i;
    for (size_t j = 0; j < adjlist_[oldvid].succ.size(); ++j) {
      const dgl_id_t oldsucc = adjlist_[oldvid].succ[j];
      if (oldv2newv.count(oldsucc)) {
        const dgl_id_t newsucc = oldv2newv[oldsucc];
        edges.push_back(adjlist_[oldvid].edge_id[j]);
464
        rst.graph->AddEdge(newvid, newsucc);
Minjie Wang's avatar
Minjie Wang committed
465
466
467
468
469
470
      }
    }
  }
  rst.induced_edges = IdArray::Empty({static_cast<int64_t>(edges.size())}, vids->dtype, vids->ctx);
  std::copy(edges.begin(), edges.end(), static_cast<int64_t*>(rst.induced_edges->data));
  return rst;
Minjie Wang's avatar
impl  
Minjie Wang committed
471
}
Minjie Wang's avatar
Minjie Wang committed
472

473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
Subgraph Graph::EdgeSubgraph(IdArray eids) const {
  CHECK(IsValidIdArray(eids)) << "Invalid vertex id array.";

  const auto len = eids->shape[0];
  std::unordered_map<dgl_id_t, dgl_id_t> oldv2newv;
  std::vector<dgl_id_t> nodes;
  const int64_t* eid_data = static_cast<int64_t*>(eids->data);

  for (int64_t i = 0; i < len; ++i) {
    dgl_id_t src_id = all_edges_src_[eid_data[i]];
    dgl_id_t dst_id = all_edges_dst_[eid_data[i]];
    if (oldv2newv.insert(std::make_pair(src_id, oldv2newv.size())).second)
      nodes.push_back(src_id);
    if (oldv2newv.insert(std::make_pair(dst_id, oldv2newv.size())).second)
      nodes.push_back(dst_id);
  }

  Subgraph rst;
491
  rst.graph = std::make_shared<Graph>(IsMultigraph());
492
  rst.induced_edges = eids;
493
  rst.graph->AddVertices(nodes.size());
494
495
496
497

  for (int64_t i = 0; i < len; ++i) {
    dgl_id_t src_id = all_edges_src_[eid_data[i]];
    dgl_id_t dst_id = all_edges_dst_[eid_data[i]];
498
    rst.graph->AddEdge(oldv2newv[src_id], oldv2newv[dst_id]);
499
500
  }

501
502
  rst.induced_vertices = IdArray::Empty(
      {static_cast<int64_t>(nodes.size())}, eids->dtype, eids->ctx);
503
504
505
  std::copy(nodes.begin(), nodes.end(), static_cast<int64_t*>(rst.induced_vertices->data));

  return rst;
Minjie Wang's avatar
Minjie Wang committed
506
507
}

508
std::vector<IdArray> Graph::GetAdj(bool transpose, const std::string &fmt) const {
509
510
  uint64_t num_edges = NumEdges();
  uint64_t num_nodes = NumVertices();
511
  if (fmt == "coo") {
512
513
514
515
    IdArray idx = IdArray::Empty(
        {2 * static_cast<int64_t>(num_edges)},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0});
516
517
518
519
520
521
522
523
    int64_t *idx_data = static_cast<int64_t*>(idx->data);
    if (transpose) {
      std::copy(all_edges_src_.begin(), all_edges_src_.end(), idx_data);
      std::copy(all_edges_dst_.begin(), all_edges_dst_.end(), idx_data + num_edges);
    } else {
      std::copy(all_edges_dst_.begin(), all_edges_dst_.end(), idx_data);
      std::copy(all_edges_src_.begin(), all_edges_src_.end(), idx_data + num_edges);
    }
524
525
526
527
    IdArray eid = IdArray::Empty(
        {static_cast<int64_t>(num_edges)},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0});
528
529
530
531
532
533
    int64_t *eid_data = static_cast<int64_t*>(eid->data);
    for (uint64_t eid = 0; eid < num_edges; ++eid) {
      eid_data[eid] = eid;
    }
    return std::vector<IdArray>{idx, eid};
  } else if (fmt == "csr") {
534
535
536
537
538
539
540
541
542
543
544
545
    IdArray indptr = IdArray::Empty(
        {static_cast<int64_t>(num_nodes) + 1},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0});
    IdArray indices = IdArray::Empty(
        {static_cast<int64_t>(num_edges)},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0});
    IdArray eid = IdArray::Empty(
        {static_cast<int64_t>(num_edges)},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0});
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
    int64_t *indptr_data = static_cast<int64_t*>(indptr->data);
    int64_t *indices_data = static_cast<int64_t*>(indices->data);
    int64_t *eid_data = static_cast<int64_t*>(eid->data);
    const AdjacencyList *adjlist;
    if (transpose) {
      // Out-edges.
      adjlist = &adjlist_;
    } else {
      // In-edges.
      adjlist = &reverse_adjlist_;
    }
    indptr_data[0] = 0;
    for (size_t i = 0; i < adjlist->size(); i++) {
      indptr_data[i + 1] = indptr_data[i] + adjlist->at(i).succ.size();
      std::copy(adjlist->at(i).succ.begin(), adjlist->at(i).succ.end(),
                indices_data + indptr_data[i]);
      std::copy(adjlist->at(i).edge_id.begin(), adjlist->at(i).edge_id.end(),
                eid_data + indptr_data[i]);
    }
    return std::vector<IdArray>{indptr, indices, eid};
  } else {
    LOG(FATAL) << "unsupported format";
    return std::vector<IdArray>();
  }
}

GraphPtr Graph::Reverse() const {
Minjie Wang's avatar
impl  
Minjie Wang committed
573
  LOG(FATAL) << "not implemented";
574
  return nullptr;
Minjie Wang's avatar
Minjie Wang committed
575
576
}

Minjie Wang's avatar
impl  
Minjie Wang committed
577
}  // namespace dgl