remap vertex ids.

dmlc · Jan 7, 2019 · e161e45 · e161e45
1 parent 5ead4cd
commit e161e45
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 5 deletions.
diff --git a/include/dgl/immutable_graph.h b/include/dgl/immutable_graph.h
@@ -408,6 +408,8 @@ class ImmutableGraph {
                                  const std::string &neigh_type,
                                  int num_hops, size_t num_neighbor) const;
 
+  void CompactSubgraph(IdArray induced_vertices);
+
   // Store the in-edges.
   std::shared_ptr<csr> in_csr_;
   // Store the out-edges.

diff --git a/src/graph/immutable_graph.cc b/src/graph/immutable_graph.cc
@@ -891,14 +891,38 @@ SampledSubgraph ImmutableGraph::SampleSubgraph(IdArray seed_arr,
   return subg;
 }
 
+void CompactSubgraph(ImmutableGraph::csr &subg,
+                     const std::unordered_map<dgl_id_t, dgl_id_t> &id_map) {
+  for (size_t i = 0; i < subg.indices.size(); i++) {
+    auto it = id_map.find(subg.indices[i]);
+    assert(it != id_map.end());
+    subg.indices[i] = it->second;
+  }
+}
+
+void ImmutableGraph::CompactSubgraph(IdArray induced_vertices) {
+  // The key is the old id, the value is the id in the subgraph.
+  std::unordered_map<dgl_id_t, dgl_id_t> id_map;
+  const dgl_id_t *vdata = static_cast<dgl_id_t *>(induced_vertices->data);
+  size_t len = induced_vertices->shape[0];
+  for (size_t i = 0; i < len; i++)
+    id_map.insert(std::pair<dgl_id_t, dgl_id_t>(vdata[i], i));
+  if (in_csr_)
+    dgl::CompactSubgraph(*in_csr_, id_map);
+  if (out_csr_)
+    dgl::CompactSubgraph(*out_csr_, id_map);
+}
+
 SampledSubgraph ImmutableGraph::NeighborUniformSample(IdArray seeds,
                                                       const std::string &neigh_type,
                                                       int num_hops, int expand_factor) const {
-  return SampleSubgraph(seeds,                 // seed vector
-                        nullptr,               // sample_id_probability
-                        neigh_type,
-                        num_hops,
-                        expand_factor);
+  auto ret = SampleSubgraph(seeds,                 // seed vector
+                            nullptr,               // sample_id_probability
+                            neigh_type,
+                            num_hops,
+                            expand_factor);
+  ret.graph.CompactSubgraph(ret.induced_vertices);
+  return ret;
 }
 
 }  // namespace dgl