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

14
15
16
17
18
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
19
20

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

22
23
24
///////////////////////////// Graph API ///////////////////////////////////

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreateMutable")
25
.set_body([] (DGLArgs args, DGLRetValue* rv) {
26
27
    bool multigraph = args[0];
    *rv = GraphRef(Graph::Create(multigraph));
Minjie Wang's avatar
Minjie Wang committed
28
  });
Minjie Wang's avatar
Minjie Wang committed
29

30

31
32
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
33
34
    const IdArray src_ids = args[0];
    const IdArray dst_ids = args[1];
35
36
37
    const int multigraph = args[2];
    const int64_t num_nodes = args[3];
    const bool readonly = args[4];
38
    if (readonly) {
39
      if (multigraph == kBoolUnknown) {
40
        *rv = GraphRef(ImmutableGraph::CreateFromCOO(num_nodes, src_ids, dst_ids));
41
      } else {
42
        *rv = GraphRef(ImmutableGraph::CreateFromCOO(num_nodes, src_ids, dst_ids, multigraph));
43
      }
44
    } else {
45
      CHECK_NE(multigraph, kBoolUnknown);
46
      *rv = GraphRef(Graph::CreateFromCOO(num_nodes, src_ids, dst_ids, multigraph));
47
    }
48
49
  });

50
51
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCSRCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
52
53
    const IdArray indptr = args[0];
    const IdArray indices = args[1];
54
    const std::string shared_mem_name = args[2];
55
    const int multigraph = args[3];
56
57
58
59
60
61
62
    const std::string edge_dir = args[4];

    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;
63
64
    if (shared_mem_name.empty()) {
      if (multigraph == kBoolUnknown) {
65
        *rv = GraphRef(ImmutableGraph::CreateFromCSR(indptr, indices, edge_ids, edge_dir));
66
      } else {
67
68
        *rv = GraphRef(ImmutableGraph::CreateFromCSR(
            indptr, indices, edge_ids, multigraph, edge_dir));
69
70
71
      }
    } else {
      if (multigraph == kBoolUnknown) {
72
73
        *rv = GraphRef(ImmutableGraph::CreateFromCSR(
            indptr, indices, edge_ids, edge_dir, shared_mem_name));
74
      } else {
75
76
        *rv = GraphRef(ImmutableGraph::CreateFromCSR(indptr, indices, edge_ids,
            multigraph, edge_dir, shared_mem_name));
77
78
      }
    }
79
80
81
82
83
84
85
  });

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];
86
    const bool multigraph = args[3];
87
    const std::string edge_dir = args[4];
88
    // TODO(minjie): how to know multigraph
89
    *rv = GraphRef(ImmutableGraph::CreateFromCSR(
90
      shared_mem_name, num_vertices, num_edges, multigraph, edge_dir));
Minjie Wang's avatar
Minjie Wang committed
91
  });
Minjie Wang's avatar
Minjie Wang committed
92

93
94
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
95
    GraphRef g = args[0];
Minjie Wang's avatar
Minjie Wang committed
96
    uint64_t num_vertices = args[1];
97
    g->AddVertices(num_vertices);
Minjie Wang's avatar
Minjie Wang committed
98
  });
Minjie Wang's avatar
Minjie Wang committed
99

100
101
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdge")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
102
    GraphRef g = args[0];
Minjie Wang's avatar
Minjie Wang committed
103
104
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
105
    g->AddEdge(src, dst);
Minjie Wang's avatar
Minjie Wang committed
106
  });
Minjie Wang's avatar
Minjie Wang committed
107

108
109
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
110
    GraphRef g = args[0];
111
112
    const IdArray src = args[1];
    const IdArray dst = args[2];
113
    g->AddEdges(src, dst);
Minjie Wang's avatar
Minjie Wang committed
114
115
  });

116
117
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphClear")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
118
119
    GraphRef g = args[0];
    g->Clear();
