graph_apis.cc 21.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
132
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
133
134
    const IdArray src_ids = args[0];
    const IdArray dst_ids = args[1];
135
136
137
    const int multigraph = args[2];
    const int64_t num_nodes = args[3];
    const bool readonly = args[4];
138
    GraphHandle ghandle;
139
    if (readonly) {
140
      if (multigraph == kBoolUnknown) {
141
        COOPtr coo(new COO(num_nodes, src_ids, dst_ids));
142
143
        ghandle = new ImmutableGraph(coo);
      } else {
144
        COOPtr coo(new COO(num_nodes, src_ids, dst_ids, multigraph));
145
146
        ghandle = new ImmutableGraph(coo);
      }
147
    } else {
148
      CHECK_NE(multigraph, kBoolUnknown);
149
150
      ghandle = new Graph(src_ids, dst_ids, num_nodes, multigraph);
    }
151
152
153
    *rv = ghandle;
  });

154
155
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCSRCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
156
157
    const IdArray indptr = args[0];
    const IdArray indices = args[1];
158
    const std::string shared_mem_name = args[2];
159
    const int multigraph = args[3];
160
    const std::string edge_dir = args[4];
161
    CSRPtr csr;
162
163
164
165
166
167

    IdArray edge_ids = IdArray::Empty({indices->shape[0]},
                                      DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
    int64_t *edge_data = static_cast<int64_t *>(edge_ids->data);
    for (size_t i = 0; i < edge_ids->shape[0]; i++)
      edge_data[i] = i;
168
169
    if (shared_mem_name.empty()) {
      if (multigraph == kBoolUnknown) {
170
        csr.reset(new CSR(indptr, indices, edge_ids));
171
      } else {
172
        csr.reset(new CSR(indptr, indices, edge_ids, multigraph));
173
174
175
176
177
178
179
180
      }
    } else {
      if (multigraph == kBoolUnknown) {
        csr.reset(new CSR(indptr, indices, edge_ids, shared_mem_name));
      } else {
        csr.reset(new CSR(indptr, indices, edge_ids, multigraph, shared_mem_name));
      }
    }
181
182
183

    GraphHandle ghandle;
    if (edge_dir == "in")
184
      ghandle = new ImmutableGraph(csr, nullptr);
185
    else
186
      ghandle = new ImmutableGraph(nullptr, csr);
187
188
189
190
191
192
193
194
    *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];
195
    const bool multigraph = args[3];
196
    const std::string edge_dir = args[4];
197
198
    // TODO(minjie): how to know multigraph
    CSRPtr csr(new CSR(shared_mem_name, num_vertices, num_edges, multigraph));
199
200
    GraphHandle ghandle;
    if (edge_dir == "in")
201
      ghandle = new ImmutableGraph(csr, nullptr);
202
    else
203
      ghandle = new ImmutableGraph(nullptr, csr);
204
205
206
    *rv = ghandle;
  });

207
208
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
209
    GraphHandle ghandle = args[0];
210
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
211
212
    delete gptr;
  });
Minjie Wang's avatar
Minjie Wang committed
213

214
215
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
216
    GraphHandle ghandle = args[0];
217
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
218
219
220
    uint64_t num_vertices = args[1];
    gptr->AddVertices(num_vertices);
  });
Minjie Wang's avatar
Minjie Wang committed
221

222
223
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdge")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
224
    GraphHandle ghandle = args[0];
225
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
226
227
228
229
    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
230

231
232
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
233
    GraphHandle ghandle = args[0];
234
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
235
236
    const IdArray src = args[1];
    const IdArray dst = args[2];
Minjie Wang's avatar
Minjie Wang committed
237
238
239
    gptr->AddEdges(src, dst);
  });

240
241
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphClear")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
242
    GraphHandle ghandle = args[0];
243
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
244
245
    gptr->Clear();
  });
Minjie Wang's avatar
Minjie Wang committed
246

247
248
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsMultigraph")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
249
250
    GraphHandle ghandle = args[0];
    // NOTE: not const since we have caches
251
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
252
253
254
    *rv = gptr->IsMultigraph();
  });

255
256
257
258
259
260
261
262
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();
  });

