"examples/vscode:/vscode.git/clone" did not exist on "2afa3598bae9ad3ca604b7dff3c0a9563beb68f8"
graph.cc 21.4 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
Graph::Graph(IdArray src_ids, IdArray dst_ids, size_t num_nodes,
18
19
20
21
22
    bool multigraph): is_multigraph_(multigraph) {
  CHECK(IsValidIdArray(src_ids));
  CHECK(IsValidIdArray(dst_ids));
  this->AddVertices(num_nodes);
  num_edges_ = src_ids->shape[0];
23
24
  CHECK(static_cast<int64_t>(num_edges_) == dst_ids->shape[0])
    << "vectors in COO must have the same length";
25
26
27
28
  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);
  all_edges_src_.reserve(num_edges_);
  all_edges_dst_.reserve(num_edges_);
29
  for (uint64_t i = 0; i < num_edges_; i++) {
30
31
32
33
34
35
    auto src = src_data[i];
    auto dst = dst_data[i];
    CHECK(HasVertex(src) && HasVertex(dst))
      << "Invalid vertices: src=" << src << " dst=" << dst;

    adjlist_[src].succ.push_back(dst);
36
    adjlist_[src].edge_id.push_back(i);
37
    reverse_adjlist_[dst].succ.push_back(src);
38
    reverse_adjlist_[dst].edge_id.push_back(i);
39
40
41
42
43
44

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

Minjie Wang's avatar
impl  
Minjie Wang committed
45
46
47
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
48
  reverse_adjlist_.resize(reverse_adjlist_.size() + num_vertices);
Minjie Wang's avatar
impl  
Minjie Wang committed
49
50
51
52
53
}

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

Minjie Wang's avatar
impl  
Minjie Wang committed
56
  dgl_id_t eid = num_edges_++;
57

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

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

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

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

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

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

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

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

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

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

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

  // 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
200
201
202
}

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

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

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

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

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

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

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

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

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

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

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

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

470
Subgraph Graph::EdgeSubgraph(IdArray eids, bool preserve_nodes) const {
471
  CHECK(IsValidIdArray(eids)) << "Invalid edge id array.";
472
473
474
475
476
  const auto len = eids->shape[0];
  std::vector<dgl_id_t> nodes;
  const int64_t* eid_data = static_cast<int64_t*>(eids->data);

  Subgraph rst;
477
478
479
480
481
482
483
484
485
486
487
  if (!preserve_nodes) {
    std::unordered_map<dgl_id_t, dgl_id_t> oldv2newv;

    for (int64_t i = 0; i < len; ++i) {
      const dgl_id_t src_id = all_edges_src_[eid_data[i]];
      const 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);
    }
488

489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
    rst.graph = std::make_shared<Graph>(IsMultigraph());
    rst.induced_edges = eids;
    rst.graph->AddVertices(nodes.size());

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

    rst.induced_vertices = IdArray::Empty(
        {static_cast<int64_t>(nodes.size())}, eids->dtype, eids->ctx);
    std::copy(nodes.begin(), nodes.end(), static_cast<int64_t*>(rst.induced_vertices->data));
  } else {
    rst.graph = std::make_shared<Graph>(IsMultigraph());
    rst.induced_edges = eids;
    rst.graph->AddVertices(NumVertices());

    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]];
      rst.graph->AddEdge(src_id, dst_id);
    }
512

513
514
515
516
517
518
519
    for (int64_t i = 0; i < NumVertices(); ++i)
      nodes.push_back(i);

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

  return rst;
Minjie Wang's avatar
Minjie Wang committed
522
523
}

524
std::vector<IdArray> Graph::GetAdj(bool transpose, const std::string &fmt) const {
525
526
  uint64_t num_edges = NumEdges();
  uint64_t num_nodes = NumVertices();
527
  if (fmt == "coo") {
528
529
530
531
    IdArray idx = IdArray::Empty(
        {2 * static_cast<int64_t>(num_edges)},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0});
532
533
534
535
536
537
538
539
    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);
    }
540
541
542
543
    IdArray eid = IdArray::Empty(
        {static_cast<int64_t>(num_edges)},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0});
544
545
546
547
548
549
    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") {
550
551
552
553
554
555
556
557
558
559
560
561
    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});
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
    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
589
  LOG(FATAL) << "not implemented";
590
  return nullptr;
Minjie Wang's avatar
Minjie Wang committed
591
592
}

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