Minjie Wang's avatar
Minjie Wang committed
120
  });
Minjie Wang's avatar
Minjie Wang committed
121

122
123
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsMultigraph")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
124
125
    GraphRef g = args[0];
    *rv = g->IsMultigraph();
126
127
  });

128
129
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsReadonly")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
130
131
    GraphRef g = args[0];
    *rv = g->IsReadonly();
132
133
  });

134
135
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
136
137
    GraphRef g = args[0];
    *rv = static_cast<int64_t>(g->NumVertices());
Minjie Wang's avatar
Minjie Wang committed
138
  });
Minjie Wang's avatar
Minjie Wang committed
139

140
141
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
142
143
    GraphRef g = args[0];
    *rv = static_cast<int64_t>(g->NumEdges());
Minjie Wang's avatar
Minjie Wang committed
144
145
  });

146
147
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertex")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
148
    GraphRef g = args[0];
Minjie Wang's avatar
Minjie Wang committed
149
    const dgl_id_t vid = args[1];
150
    *rv = g->HasVertex(vid);
Minjie Wang's avatar
Minjie Wang committed
151
152
  });

153
154
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
155
    GraphRef g = args[0];
156
    const IdArray vids = args[1];
157
    *rv = g->HasVertices(vids);
158
159
  });

160
161
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgeBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
162
    GraphRef g = args[0];
Minjie Wang's avatar
Minjie Wang committed
163
164
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
165
    *rv = g->HasEdgeBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
166
167
  });

168
169
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgesBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
170
    GraphRef g = args[0];
171
172
    const IdArray src = args[1];
    const IdArray dst = args[2];
173
    *rv = g->HasEdgesBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
174
175
  });

176
177
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphPredecessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
178
    GraphRef g = args[0];
Minjie Wang's avatar
Minjie Wang committed
179
180
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
181
    *rv = g->Predecessors(vid, radius);
Minjie Wang's avatar
Minjie Wang committed
182
183
  });

184
185
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphSuccessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
186
    GraphRef g = args[0];
Minjie Wang's avatar
Minjie Wang committed
187
188
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
189
    *rv = g->Successors(vid, radius);
Minjie Wang's avatar
Minjie Wang committed
190
191
  });

192
193
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeId")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
194
    GraphRef g = args[0];
Minjie Wang's avatar
Minjie Wang committed
195
196
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
197
    *rv = g->EdgeId(src, dst);
Minjie Wang's avatar
Minjie Wang committed
198
199
  });

200
201
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeIds")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
202
    GraphRef g = args[0];
203
204
    const IdArray src = args[1];
    const IdArray dst = args[2];
205
    *rv = ConvertEdgeArrayToPackedFunc(g->EdgeIds(src, dst));
206
207
  });

208
209
210
211
212
213
214
215
216
217
218
219
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFindEdge")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphRef g = args[0];
    const dgl_id_t eid = args[1];
    const auto& pair = g->FindEdge(eid);
    *rv = PackedFunc([pair] (DGLArgs args, DGLRetValue* rv) {
        const int choice = args[0];
        const int64_t ret = (choice == 0? pair.first : pair.second);
        *rv = ret;
      });
  });

220
221
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFindEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
222
    GraphRef g = args[0];
223
    const IdArray eids = args[1];
224
    *rv = ConvertEdgeArrayToPackedFunc(g->FindEdges(eids));
Minjie Wang's avatar
Minjie Wang committed
225
226
  });

227
228
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
229
    GraphRef g = args[0];
Minjie Wang's avatar
Minjie Wang committed
230
    const dgl_id_t vid = args[1];
231
    *rv = ConvertEdgeArrayToPackedFunc(g->InEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
232
233
  });

234
235
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
236
    GraphRef g = args[0];
237
    const IdArray vids = args[1];
238
    *rv = ConvertEdgeArrayToPackedFunc(g->InEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
239
240
  });

