graph.cc 20.6 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
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

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];
  CHECK(num_edges_ == dst_ids->shape[0]) << "vectors in COO must have the same length";
  CHECK(num_edges_ == edge_ids->shape[0]) << "vectors in COO must have the same length";
  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_);
  for (int64_t i = 0; i < num_edges_; i++) {
    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
47
48
49
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
50
  reverse_adjlist_.resize(reverse_adjlist_.size() + num_vertices);
Minjie Wang's avatar
impl  
Minjie Wang committed
51
52
53
54
55
}

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))
56
    << "Invalid vertices: src=" << src << " dst=" << dst;
57

Minjie Wang's avatar
impl  
Minjie Wang committed
58
  dgl_id_t eid = num_edges_++;
59

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

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

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
74
  const auto dstlen = dst_ids->shape[0];
Minjie Wang's avatar
impl  
Minjie Wang committed
75
76
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
  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
102
  const int64_t nverts = NumVertices();
Minjie Wang's avatar
impl  
Minjie Wang committed
103
104
105
106
107
108
109
  for (int64_t i = 0; i < len; ++i) {
    rst_data[i] = (vid_data[i] < nverts)? 1 : 0;
  }
  return rst;
}

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

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

// The data is copy-out; support zero-copy?
Minjie Wang's avatar
Minjie Wang committed
148
IdArray Graph::Predecessors(dgl_id_t vid, uint64_t radius) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
149
  CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
Minjie Wang's avatar
Minjie Wang committed
150
  CHECK(radius >= 1) << "invalid radius: " << radius;
151
152
153
154
155
156
  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
157
158
  IdArray rst = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  int64_t* rst_data = static_cast<int64_t*>(rst->data);
159
160

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

// The data is copy-out; support zero-copy?
Minjie Wang's avatar
Minjie Wang committed
165
IdArray Graph::Successors(dgl_id_t vid, uint64_t radius) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
166
  CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
Minjie Wang's avatar
Minjie Wang committed
167
  CHECK(radius >= 1) << "invalid radius: " << radius;
168
169
170
171
172
173
  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
174
175
  IdArray rst = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  int64_t* rst_data = static_cast<int64_t*>(rst->data);
176
177

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

// O(E)
182
183
184
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
185
  const auto& succ = adjlist_[src].succ;
186
187
  std::vector<dgl_id_t> edgelist;

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

  // 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
202
203
204
}

// O(E*k) pretty slow
205
Graph::EdgeArray Graph::EdgeIds(IdArray src_ids, IdArray dst_ids) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
206
207
208
  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
209
  const auto dstlen = dst_ids->shape[0];
210
211
212
213
214
215
216
  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
217
218
  const int64_t* src_data = static_cast<int64_t*>(src_ids->data);
  const int64_t* dst_data = static_cast<int64_t*>(dst_ids->data);
219
220
221
222
223

  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];
224
225
    CHECK(HasVertex(src_id) && HasVertex(dst_id)) <<
        "invalid edge: " << src_id << " -> " << dst_id;
226
227
228
    const auto& succ = adjlist_[src_id].succ;
    for (size_t k = 0; k < succ.size(); ++k) {
      if (succ[k] == dst_id) {
229
230
231
        src.push_back(src_id);
        dst.push_back(dst_id);
        eid.push_back(adjlist_[src_id].edge_id[k]);
232
      }
Minjie Wang's avatar
impl  
Minjie Wang committed
233
234
    }
  }
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

  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
273
274
275
}

// O(E)
276
277
Graph::EdgeArray Graph::InEdges(dgl_id_t vid) const {
  CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
Minjie Wang's avatar
Minjie Wang committed
278
  const int64_t len = reverse_adjlist_[vid].succ.size();
279
280
281
282
  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
283
  int64_t* dst_data = static_cast<int64_t*>(dst->data);
284
285
  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
286
287
    src_data[i] = reverse_adjlist_[vid].succ[i];
    eid_data[i] = reverse_adjlist_[vid].edge_id[i];
288
289
290
  }
  std::fill(dst_data, dst_data + len, vid);
  return EdgeArray{src, dst, eid};
Minjie Wang's avatar
impl  
Minjie Wang committed
291
292
293
}

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

// O(E)
322
323
324
325
326
327
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
328
  int64_t* src_data = static_cast<int64_t*>(src->data);
329
330
331
332
  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
333
    eid_data[i] = adjlist_[vid].edge_id[i];
334
335
336
  }
  std::fill(src_data, src_data + len, vid);
  return EdgeArray{src, dst, eid};
