graph.cc 21 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>
Da Zheng's avatar
Da Zheng committed
7
#include <dgl/sampler.h>
Minjie Wang's avatar
impl  
Minjie Wang committed
8
#include <algorithm>
Minjie Wang's avatar
Minjie Wang committed
9
#include <unordered_map>
10
11
#include <set>
#include <functional>
GaiYu0's avatar
GaiYu0 committed
12
13
#include <tuple>
#include "../c_api_common.h"
Minjie Wang's avatar
Minjie Wang committed
14

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

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];
24
25
26
27
  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";
28
29
30
31
32
  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_);
33
  for (uint64_t i = 0; i < num_edges_; i++) {
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
    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
50
51
52
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
53
  reverse_adjlist_.resize(reverse_adjlist_.size() + num_vertices);
Minjie Wang's avatar
impl  
Minjie Wang committed
54
55
56
57
58
}

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

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

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

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

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

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

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

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

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

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

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

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

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

  // 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
205
206
207
}

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

  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];
227
228
    CHECK(HasVertex(src_id) && HasVertex(dst_id)) <<
        "invalid edge: " << src_id << " -> " << dst_id;
229
230
231
    const auto& succ = adjlist_[src_id].succ;
    for (size_t k = 0; k < succ.size(); ++k) {
      if (succ[k] == dst_id) {
232
233
234
        src.push_back(src_id);
        dst.push_back(dst_id);
        eid.push_back(adjlist_[src_id].edge_id[k]);
235
      }
Minjie Wang's avatar
impl  
Minjie Wang committed
236
237
    }
  }
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254

  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 {
255
  CHECK(IsValidIdArray(eids)) << "Invalid edge id array";
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
  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
277
278
279
}

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

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

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

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

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

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

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

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

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

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

  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;
493
  rst.graph = std::make_shared<Graph>(IsMultigraph());
494
  rst.induced_edges = eids;
495
  rst.graph->AddVertices(nodes.size());
496
497
498
499

  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]];
500
    rst.graph->AddEdge(oldv2newv[src_id], oldv2newv[dst_id]);
501
502
  }

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

  return rst;
Minjie Wang's avatar
Minjie Wang committed
508
509
}

510
std::vector<IdArray> Graph::GetAdj(bool transpose, const std::string &fmt) const {
511
512
  uint64_t num_edges = NumEdges();
  uint64_t num_nodes = NumVertices();
513
  if (fmt == "coo") {
514
515
516
517
    IdArray idx = IdArray::Empty(
        {2 * static_cast<int64_t>(num_edges)},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0});
518
519
520
521
522
523
524
525
    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);
    }
526
527
528
529
    IdArray eid = IdArray::Empty(
        {static_cast<int64_t>(num_edges)},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0});
530
531
532
533
534
535
    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") {
536
537
538
539
540
541
542
543
544
545
546
547
    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});
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
573
574
    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
575
  LOG(FATAL) << "not implemented";
576
  return nullptr;
Minjie Wang's avatar
Minjie Wang committed
577
578
}

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