graph_apis.cc 18.1 KB
Newer Older
1
2
3
4
5
/*!
 *  Copyright (c) 2018 by Contributors
 * \file graph/graph.cc
 * \brief DGL graph index APIs
 */
Minjie Wang's avatar
Minjie Wang committed
6
#include <dgl/graph.h>
7
#include <dgl/immutable_graph.h>
Minjie Wang's avatar
Minjie Wang committed
8
#include <dgl/graph_op.h>
Da Zheng's avatar
Da Zheng committed
9
#include <dgl/sampler.h>
10
#include <dgl/nodeflow.h>
Lingfan Yu's avatar
Lingfan Yu committed
11
#include "../c_api_common.h"
Minjie Wang's avatar
Minjie Wang committed
12

13
14
15
16
17
using dgl::runtime::DGLArgs;
using dgl::runtime::DGLArgValue;
using dgl::runtime::DGLRetValue;
using dgl::runtime::PackedFunc;
using dgl::runtime::NDArray;
Minjie Wang's avatar
Minjie Wang committed
18
19

namespace dgl {
Minjie Wang's avatar
Minjie Wang committed
20

Minjie Wang's avatar
Minjie Wang committed
21
namespace {
Minjie Wang's avatar
Minjie Wang committed
22
// Convert EdgeArray structure to PackedFunc.
23
24
template<class EdgeArray>
PackedFunc ConvertEdgeArrayToPackedFunc(const EdgeArray& ea) {
25
  auto body = [ea] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
26
      const int which = args[0];
27
      if (which == 0) {
Minjie Wang's avatar
Minjie Wang committed
28
29
30
31
32
33
34
35
36
37
38
39
        *rv = std::move(ea.src);
      } else if (which == 1) {
        *rv = std::move(ea.dst);
      } else if (which == 2) {
        *rv = std::move(ea.id);
      } else {
        LOG(FATAL) << "invalid choice";
      }
    };
  return PackedFunc(body);
}

40
41
42
43
44
45
46
47
48
49
50
51
52
// Convert CSRArray structure to PackedFunc.
PackedFunc ConvertAdjToPackedFunc(const std::vector<IdArray>& ea) {
  auto body = [ea] (DGLArgs args, DGLRetValue* rv) {
      const int which = args[0];
      if ((size_t) which < ea.size()) {
        *rv = std::move(ea[which]);
      } else {
        LOG(FATAL) << "invalid choice";
      }
    };
  return PackedFunc(body);
}

Minjie Wang's avatar
Minjie Wang committed
53
54
// Convert Subgraph structure to PackedFunc.
PackedFunc ConvertSubgraphToPackedFunc(const Subgraph& sg) {
55
  auto body = [sg] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
56
      const int which = args[0];
Minjie Wang's avatar
Minjie Wang committed
57
      if (which == 0) {
58
        GraphInterface* gptr = sg.graph->Reset();
Minjie Wang's avatar
Minjie Wang committed
59
60
        GraphHandle ghandle = gptr;
        *rv = ghandle;
61
      } else if (which == 1) {
Minjie Wang's avatar
Minjie Wang committed
62
        *rv = std::move(sg.induced_vertices);
63
      } else if (which == 2) {
Minjie Wang's avatar
Minjie Wang committed
64
        *rv = std::move(sg.induced_edges);
65
66
67
      } else {
        LOG(FATAL) << "invalid choice";
      }
Minjie Wang's avatar
Minjie Wang committed
68
69
70
    };
  return PackedFunc(body);
}
Minjie Wang's avatar
Minjie Wang committed
71

Minjie Wang's avatar
Minjie Wang committed
72
}  // namespace
Minjie Wang's avatar
Minjie Wang committed
73

