graph_apis.cc 16.5 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
///////////////////////////// Graph API ///////////////////////////////////

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreateMutable")
77
.set_body([] (DGLArgs args, DGLRetValue* rv) {
78
79
    bool multigraph = static_cast<bool>(args[0]);
    GraphHandle ghandle = new Graph(multigraph);
Minjie Wang's avatar
Minjie Wang committed
80
81
    *rv = ghandle;
  });
Minjie Wang's avatar
Minjie Wang committed
82

83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    const IdArray src_ids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[0]));
    const IdArray dst_ids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray edge_ids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
    const bool multigraph = static_cast<bool>(args[3]);
    const int64_t num_nodes = static_cast<int64_t>(args[4]);
    const bool readonly = static_cast<bool>(args[5]);
    GraphHandle ghandle;
    if (readonly)
      ghandle = new ImmutableGraph(src_ids, dst_ids, edge_ids, num_nodes, multigraph);
    else
      ghandle = new Graph(src_ids, dst_ids, edge_ids, num_nodes, multigraph);
    *rv = ghandle;
  });

99
100
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
101
    GraphHandle ghandle = args[0];
102
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
103
104
    delete gptr;
  });
Minjie Wang's avatar
Minjie Wang committed
105

106
107
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
108
    GraphHandle ghandle = args[0];
109
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
110
111
112
    uint64_t num_vertices = args[1];
    gptr->AddVertices(num_vertices);
  });
Minjie Wang's avatar
Minjie Wang committed
113

114
115
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdge")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
116
    GraphHandle ghandle = args[0];
117
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
118
119
120
121
    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
122

123
124
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
125
    GraphHandle ghandle = args[0];
126
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
127
128
    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
Minjie Wang's avatar
Minjie Wang committed
129
130
131
    gptr->AddEdges(src, dst);
  });

132
133
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphClear")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
134
    GraphHandle ghandle = args[0];
135
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
136
137
    gptr->Clear();
  });
Minjie Wang's avatar
Minjie Wang committed
138

139
140
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsMultigraph")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
141
142
    GraphHandle ghandle = args[0];
    // NOTE: not const since we have caches
143
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
144
145
146
    *rv = gptr->IsMultigraph();
  });

147
148
149
150
151
152
153
154
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();
  });

155
156
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
157
    GraphHandle ghandle = args[0];
158
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
159
160
    *rv = static_cast<int64_t>(gptr->NumVertices());
  });
Minjie Wang's avatar
Minjie Wang committed
161

162
163
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
164
    GraphHandle ghandle = args[0];
165
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
166
167
168
    *rv = static_cast<int64_t>(gptr->NumEdges());
  });

169
170
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertex")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
171
    GraphHandle ghandle = args[0];
172
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
173
174
175
176
    const dgl_id_t vid = args[1];
    *rv = gptr->HasVertex(vid);
  });

177
178
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
179
    GraphHandle ghandle = args[0];
180
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
181
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
Minjie Wang's avatar
Minjie Wang committed
182
183
184
    *rv = gptr->HasVertices(vids);
  });

185
186
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLMapSubgraphNID")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
187
188
189
190
191
    const IdArray parent_vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[0]));
    const IdArray query = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = GraphOp::MapParentIdToSubgraphId(parent_vids, query);
  });

192
193
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgeBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
194
    GraphHandle ghandle = args[0];
195
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
196
197
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
198
    *rv = gptr->HasEdgeBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
199
200
  });

201
202
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgesBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
203
    GraphHandle ghandle = args[0];
204
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
205
206
    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
207
    *rv = gptr->HasEdgesBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
208
209
  });

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

219
220
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphSuccessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
221
    GraphHandle ghandle = args[0];
222
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
223
224
225
226
227
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
    *rv = gptr->Successors(vid, radius);
  });

228
229
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeId")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
230
    GraphHandle ghandle = args[0];
231
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
232
233
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
234
    *rv = gptr->EdgeId(src, dst);
Minjie Wang's avatar
Minjie Wang committed
235
236
  });

237
238
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeIds")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
239
    GraphHandle ghandle = args[0];
240
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
241
242
    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
243
244
245
    *rv = ConvertEdgeArrayToPackedFunc(gptr->EdgeIds(src, dst));
  });

246
247
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFindEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
248
    GraphHandle ghandle = args[0];
249
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
250
251
    const IdArray eids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = ConvertEdgeArrayToPackedFunc(gptr->FindEdges(eids));
Minjie Wang's avatar
Minjie Wang committed
252
253
  });

254
255
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
256
    GraphHandle ghandle = args[0];
257
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
258
    const dgl_id_t vid = args[1];
259
    *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
260
261
  });

262
263
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
264
    GraphHandle ghandle = args[0];
265
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
266
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
267
    *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
268
269
  });

270
271
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_1")
.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
    const dgl_id_t vid = args[1];
275
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
276
277
  });

278
279
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
280
    GraphHandle ghandle = args[0];
281
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
282
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
283
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
284
285
  });

286
287
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
288
    GraphHandle ghandle = args[0];