241
242
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
243
    GraphRef g = args[0];
Minjie Wang's avatar
Minjie Wang committed
244
    const dgl_id_t vid = args[1];
245
    *rv = ConvertEdgeArrayToPackedFunc(g->OutEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
246
247
  });

248
249
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
250
    GraphRef g = args[0];
251
    const IdArray vids = args[1];
252
    *rv = ConvertEdgeArrayToPackedFunc(g->OutEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
253
254
  });

255
256
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
257
    GraphRef g = args[0];
258
    std::string order = args[1];
259
    *rv = ConvertEdgeArrayToPackedFunc(g->Edges(order));
Minjie Wang's avatar
Minjie Wang committed
260
261
  });

262
263
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
264
    GraphRef g = args[0];
Minjie Wang's avatar
Minjie Wang committed
265
    const dgl_id_t vid = args[1];
266
    *rv = static_cast<int64_t>(g->InDegree(vid));
Minjie Wang's avatar
Minjie Wang committed
267
268
  });

269
270
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
271
    GraphRef g = args[0];
272
    const IdArray vids = args[1];
273
    *rv = g->InDegrees(vids);
Minjie Wang's avatar
Minjie Wang committed
274
275
  });

276
277
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
278
    GraphRef g = args[0];
Minjie Wang's avatar
Minjie Wang committed
279
    const dgl_id_t vid = args[1];
280
    *rv = static_cast<int64_t>(g->OutDegree(vid));
Minjie Wang's avatar
Minjie Wang committed
281
282
  });

283
284
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
285
    GraphRef g = args[0];
286
    const IdArray vids = args[1];
287
    *rv = g->OutDegrees(vids);
Minjie Wang's avatar
Minjie Wang committed
288
  });
Minjie Wang's avatar
Minjie Wang committed
289

290
291
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphVertexSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
292
    GraphRef g = args[0];
293
    const IdArray vids = args[1];
Da Zheng's avatar
Da Zheng committed
294
295
    std::shared_ptr<Subgraph> subg(new Subgraph(g->VertexSubgraph(vids)));
    *rv = SubgraphRef(subg);
Minjie Wang's avatar
Minjie Wang committed
296
297
  });

298
299
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
300
    GraphRef g = args[0];
301
    const IdArray eids = args[1];
302
    bool preserve_nodes = args[2];
Da Zheng's avatar
Da Zheng committed
303
304
305
    std::shared_ptr<Subgraph> subg(
        new Subgraph(g->EdgeSubgraph(eids, preserve_nodes)));
    *rv = SubgraphRef(subg);
GaiYu0's avatar
cpp lg  
GaiYu0 committed
306
  });
GaiYu0's avatar
GaiYu0 committed
307

308
309
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphGetAdj")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
310
    GraphRef g = args[0];
311
312
    bool transpose = args[1];
    std::string format = args[2];
313
    auto res = g->GetAdj(transpose, format);
314
    *rv = ConvertNDArrayVectorToPackedFunc(res);
315
316
  });

317
318
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphContext")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
319
320
    GraphRef g = args[0];
    *rv = g->Context();
321
322
  });

323
324
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumBits")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
325
326
    GraphRef g = args[0];
    *rv = g->NumBits();
327
328
  });

Da Zheng's avatar
Da Zheng committed
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
// Subgraph C APIs

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLSubgraphGetGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    SubgraphRef subg = args[0];
    *rv = GraphRef(subg->graph);
  });

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLSubgraphGetInducedVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    SubgraphRef subg = args[0];
    *rv = subg->induced_vertices;
  });

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLSubgraphGetInducedEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    SubgraphRef subg = args[0];
    *rv = subg->induced_edges;
  });

Da Zheng's avatar
Da Zheng committed
349
350
351
352
353
354
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLSortAdj")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphRef g = args[0];
    g->SortCSR();
  });

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