74
75
76
77
78
79
namespace {
// This namespace contains template functions for batching
// and unbatching over graph and immutable graph
template<typename T>
void DGLDisjointPartitionByNum(const T *gptr, DGLArgs args, DGLRetValue *rv) {
  int64_t num = args[1];
80
  std::vector<T> &&rst = GraphOp::DisjointPartitionByNum(gptr, num);
81
82
83
84
85
  // return the pointer array as an integer array
  const int64_t len = rst.size();
  NDArray ptr_array = NDArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  int64_t *ptr_array_data = static_cast<int64_t *>(ptr_array->data);
  for (size_t i = 0; i < rst.size(); ++i) {
86
    GraphInterface *ptr = rst[i].Reset();
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
    ptr_array_data[i] = reinterpret_cast<std::intptr_t>(ptr);
  }
  *rv = ptr_array;
}

template<typename T>
void DGLDisjointUnion(GraphHandle *inhandles, int list_size, DGLRetValue *rv) {
  std::vector<const T *> graphs;
  for (int i = 0; i < list_size; ++i) {
    const GraphInterface *ptr = static_cast<const GraphInterface *>(inhandles[i]);
    const T *gr = dynamic_cast<const T *>(ptr);
    CHECK(gr) << "Error: Attempted to batch MutableGraph with ImmutableGraph";
    graphs.push_back(gr);
  }

102
  GraphHandle ghandle = GraphOp::DisjointUnion(std::move(graphs)).Reset();
103
104
105
106
107
108
109
110
111
112
113
  *rv = ghandle;
}

template<typename T>
void DGLDisjointPartitionBySizes(const T *gptr, const IdArray sizes, DGLRetValue *rv) {
  std::vector<T> &&rst = GraphOp::DisjointPartitionBySizes(gptr, sizes);
  // return the pointer array as an integer array
  const int64_t len = rst.size();
  NDArray ptr_array = NDArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  int64_t *ptr_array_data = static_cast<int64_t *>(ptr_array->data);
  for (size_t i = 0; i < rst.size(); ++i) {
114
    GraphInterface *ptr = rst[i].Reset();
115
116
117
118
119
120
    ptr_array_data[i] = reinterpret_cast<std::intptr_t>(ptr);
  }
  *rv = ptr_array;
}
}  // namespace

121
122
123
///////////////////////////// Graph API ///////////////////////////////////

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreateMutable")
124
.set_body([] (DGLArgs args, DGLRetValue* rv) {
125
126
    bool multigraph = static_cast<bool>(args[0]);
    GraphHandle ghandle = new Graph(multigraph);
Minjie Wang's avatar
Minjie Wang committed
127
128
    *rv = ghandle;
  });
Minjie Wang's avatar
Minjie Wang committed
129

130
131
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
132
133
134
135
136
    const IdArray src_ids = args[0];
    const IdArray dst_ids = args[1];
    const bool multigraph = static_cast<bool>(args[2]);
    const int64_t num_nodes = static_cast<int64_t>(args[3]);
    const bool readonly = static_cast<bool>(args[4]);
137
    GraphHandle ghandle;
138
139
140
141
142
143
144
145
146
147
    if (readonly) {
      // TODO(minjie): The array copy here is unnecessary and adds extra overhead.
      //   However, with MXNet backend, the memory would be corrupted if we directly
      //   save the passed-in ndarrays into DGL's graph object. We hope MXNet team
      //   could help look into this.
      COOPtr coo(new COO(num_nodes, Clone(src_ids), Clone(dst_ids), multigraph));
      ghandle = new ImmutableGraph(coo);
    } else {
      ghandle = new Graph(src_ids, dst_ids, num_nodes, multigraph);
    }
148
149
150
    *rv = ghandle;
  });

151
152
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCSRCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
153
154
155
    const IdArray indptr = args[0];
    const IdArray indices = args[1];
    const IdArray edge_ids = args[2];
156
157
158
    const std::string shared_mem_name = args[3];
    const bool multigraph = static_cast<bool>(args[4]);
    const std::string edge_dir = args[5];
159
    CSRPtr csr;
160
    if (shared_mem_name.empty())
161
162
163
164
165
      // TODO(minjie): The array copy here is unnecessary and adds extra overhead.
      //   However, with MXNet backend, the memory would be corrupted if we directly
      //   save the passed-in ndarrays into DGL's graph object. We hope MXNet team
      //   could help look into this.
      csr.reset(new CSR(Clone(indptr), Clone(indices), Clone(edge_ids), multigraph));
166
    else
167
      csr.reset(new CSR(indptr, indices, edge_ids, multigraph, shared_mem_name));
168
169
170

    GraphHandle ghandle;
    if (edge_dir == "in")
171
      ghandle = new ImmutableGraph(csr, nullptr);
172
    else
173
      ghandle = new ImmutableGraph(nullptr, csr);
174
175
176
177
178
179
180
181
182
183
    *rv = ghandle;
  });

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCSRCreateMMap")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    const std::string shared_mem_name = args[0];
    const int64_t num_vertices = args[1];
    const int64_t num_edges = args[2];
    const bool multigraph = static_cast<bool>(args[3]);
    const std::string edge_dir = args[4];
184
185
    // TODO(minjie): how to know multigraph
    CSRPtr csr(new CSR(shared_mem_name, num_vertices, num_edges, multigraph));
186
187
    GraphHandle ghandle;
    if (edge_dir == "in")
188
      ghandle = new ImmutableGraph(csr, nullptr);
189
    else
190
      ghandle = new ImmutableGraph(nullptr, csr);
191
192
193
    *rv = ghandle;
  });

194
195
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
196
    GraphHandle ghandle = args[0];
197
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
198
199
    delete gptr;
  });
Minjie Wang's avatar
Minjie Wang committed
200

201
202
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
203
    GraphHandle ghandle = args[0];
