graph_apis.cc 18.2 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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCSRCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    const IdArray indptr = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[0]));
    const IdArray indices = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray edge_ids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
    const std::string shared_mem_name = args[3];
    const bool multigraph = static_cast<bool>(args[4]);
    const std::string edge_dir = args[5];
    ImmutableGraph::CSR::Ptr csr;
    if (shared_mem_name.empty())
      csr.reset(new ImmutableGraph::CSR(indptr, indices, edge_ids));
    else
      csr.reset(new ImmutableGraph::CSR(indptr, indices, edge_ids, shared_mem_name));

    GraphHandle ghandle;
    if (edge_dir == "in")
      ghandle = new ImmutableGraph(csr, nullptr, multigraph);
    else
      ghandle = new ImmutableGraph(nullptr, csr, multigraph);
    *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];
    ImmutableGraph::CSR::Ptr csr(new ImmutableGraph::CSR(shared_mem_name,
                                                         num_vertices, num_edges));
    GraphHandle ghandle;
    if (edge_dir == "in")
      ghandle = new ImmutableGraph(csr, nullptr, multigraph);
    else
      ghandle = new ImmutableGraph(nullptr, csr, multigraph);
    *rv = ghandle;
  });

138
139
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
140
    GraphHandle ghandle = args[0];
141
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
142
143
    delete gptr;
  });
Minjie Wang's avatar
Minjie Wang committed
144

145
146
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
147
    GraphHandle ghandle = args[0];
148
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
149
150
151
    uint64_t num_vertices = args[1];
    gptr->AddVertices(num_vertices);
  });
Minjie Wang's avatar
Minjie Wang committed
152

153
154
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdge")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
155
    GraphHandle ghandle = args[0];
156
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
157
158
159
160
    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
161

162
163
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
164
    GraphHandle ghandle = args[0];
165
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
166
167
    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
Minjie Wang's avatar
Minjie Wang committed
168
169
170
    gptr->AddEdges(src, dst);
  });

171
172
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphClear")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
173
    GraphHandle ghandle = args[0];
174
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
175
176
    gptr->Clear();
  });
Minjie Wang's avatar
Minjie Wang committed
177

178
179
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsMultigraph")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
180
181
    GraphHandle ghandle = args[0];
    // NOTE: not const since we have caches
182
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
183
184
185
    *rv = gptr->IsMultigraph();
  });

186
187
188
189
190
191
192
193
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();
  });

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

201
202
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumEdges")
.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
207
    *rv = static_cast<int64_t>(gptr->NumEdges());
  });

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

216
217
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
218
    GraphHandle ghandle = args[0];
219
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
220
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
Minjie Wang's avatar
Minjie Wang committed
221
222
223
    *rv = gptr->HasVertices(vids);
  });

224
225
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLMapSubgraphNID")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
226
227
228
229
230
    const IdArray parent_vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[0]));
    const IdArray query = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = GraphOp::MapParentIdToSubgraphId(parent_vids, query);
  });

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

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

249
250
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphPredecessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
251
    GraphHandle ghandle = args[0];
252
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
253
254
255
256
257
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
    *rv = gptr->Predecessors(vid, radius);
  });

258
259
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphSuccessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
260
    GraphHandle ghandle = args[0];
261
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
262
263
264
265
266
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
    *rv = gptr->Successors(vid, radius);
  });

267
268
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeId")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
269
    GraphHandle ghandle = args[0];
270
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
271
272
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
273
    *rv = gptr->EdgeId(src, dst);
Minjie Wang's avatar
Minjie Wang committed
274
275
  });

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

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

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

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

309
310
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_1")
.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);
Minjie Wang's avatar
Minjie Wang committed
313
    const dgl_id_t vid = args[1];
314
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
315
316
  });

317
318
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
319
    GraphHandle ghandle = args[0];
320
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
321
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
322
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
323
324
  });

325
326
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
327
    GraphHandle ghandle = args[0];
328
329
330
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
    std::string order = args[1];
    *rv = ConvertEdgeArrayToPackedFunc(gptr->Edges(order));
Minjie Wang's avatar
Minjie Wang committed
331
332
  });

333
334
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
335
    GraphHandle ghandle = args[0];
336
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
337
338
339
340
    const dgl_id_t vid = args[1];
    *rv = static_cast<int64_t>(gptr->InDegree(vid));
  });

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

349
350
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegree")
.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
354
355
356
    const dgl_id_t vid = args[1];
    *rv = static_cast<int64_t>(gptr->OutDegree(vid));
  });

357
358
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegrees")
.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);
Minjie Wang's avatar
Minjie Wang committed
361
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
Minjie Wang's avatar
Minjie Wang committed
362
363
    *rv = gptr->OutDegrees(vids);
  });
Minjie Wang's avatar
Minjie Wang committed
364

365
366
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphVertexSubgraph")
.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
370
371
372
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = ConvertSubgraphToPackedFunc(gptr->VertexSubgraph(vids));
  });

373
374
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
375
    GraphHandle ghandle = args[0];
376
    const GraphInterface *gptr = static_cast<GraphInterface*>(ghandle);
377
378
379
380
    const IdArray eids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = ConvertSubgraphToPackedFunc(gptr->EdgeSubgraph(eids));
  });

381
382
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointUnion")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
383
384
385
386
387
    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) {
388
389
390
      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";
391
392
393
      graphs.push_back(gr);
    }
    Graph* gptr = new Graph();
Minjie Wang's avatar
Minjie Wang committed
394
    *gptr = GraphOp::DisjointUnion(std::move(graphs));
395
396
397
398
    GraphHandle ghandle = gptr;
    *rv = ghandle;
  });

399
400
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionByNum")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
401
    GraphHandle ghandle = args[0];
402
403
404
    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
405
406
407
408
409
410
411
412
413
414
415
416
417
418
    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;
  });

419
420
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionBySizes")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
421
    GraphHandle ghandle = args[0];
422
423
424
    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
425
426
427
428
429
430
431
432
433
434
435
436
437
    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
438

439
440
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphLineGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
GaiYu0's avatar
cpp lg  
GaiYu0 committed
441
442
    GraphHandle ghandle = args[0];
    bool backtracking = args[1];
443
444
445
    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
446
447
448
449
450
    Graph* lgptr = new Graph();
    *lgptr = GraphOp::LineGraph(gptr, backtracking);
    GraphHandle lghandle = lgptr;
    *rv = lghandle;
  });
GaiYu0's avatar
GaiYu0 committed
451

452
453
454
455
456
457
458
459
460
461
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);
  });

462
463
464
465
466
467
468
469
470
471
472
473
474
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
475
}  // namespace dgl