263
264
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
265
    GraphHandle ghandle = args[0];
266
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
267
268
    *rv = static_cast<int64_t>(gptr->NumVertices());
  });
Minjie Wang's avatar
Minjie Wang committed
269

270
271
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
272
    GraphHandle ghandle = args[0];
273
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
274
275
276
    *rv = static_cast<int64_t>(gptr->NumEdges());
  });

277
278
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertex")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
279
    GraphHandle ghandle = args[0];
280
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
281
282
283
284
    const dgl_id_t vid = args[1];
    *rv = gptr->HasVertex(vid);
  });

285
286
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
287
    GraphHandle ghandle = args[0];
288
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
289
    const IdArray vids = args[1];
Minjie Wang's avatar
Minjie Wang committed
290
291
292
    *rv = gptr->HasVertices(vids);
  });

293
294
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLMapSubgraphNID")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
295
296
    const IdArray parent_vids = args[0];
    const IdArray query = args[1];
297
298
299
    *rv = GraphOp::MapParentIdToSubgraphId(parent_vids, query);
  });

300
301
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgeBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
302
    GraphHandle ghandle = args[0];
303
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
304
305
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
306
    *rv = gptr->HasEdgeBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
307
308
  });

309
310
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgesBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
311
    GraphHandle ghandle = args[0];
312
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
313
314
    const IdArray src = args[1];
    const IdArray dst = args[2];
315
    *rv = gptr->HasEdgesBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
316
317
  });

318
319
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphPredecessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
320
    GraphHandle ghandle = args[0];
321
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
322
323
324
325
326
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
    *rv = gptr->Predecessors(vid, radius);
  });

327
328
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphSuccessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
329
    GraphHandle ghandle = args[0];
330
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
331
332
333
334
335
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
    *rv = gptr->Successors(vid, radius);
  });

336
337
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeId")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
338
    GraphHandle ghandle = args[0];
339
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
340
341
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
342
    *rv = gptr->EdgeId(src, dst);
Minjie Wang's avatar
Minjie Wang committed
343
344
  });

345
346
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeIds")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
347
    GraphHandle ghandle = args[0];
348
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
349
350
    const IdArray src = args[1];
    const IdArray dst = args[2];
351
352
353
    *rv = ConvertEdgeArrayToPackedFunc(gptr->EdgeIds(src, dst));
  });

354
355
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFindEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
356
    GraphHandle ghandle = args[0];
357
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
358
    const IdArray eids = args[1];
359
    *rv = ConvertEdgeArrayToPackedFunc(gptr->FindEdges(eids));
Minjie Wang's avatar
Minjie Wang committed
360
361
  });

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

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

378
379
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
380
    GraphHandle ghandle = args[0];
381
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
382
    const dgl_id_t vid = args[1];
383
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
384
385
  });

386
387
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
388
    GraphHandle ghandle = args[0];
389
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
390
    const IdArray vids = args[1];
391
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
392
393
  });

394
395
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
396
    GraphHandle ghandle = args[0];
397
398
399
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
    std::string order = args[1];
    *rv = ConvertEdgeArrayToPackedFunc(gptr->Edges(order));
Minjie Wang's avatar
Minjie Wang committed
400
401
  });

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

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

418
419
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
420
    GraphHandle ghandle = args[0];
421
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
422
423
424
425
    const dgl_id_t vid = args[1];
    *rv = static_cast<int64_t>(gptr->OutDegree(vid));
  });

426
427
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
428
    GraphHandle ghandle = args[0];
429
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
430
    const IdArray vids = args[1];
Minjie Wang's avatar
Minjie Wang committed
431
432
    *rv = gptr->OutDegrees(vids);
  });
Minjie Wang's avatar
Minjie Wang committed
433

434
435
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphVertexSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
436
    GraphHandle ghandle = args[0];
437
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
438
    const IdArray vids = args[1];
Minjie Wang's avatar
Minjie Wang committed
439
440
441
    *rv = ConvertSubgraphToPackedFunc(gptr->VertexSubgraph(vids));
  });

442
443
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
444
    GraphHandle ghandle = args[0];