204
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
205
206
207
    uint64_t num_vertices = args[1];
    gptr->AddVertices(num_vertices);
  });
Minjie Wang's avatar
Minjie Wang committed
208

209
210
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdge")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
211
    GraphHandle ghandle = args[0];
212
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
213
214
215
216
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
    gptr->AddEdge(src, dst);
  });
Minjie Wang's avatar
Minjie Wang committed
217

218
219
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
220
    GraphHandle ghandle = args[0];
221
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
222
223
    const IdArray src = args[1];
    const IdArray dst = args[2];
Minjie Wang's avatar
Minjie Wang committed
224
225
226
    gptr->AddEdges(src, dst);
  });

227
228
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphClear")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
229
    GraphHandle ghandle = args[0];
230
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
231
232
    gptr->Clear();
  });
Minjie Wang's avatar
Minjie Wang committed
233

234
235
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsMultigraph")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
236
237
    GraphHandle ghandle = args[0];
    // NOTE: not const since we have caches
238
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
239
240
241
    *rv = gptr->IsMultigraph();
  });

242
243
244
245
246
247
248
249
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsReadonly")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
    GraphHandle ghandle = args[0];
    // NOTE: not const since we have caches
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
    *rv = gptr->IsReadonly();
  });

250
251
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
252
    GraphHandle ghandle = args[0];
253
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
254
255
    *rv = static_cast<int64_t>(gptr->NumVertices());
  });
Minjie Wang's avatar
Minjie Wang committed
256

257
258
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
259
    GraphHandle ghandle = args[0];
260
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
261
262
263
    *rv = static_cast<int64_t>(gptr->NumEdges());
  });

264
265
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertex")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
266
    GraphHandle ghandle = args[0];
267
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
268
269
270
271
    const dgl_id_t vid = args[1];
    *rv = gptr->HasVertex(vid);
  });

272
273
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
274
    GraphHandle ghandle = args[0];
275
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
276
    const IdArray vids = args[1];
Minjie Wang's avatar
Minjie Wang committed
277
278
279
    *rv = gptr->HasVertices(vids);
  });

280
281
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLMapSubgraphNID")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
282
283
    const IdArray parent_vids = args[0];
    const IdArray query = args[1];
284
285
286
    *rv = GraphOp::MapParentIdToSubgraphId(parent_vids, query);
  });

287
288
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgeBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
289
    GraphHandle ghandle = args[0];
290
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
291
292
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
293
    *rv = gptr->HasEdgeBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
294
295
  });

296
297
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgesBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
298
    GraphHandle ghandle = args[0];
299
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
300
301
    const IdArray src = args[1];
    const IdArray dst = args[2];
302
    *rv = gptr->HasEdgesBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
303
304
  });

305
306
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphPredecessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
307
    GraphHandle ghandle = args[0];
308
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
309
310
311
312
313
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
    *rv = gptr->Predecessors(vid, radius);
  });

314
315
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphSuccessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
316
    GraphHandle ghandle = args[0];
317
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
318
319
320
321
322
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
    *rv = gptr->Successors(vid, radius);
  });

323
324
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeId")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
325
    GraphHandle ghandle = args[0];
326
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
327
328
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
329
    *rv = gptr->EdgeId(src, dst);
Minjie Wang's avatar
Minjie Wang committed
330
331
  });

332
333
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeIds")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
334
    GraphHandle ghandle = args[0];
335
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
336
337
    const IdArray src = args[1];
    const IdArray dst = args[2];
338
339
340
    *rv = ConvertEdgeArrayToPackedFunc(gptr->EdgeIds(src, dst));
  });

341
342
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFindEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
343
    GraphHandle ghandle = args[0];
344
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
345
    const IdArray eids = args[1];
346
    *rv = ConvertEdgeArrayToPackedFunc(gptr->FindEdges(eids));
Minjie Wang's avatar
Minjie Wang committed
347
348
  });

349
350
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
351
    GraphHandle ghandle = args[0];
352
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
353
    const dgl_id_t vid = args[1];
354
    *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
355
356
  });

357
358
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
359
    GraphHandle ghandle = args[0];
360
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
361
    const IdArray vids = args[1];
362
    *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
363
364
  });

365
366
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
367
    GraphHandle ghandle = args[0];
368
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
369
    const dgl_id_t vid = args[1];
370
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
371
372
  });

373
374
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
375
    GraphHandle ghandle = args[0];
376
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
377
    const IdArray vids = args[1];
378
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
379
380
  });

381
382
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
383
    GraphHandle ghandle = args[0];
384
385
386
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
    std::string order = args[1];
    *rv = ConvertEdgeArrayToPackedFunc(gptr->Edges(order));
Minjie Wang's avatar
Minjie Wang committed
387
388
  });

389
390
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
391
    GraphHandle ghandle = args[0];