289
290
291
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
    std::string order = args[1];
    *rv = ConvertEdgeArrayToPackedFunc(gptr->Edges(order));
Minjie Wang's avatar
Minjie Wang committed
292
293
  });

294
295
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
296
    GraphHandle ghandle = args[0];
297
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
298
299
300
301
    const dgl_id_t vid = args[1];
    *rv = static_cast<int64_t>(gptr->InDegree(vid));
  });

302
303
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
304
    GraphHandle ghandle = args[0];
305
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
306
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
Minjie Wang's avatar
Minjie Wang committed
307
308
309
    *rv = gptr->InDegrees(vids);
  });

310
311
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
312
    GraphHandle ghandle = args[0];
313
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
314
315
316
317
    const dgl_id_t vid = args[1];
    *rv = static_cast<int64_t>(gptr->OutDegree(vid));
  });

318
319
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegrees")
.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
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
Minjie Wang's avatar
Minjie Wang committed
323
324
    *rv = gptr->OutDegrees(vids);
  });
Minjie Wang's avatar
Minjie Wang committed
325

326
327
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphVertexSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
328
    GraphHandle ghandle = args[0];
329
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
330
331
332
333
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = ConvertSubgraphToPackedFunc(gptr->VertexSubgraph(vids));
  });

334
335
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
336
    GraphHandle ghandle = args[0];
337
    const GraphInterface *gptr = static_cast<GraphInterface*>(ghandle);
338
339
340
341
    const IdArray eids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = ConvertSubgraphToPackedFunc(gptr->EdgeSubgraph(eids));
  });

342
343
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointUnion")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
344
345
346
347
348
    void* list = args[0];
    GraphHandle* inhandles = static_cast<GraphHandle*>(list);
    int list_size = args[1];
    std::vector<const Graph*> graphs;
    for (int i = 0; i < list_size; ++i) {
349
350
351
      const GraphInterface *ptr = static_cast<const GraphInterface *>(inhandles[i]);
      const Graph* gr = dynamic_cast<const Graph*>(ptr);
      CHECK(gr) << "_CAPI_DGLDisjointUnion isn't implemented in immutable graph";
352
353
354
      graphs.push_back(gr);
    }
    Graph* gptr = new Graph();
Minjie Wang's avatar
Minjie Wang committed
355
    *gptr = GraphOp::DisjointUnion(std::move(graphs));
356
357
358
359
    GraphHandle ghandle = gptr;
    *rv = ghandle;
  });

360
361
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionByNum")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
362
    GraphHandle ghandle = args[0];
363
364
365
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const Graph* gptr = dynamic_cast<const Graph*>(ptr);
    CHECK(gptr) << "_CAPI_DGLDisjointPartitionByNum isn't implemented in immutable graph";
Minjie Wang's avatar
Minjie Wang committed
366
367
368
369
370
371
372
373
374
375
376
377
378
379
    int64_t num = args[1];
    std::vector<Graph>&& rst = GraphOp::DisjointPartitionByNum(gptr, num);
    // 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) {
      Graph* ptr = new Graph();
      *ptr = std::move(rst[i]);
      ptr_array_data[i] = reinterpret_cast<std::intptr_t>(ptr);
    }
    *rv = ptr_array;
  });

380
381
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionBySizes")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
382
    GraphHandle ghandle = args[0];
383
384
385
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const Graph* gptr = dynamic_cast<const Graph*>(ptr);
    CHECK(gptr) << "_CAPI_DGLDisjointPartitionBySizes isn't implemented in immutable graph";
Minjie Wang's avatar
Minjie Wang committed
386
387
388
389
390
391
392
393
394
395
396
397
398
    const IdArray sizes = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    std::vector<Graph>&& 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) {
      Graph* ptr = new Graph();
      *ptr = std::move(rst[i]);
      ptr_array_data[i] = reinterpret_cast<std::intptr_t>(ptr);
    }
    *rv = ptr_array;
  });
GaiYu0's avatar
cpp lg  
GaiYu0 committed
399

400
401
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphLineGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
GaiYu0's avatar
cpp lg  
GaiYu0 committed
402
403
    GraphHandle ghandle = args[0];
    bool backtracking = args[1];
404
405
406
    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
407
408
409
410
411
    Graph* lgptr = new Graph();
    *lgptr = GraphOp::LineGraph(gptr, backtracking);
    GraphHandle lghandle = lgptr;
    *rv = lghandle;
  });
GaiYu0's avatar
GaiYu0 committed
412

413
414
415
416
417
418
419
420
421
422
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);
  });

423
424
425
426
427
428
429
430
431
432
433
434
435
DGL_REGISTER_GLOBAL("nodeflow._CAPI_NodeFlowGetBlockAdj")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    std::string format = args[1];
    int64_t layer0_size = args[2];
    int64_t start = args[3];
    int64_t end = args[4];
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const ImmutableGraph* gptr = dynamic_cast<const ImmutableGraph*>(ptr);
    auto res = GetNodeFlowSlice(*gptr, format, layer0_size, start, end, true);
    *rv = ConvertAdjToPackedFunc(res);
  });

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