Minjie Wang's avatar
impl  
Minjie Wang committed
337
338
339
}

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

Minjie Wang's avatar
Minjie Wang committed
367
// O(E*log(E)) if sort is required; otherwise, O(E)
368
Graph::EdgeArray Graph::Edges(const std::string &order) const {
Minjie Wang's avatar
impl  
Minjie Wang committed
369
370
371
  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});
372
  IdArray eid = IdArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
Minjie Wang's avatar
Minjie Wang committed
373

374
  if (order == "srcdst") {
Minjie Wang's avatar
Minjie Wang committed
375
376
377
    typedef std::tuple<int64_t, int64_t, int64_t> Tuple;
    std::vector<Tuple> tuples;
    tuples.reserve(len);
Minjie Wang's avatar
Minjie Wang committed
378
379
    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
380
    }
Minjie Wang's avatar
Minjie Wang committed
381
    // sort according to src and dst ids
Minjie Wang's avatar
Minjie Wang committed
382
383
    std::sort(tuples.begin(), tuples.end(),
        [] (const Tuple& t1, const Tuple& t2) {
Minjie Wang's avatar
Minjie Wang committed
384
385
          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
386
        });
387

Minjie Wang's avatar
Minjie Wang committed
388
389
390
    // make return arrays
    int64_t* src_ptr = static_cast<int64_t*>(src->data);
    int64_t* dst_ptr = static_cast<int64_t*>(dst->data);
391
    int64_t* eid_ptr = static_cast<int64_t*>(eid->data);
Minjie Wang's avatar
Minjie Wang committed
392
    for (size_t i = 0; i < tuples.size(); ++i) {
Minjie Wang's avatar
Minjie Wang committed
393
394
      src_ptr[i] = std::get<0>(tuples[i]);
      dst_ptr[i] = std::get<1>(tuples[i]);
395
      eid_ptr[i] = std::get<2>(tuples[i]);
Minjie Wang's avatar
Minjie Wang committed
396
397
398
399
    }
  } else {
    int64_t* src_ptr = static_cast<int64_t*>(src->data);
    int64_t* dst_ptr = static_cast<int64_t*>(dst->data);
400
    int64_t* eid_ptr = static_cast<int64_t*>(eid->data);
Minjie Wang's avatar
Minjie Wang committed
401
402
403
404
    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
405
    }
Minjie Wang's avatar
impl  
Minjie Wang committed
406
407
  }

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

// 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
421
    rst_data[i] = reverse_adjlist_[vid].succ.size();
Minjie Wang's avatar
impl  
Minjie Wang committed
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
  }
  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
441
442
443
444
445
446
447
448
449
450
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;
451
  rst.graph = std::make_shared<Graph>(IsMultigraph());
Minjie Wang's avatar
Minjie Wang committed
452
  rst.induced_vertices = vids;
453
  rst.graph->AddVertices(len);
Minjie Wang's avatar
Minjie Wang committed
454
455
456
457
458
459
460
461
  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]);
462
        rst.graph->AddEdge(newvid, newsucc);
Minjie Wang's avatar
Minjie Wang committed
463
464
465
466
467
468
      }
    }
  }
  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
469
}
Minjie Wang's avatar
Minjie Wang committed
470

471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
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;
489
  rst.graph = std::make_shared<Graph>(IsMultigraph());
490
  rst.induced_edges = eids;
491
  rst.graph->AddVertices(nodes.size());
492
493
494
495

  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]];
496
    rst.graph->AddEdge(oldv2newv[src_id], oldv2newv[dst_id]);
497
498
  }

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

  return rst;
Minjie Wang's avatar
Minjie Wang committed
504
505
}

506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
std::vector<IdArray> Graph::GetAdj(bool transpose, const std::string &fmt) const {
  int64_t num_edges = NumEdges();
  int64_t num_nodes = NumVertices();
  if (fmt == "coo") {
    IdArray idx = IdArray::Empty({2 * num_edges}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
    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);
    }
    IdArray eid = IdArray::Empty({num_edges}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
    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") {
    IdArray indptr = IdArray::Empty({num_nodes + 1}, DLDataType{kDLInt, 64, 1},
                                    DLContext{kDLCPU, 0});
    IdArray indices = IdArray::Empty({num_edges}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
    IdArray eid = IdArray::Empty({num_edges}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
    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
557
  LOG(FATAL) << "not implemented";
558
  return nullptr;
Minjie Wang's avatar
Minjie Wang committed
559
560
}

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