392
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
393
394
395
396
    const dgl_id_t vid = args[1];
    *rv = static_cast<int64_t>(gptr->InDegree(vid));
  });

397
398
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
399
    GraphHandle ghandle = args[0];
400
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
401
    const IdArray vids = args[1];
Minjie Wang's avatar
Minjie Wang committed
402
403
404
    *rv = gptr->InDegrees(vids);
  });

405
406
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
407
    GraphHandle ghandle = args[0];
408
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
409
410
411
412
    const dgl_id_t vid = args[1];
    *rv = static_cast<int64_t>(gptr->OutDegree(vid));
  });

413
414
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
415
    GraphHandle ghandle = args[0];
416
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
417
    const IdArray vids = args[1];
Minjie Wang's avatar
Minjie Wang committed
418
419
    *rv = gptr->OutDegrees(vids);
  });
Minjie Wang's avatar
Minjie Wang committed
420

421
422
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphVertexSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
423
    GraphHandle ghandle = args[0];
424
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
425
    const IdArray vids = args[1];
Minjie Wang's avatar
Minjie Wang committed
426
427
428
    *rv = ConvertSubgraphToPackedFunc(gptr->VertexSubgraph(vids));
  });

429
430
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
431
    GraphHandle ghandle = args[0];
432
    const GraphInterface *gptr = static_cast<GraphInterface*>(ghandle);
433
    const IdArray eids = args[1];
434
435
436
    *rv = ConvertSubgraphToPackedFunc(gptr->EdgeSubgraph(eids));
  });

437
438
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointUnion")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
439
440
441
    void* list = args[0];
    GraphHandle* inhandles = static_cast<GraphHandle*>(list);
    int list_size = args[1];
442
443
444
445
446
447
448
449
    const GraphInterface *ptr = static_cast<const GraphInterface *>(inhandles[0]);
    const ImmutableGraph *im_gr = dynamic_cast<const ImmutableGraph *>(ptr);
    const Graph *gr = dynamic_cast<const Graph *>(ptr);
    if (gr) {
      DGLDisjointUnion<Graph>(inhandles, list_size, rv);
    } else {
      CHECK(im_gr) << "Args[0] is not a list of valid DGLGraph";
      DGLDisjointUnion<ImmutableGraph>(inhandles, list_size, rv);
450
451
452
    }
  });

453
454
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionByNum")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
455
    GraphHandle ghandle = args[0];
456
457
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const Graph* gptr = dynamic_cast<const Graph*>(ptr);
458
459
460
461
462
463
    const ImmutableGraph* im_gptr = dynamic_cast<const ImmutableGraph*>(ptr);
    if (gptr) {
      DGLDisjointPartitionByNum(gptr, args, rv);
    } else {
      CHECK(im_gptr) << "Args[0] is not a valid DGLGraph";
      DGLDisjointPartitionByNum(im_gptr, args, rv);
Minjie Wang's avatar
Minjie Wang committed
464
465
466
    }
  });

467
468
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionBySizes")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
469
    GraphHandle ghandle = args[0];
470
    const IdArray sizes = args[1];
471
472
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const Graph* gptr = dynamic_cast<const Graph*>(ptr);
473
474
475
476
477
478
    const ImmutableGraph* im_gptr = dynamic_cast<const ImmutableGraph*>(ptr);
    if (gptr) {
      DGLDisjointPartitionBySizes(gptr, sizes, rv);
    } else {
      CHECK(im_gptr) << "Args[0] is not a valid DGLGraph";
      DGLDisjointPartitionBySizes(im_gptr, sizes, rv);
Minjie Wang's avatar
Minjie Wang committed
479
    }
480
});
GaiYu0's avatar
cpp lg  
GaiYu0 committed
481

482
483
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphLineGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
GaiYu0's avatar
cpp lg  
GaiYu0 committed
484
485
    GraphHandle ghandle = args[0];
    bool backtracking = args[1];
486
487
488
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const Graph* gptr = dynamic_cast<const Graph*>(ptr);
    CHECK(gptr) << "_CAPI_DGLGraphLineGraph isn't implemented in immutable graph";
GaiYu0's avatar
cpp lg  
GaiYu0 committed
489
490
491
492
493
    Graph* lgptr = new Graph();
    *lgptr = GraphOp::LineGraph(gptr, backtracking);
    GraphHandle lghandle = lgptr;
    *rv = lghandle;
  });
GaiYu0's avatar
GaiYu0 committed
494

495
496
497
498
499
500
501
502
503
504
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphGetAdj")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    bool transpose = args[1];
    std::string format = args[2];
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    auto res = ptr->GetAdj(transpose, format);
    *rv = ConvertAdjToPackedFunc(res);
  });

Minjie Wang's avatar
Minjie Wang committed
505
}  // namespace dgl