445
    const GraphInterface *gptr = static_cast<GraphInterface*>(ghandle);
446
    const IdArray eids = args[1];
447
448
    bool preserve_nodes = args[2];
    *rv = ConvertSubgraphToPackedFunc(gptr->EdgeSubgraph(eids, preserve_nodes));
449
450
  });

451
452
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointUnion")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
453
454
455
    void* list = args[0];
    GraphHandle* inhandles = static_cast<GraphHandle*>(list);
    int list_size = args[1];
456
457
458
459
460
461
462
463
    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);
464
465
466
    }
  });

467
468
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionByNum")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
469
    GraphHandle ghandle = args[0];
470
471
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const Graph* gptr = dynamic_cast<const Graph*>(ptr);
472
473
474
475
476
477
    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
478
479
480
    }
  });

481
482
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionBySizes")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
483
    GraphHandle ghandle = args[0];
484
    const IdArray sizes = args[1];
485
486
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const Graph* gptr = dynamic_cast<const Graph*>(ptr);
487
488
489
490
491
492
    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
493
    }
494
});
GaiYu0's avatar
cpp lg  
GaiYu0 committed
495

496
497
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphLineGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
GaiYu0's avatar
cpp lg  
GaiYu0 committed
498
499
    GraphHandle ghandle = args[0];
    bool backtracking = args[1];
500
501
502
    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
503
504
505
506
507
    Graph* lgptr = new Graph();
    *lgptr = GraphOp::LineGraph(gptr, backtracking);
    GraphHandle lghandle = lgptr;
    *rv = lghandle;
  });
GaiYu0's avatar
GaiYu0 committed
508

509
510
511
512
513
514
515
516
517
518
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);
  });

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
557
558
559
560
561
562
563
564
565
566
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLToImmutable")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    const GraphInterface *ptr = static_cast<GraphInterface *>(ghandle);
    GraphHandle newhandle = new ImmutableGraph(ImmutableGraph::ToImmutable(ptr));
    *rv = newhandle;
  });

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphContext")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    const GraphInterface *ptr = static_cast<GraphInterface *>(ghandle);
    *rv = ptr->Context();
  });

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLImmutableGraphCopyTo")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    const int device_type = args[1];
    const int device_id = args[2];
    DLContext ctx;
    ctx.device_type = static_cast<DLDeviceType>(device_type);
    ctx.device_id = device_id;
    const GraphInterface *ptr = static_cast<GraphInterface *>(ghandle);
    const ImmutableGraph *ig = dynamic_cast<const ImmutableGraph*>(ptr);
    CHECK(ig) << "Invalid argument: must be an immutable graph object.";
    GraphHandle newhandle = new ImmutableGraph(ig->CopyTo(ctx));
    *rv = newhandle;
  });

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumBits")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    const GraphInterface *ptr = static_cast<GraphInterface *>(ghandle);
    *rv = ptr->NumBits();
  });

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLImmutableGraphAsNumBits")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    int bits = args[1];
    const GraphInterface *ptr = static_cast<GraphInterface *>(ghandle);
    const ImmutableGraph *ig = dynamic_cast<const ImmutableGraph*>(ptr);
    CHECK(ig) << "Invalid argument: must be an immutable graph object.";
    GraphHandle newhandle = new ImmutableGraph(ig->AsNumBits(bits));
    *rv = newhandle;
  });

567
568
569
570
571
572
573
574
DGL_REGISTER_GLOBAL("transform._CAPI_DGLToSimpleGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    GraphHandle ret = GraphOp::ToSimpleGraph(ptr).Reset();
    *rv = ret;
  });

575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
DGL_REGISTER_GLOBAL("transform._CAPI_DGLToBidirectedMutableGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    Graph* bgptr = new Graph();
    *bgptr = GraphOp::ToBidirectedMutableGraph(ptr);
    GraphHandle bghandle = bgptr;
    *rv = bghandle;
  });

DGL_REGISTER_GLOBAL("transform._CAPI_DGLToBidirectedImmutableGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    GraphHandle bghandle = GraphOp::ToBidirectedImmutableGraph(ptr).Reset();
    *rv = bghandle;
  });

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