[Hetero] Heterograph C++ implementation; Bipartite and Python wrapper (#725)

* finish bipartite graph implementation; compiled

* finished heterograph implementation; compiled

* WIP: apis

* C API codes

* compiled

* WIP: python

* HeteroGraphIndex

* WIP: test

* add DGLContext support in ffi

* fix bug in has edge

* unittests except edge subgraph

* edge subgraph

* fix lint

* address comments

* poke ci

* try fix

* fix msvc

include/dgl/heterograph_interface.h → include/dgl/base_heterograph.h

@@ -4,27 +4,28 @@
  * \brief DGL heterogeneous graph index class.
  */
 
-#ifndef DGL_HETEROGRAPH_INTERFACE_H_
-#define DGL_HETEROGRAPH_INTERFACE_H_
+#ifndef DGL_BASE_HETEROGRAPH_H_
+#define DGL_BASE_HETEROGRAPH_H_
 
 #include <string>
 #include <vector>
 #include <utility>
 #include <algorithm>
+#include <memory>
 
+#include "./runtime/object.h"
 #include "graph_interface.h"
 #include "array.h"
 
 namespace dgl {
 
 // Forward declaration
+class BaseHeteroGraph;
+typedef std::shared_ptr<BaseHeteroGraph> HeteroGraphPtr;
 struct HeteroSubgraph;
-class HeteroGraphInterface;
-
-typedef std::shared_ptr<HeteroGraphInterface> HeteroGraphPtr;
 
 /*!
- * \brief Heterogenous graph APIs
+ * \brief Base heterogenous graph.
  *
  * In heterograph, nodes represent entities and edges represent relations.
  * Nodes and edges are associated with types. The same pair of entity types
@@ -36,9 +37,11 @@ typedef std::shared_ptr<HeteroGraphInterface> HeteroGraphPtr;
  *  - A dictionary of relation type to the bipartite graph representing the
  *    actual connections among entity nodes.
  */
-class HeteroGraphInterface {
+class BaseHeteroGraph : public runtime::Object {
  public:
-  virtual ~HeteroGraphInterface() = default;
+  explicit BaseHeteroGraph(GraphPtr meta_graph): meta_graph_(meta_graph) {}
+
+  virtual ~BaseHeteroGraph() = default;
 
   ////////////////////////// query/operations on meta graph ////////////////////////
 
@@ -49,14 +52,16 @@ class HeteroGraphInterface {
   virtual uint64_t NumEdgeTypes() const = 0;
 
   /*! \return the meta graph */
-  virtual const GraphInterface& GetMetaGraph() const = 0;
+  virtual GraphPtr meta_graph() const {
+    return meta_graph_;
+  }
 
   /*!
    * \brief Return the bipartite graph of the given edge type.
    * \param etype The edge type.
    * \return The bipartite graph.
    */
-  virtual const HeteroGraphInterface& GetRelationGraph(dgl_type_t etype) const = 0;
+  virtual HeteroGraphPtr GetRelationGraph(dgl_type_t etype) const = 0;
 
   ////////////////////////// query/operations on realized graph ////////////////////////
 
@@ -102,48 +107,13 @@ class HeteroGraphInterface {
   virtual bool HasVertex(dgl_type_t vtype, dgl_id_t vid) const = 0;
 
   /*! \return a 0-1 array indicating whether the given vertices are in the graph.*/
-  virtual BoolArray HasVertices(dgl_type_t vtype, IdArray vids) const {
-    const auto len = vids->shape[0];
-    BoolArray rst = aten::NewBoolArray(len);
-    const dgl_id_t* vid_data = static_cast<dgl_id_t*>(vids->data);
-    dgl_id_t* rst_data = static_cast<dgl_id_t*>(rst->data);
-    for (int64_t i = 0; i < len; ++i) {
-      rst_data[i] = HasVertex(vtype, vid_data[i])? 1 : 0;
-    }
-    return rst;
-  }
+  virtual BoolArray HasVertices(dgl_type_t vtype, IdArray vids) const = 0;
 
   /*! \return true if the given edge is in the graph.*/
   virtual bool HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const = 0;
 
   /*! \return a 0-1 array indicating whether the given edges are in the graph.*/
-  virtual BoolArray HasEdgesBetween(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) const {
-    const auto srclen = src_ids->shape[0];
-    const auto dstlen = dst_ids->shape[0];
-    const auto rstlen = std::max(srclen, dstlen);
-    BoolArray rst = aten::NewBoolArray(rstlen);
-    dgl_id_t* rst_data = static_cast<dgl_id_t*>(rst->data);
-    const dgl_id_t* src_data = static_cast<dgl_id_t*>(src_ids->data);
-    const dgl_id_t* dst_data = static_cast<dgl_id_t*>(dst_ids->data);
-    if (srclen == 1) {
-      // one-many
-      for (int64_t i = 0; i < dstlen; ++i) {
-        rst_data[i] = HasEdgeBetween(etype, src_data[0], dst_data[i])? 1 : 0;
-      }
-    } else if (dstlen == 1) {
-      // many-one
-      for (int64_t i = 0; i < srclen; ++i) {
-        rst_data[i] = HasEdgeBetween(etype, src_data[i], dst_data[0])? 1 : 0;
-      }
-    } else {
-      // many-many
-      CHECK(srclen == dstlen) << "Invalid src and dst id array.";
-      for (int64_t i = 0; i < srclen; ++i) {
-        rst_data[i] = HasEdgeBetween(etype, src_data[i], dst_data[i])? 1 : 0;
-      }
-    }
-    return rst;
-  }
+  virtual BoolArray HasEdgesBetween(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) const = 0;
 
   /*!
    * \brief Find the predecessors of a vertex.
@@ -350,7 +320,7 @@ class HeteroGraphInterface {
    * \return a vector of IdArrays.
    */
   virtual std::vector<IdArray> GetAdj(
-      dgl_id_t etype, bool transpose, const std::string &fmt) const = 0;
+      dgl_type_t etype, bool transpose, const std::string &fmt) const = 0;
 
   /*!
    * \brief Extract the induced subgraph by the given vertices.
@@ -380,10 +350,20 @@ class HeteroGraphInterface {
    */
   virtual HeteroSubgraph EdgeSubgraph(
       const std::vector<IdArray>& eids, bool preserve_nodes = false) const = 0;
+
+  static constexpr const char* _type_key = "graph.HeteroGraph";
+  DGL_DECLARE_OBJECT_TYPE_INFO(BaseHeteroGraph, runtime::Object);
+
+ protected:
+  /*! \brief meta graph */
+  GraphPtr meta_graph_;
 };
 
+// Define HeteroGraphRef
+DGL_DEFINE_OBJECT_REF(HeteroGraphRef, BaseHeteroGraph);
+
 /*! \brief Heter-subgraph data structure */
-struct HeteroSubgraph {
+struct HeteroSubgraph : public runtime::Object {
   /*! \brief The heterograph. */
   HeteroGraphPtr graph;
   /*!
@@ -396,8 +376,29 @@ struct HeteroSubgraph {
    * The vector length is equal to the number of vertex types in the parent graph.
    */
   std::vector<IdArray> induced_edges;
+
+  static constexpr const char* _type_key = "graph.HeteroSubgraph";
+  DGL_DECLARE_OBJECT_TYPE_INFO(HeteroSubgraph, runtime::Object);
 };
 
+// Define HeteroSubgraphRef
+DGL_DEFINE_OBJECT_REF(HeteroSubgraphRef, HeteroSubgraph);
+
+// creators
+
+/*! \brief Create a bipartite graph from COO arrays */
+HeteroGraphPtr CreateBipartiteFromCOO(
+    int64_t num_src, int64_t num_dst, IdArray row, IdArray col);
+
+/*! \brief Create a bipartite graph from (out) CSR arrays */
+HeteroGraphPtr CreateBipartiteFromCSR(
+    int64_t num_src, int64_t num_dst,
+    IdArray indptr, IdArray indices, IdArray edge_ids);
+
+/*! \brief Create a heterograph from meta graph and a list of bipartite graph */
+HeteroGraphPtr CreateHeteroGraph(
+    GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& rel_graphs);
+
 };  // namespace dgl
 
-#endif  // DGL_HETEROGRAPH_INTERFACE_H_
+#endif  // DGL_BASE_HETEROGRAPH_H_

python/dgl/_ffi/_ctypes/types.py

@@ -63,7 +63,8 @@ RETURN_SWITCH = {
     TypeCode.HANDLE: _return_handle,
     TypeCode.NULL: lambda x: None,
     TypeCode.STR: lambda x: py_str(x.v_str),
-    TypeCode.BYTES: _return_bytes
+    TypeCode.BYTES: _return_bytes,
+    TypeCode.DGL_CONTEXT: lambda x: DGLContext(x.v_ctx.device_type, x.v_ctx.device_id),
 }
 
 C_TO_PY_ARG_SWITCH = {
@@ -72,5 +73,6 @@ C_TO_PY_ARG_SWITCH = {
     TypeCode.HANDLE: _return_handle,
     TypeCode.NULL: lambda x: None,
     TypeCode.STR: lambda x: py_str(x.v_str),
-    TypeCode.BYTES: _return_bytes
+    TypeCode.BYTES: _return_bytes,
+    TypeCode.DGL_CONTEXT: lambda x: DGLContext(x.v_ctx.device_type, x.v_ctx.device_id),
 }

python/dgl/container.py

@@ -69,3 +69,11 @@ class StrMap(Map):
         """Get the items from the map"""
         akvs = _api_internal._MapItems(self)
         return [(akvs[i].value, akvs[i+1]) for i in range(0, len(akvs), 2)]
+
+@register_object
+class Value(ObjectBase):
+    """Object wrapper for various values."""
+    @property
+    def data(self):
+        """Return the value data."""
+        return _api_internal._ValueGet(self)

src/c_api_common.cc

@@ -3,6 +3,7 @@
  * \file c_runtime_api.cc
  * \brief DGL C API common implementations
  */
+#include <dgl/graph_interface.h>
 #include "c_api_common.h"
 
 using dgl::runtime::DGLArgs;
@@ -25,4 +26,20 @@ PackedFunc ConvertNDArrayVectorToPackedFunc(const std::vector<NDArray>& vec) {
     return PackedFunc(body);
 }
 
+PackedFunc ConvertEdgeArrayToPackedFunc(const EdgeArray& ea) {
+  auto body = [ea] (DGLArgs args, DGLRetValue* rv) {
+      const int which = args[0];
+      if (which == 0) {
+        *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);
+}
+
 }  // namespace dgl

src/c_api_common.h

@@ -9,6 +9,8 @@
 #include <dgl/runtime/ndarray.h>
 #include <dgl/runtime/packed_func.h>
 #include <dgl/runtime/registry.h>
+#include <dgl/array.h>
+#include <dgl/graph_interface.h>
 #include <algorithm>
 #include <vector>
 
@@ -70,6 +72,8 @@ dgl::runtime::NDArray CopyVectorToNDArray(
   return a;
 }
 
+runtime::PackedFunc ConvertEdgeArrayToPackedFunc(const EdgeArray& ea);
+
 }  // namespace dgl
 
 #endif  // DGL_C_API_COMMON_H_

src/graph/bipartite.h

+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file graph/bipartite.h
+ * \brief Bipartite graph
+ */
+
+#ifndef DGL_GRAPH_BIPARTITE_H_
+#define DGL_GRAPH_BIPARTITE_H_
+
+#include <dgl/graph_interface.h>
+#include <dgl/base_heterograph.h>
+#include <vector>
+#include <string>
+#include <utility>
+#include <memory>
+
+namespace dgl {
+
+/*!
+ * \brief Bipartite graph
+ *
+ * Bipartite graph is a special type of heterograph which has two types
+ * of nodes: "Src" and "Dst". All the edges are from "Src" type nodes to
+ * "Dst" type nodes, so there is no edge among nodes of the same type.
+ */
+class Bipartite : public BaseHeteroGraph {
+ public:
+  /*! \brief source node group type */
+  static constexpr dgl_type_t kSrcVType = 0;
+  /*! \brief destination node group type */
+  static constexpr dgl_type_t kDstVType = 1;
+  /*! \brief edge group type */
+  static constexpr dgl_type_t kEType = 0;
+
+  uint64_t NumVertexTypes() const override {
+    return 2;
+  }
+
+  uint64_t NumEdgeTypes() const override {
+    return 1;
+  }
+
+  HeteroGraphPtr GetRelationGraph(dgl_type_t etype) const override {
+    LOG(FATAL) << "The method shouldn't be called for Bipartite graph. "
+      << "The relation graph is simply this graph itself.";
+    return {};
+  }
+
+  void AddVertices(dgl_type_t vtype, uint64_t num_vertices) override {
+    LOG(FATAL) << "Bipartite graph is not mutable.";
+  }
+
+  void AddEdge(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) override {
+    LOG(FATAL) << "Bipartite graph is not mutable.";
+  }
+
+  void AddEdges(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) override {
+    LOG(FATAL) << "Bipartite graph is not mutable.";
+  }
+
+  void Clear() override {
+    LOG(FATAL) << "Bipartite graph is not mutable.";
+  }
+
+  DLContext Context() const override;
+
+  uint8_t NumBits() const override;
+
+  bool IsMultigraph() const override;
+
+  bool IsReadonly() const override {
+    return true;
+  }
+
+  uint64_t NumVertices(dgl_type_t vtype) const override;
+
+  uint64_t NumEdges(dgl_type_t etype) const override;
+
+  bool HasVertex(dgl_type_t vtype, dgl_id_t vid) const override;
+
+  BoolArray HasVertices(dgl_type_t vtype, IdArray vids) const override;
+
+  bool HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override;
+
+  BoolArray HasEdgesBetween(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) const override;
+
+  IdArray Predecessors(dgl_type_t etype, dgl_id_t dst) const override;
+
+  IdArray Successors(dgl_type_t etype, dgl_id_t src) const override;
+
+  IdArray EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override;
+
+  EdgeArray EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const override;
+
+  std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_type_t etype, dgl_id_t eid) const override;
+
+  EdgeArray FindEdges(dgl_type_t etype, IdArray eids) const override;
+
+  EdgeArray InEdges(dgl_type_t etype, dgl_id_t vid) const override;
+
+  EdgeArray InEdges(dgl_type_t etype, IdArray vids) const override;
+
+  EdgeArray OutEdges(dgl_type_t etype, dgl_id_t vid) const override;
+
+  EdgeArray OutEdges(dgl_type_t etype, IdArray vids) const override;
+
+  EdgeArray Edges(dgl_type_t etype, const std::string &order = "") const override;
+
+  uint64_t InDegree(dgl_type_t etype, dgl_id_t vid) const override;
+
+  DegreeArray InDegrees(dgl_type_t etype, IdArray vids) const override;
+
+  uint64_t OutDegree(dgl_type_t etype, dgl_id_t vid) const override;
+
+  DegreeArray OutDegrees(dgl_type_t etype, IdArray vids) const override;
+
+  DGLIdIters SuccVec(dgl_type_t etype, dgl_id_t vid) const override;
+
+  DGLIdIters OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const override;
+
+  DGLIdIters PredVec(dgl_type_t etype, dgl_id_t vid) const override;
+
+  DGLIdIters InEdgeVec(dgl_type_t etype, dgl_id_t vid) const override;
+
+  std::vector<IdArray> GetAdj(
+      dgl_type_t etype, bool transpose, const std::string &fmt) const override;
+
+  HeteroSubgraph VertexSubgraph(const std::vector<IdArray>& vids) const override;
+
+  HeteroSubgraph EdgeSubgraph(
+      const std::vector<IdArray>& eids, bool preserve_nodes = false) const override;
+
+  // creators
+  /*! \brief Create a bipartite graph from COO arrays */
+  static HeteroGraphPtr CreateFromCOO(int64_t num_src, int64_t num_dst,
+      IdArray row, IdArray col);
+
+  /*! \brief Create a bipartite graph from (out) CSR arrays */
+  static HeteroGraphPtr CreateFromCSR(
+      int64_t num_src, int64_t num_dst,
+      IdArray indptr, IdArray indices, IdArray edge_ids);
+
+ private:
+  // internal data structure
+  class COO;
+  class CSR;
+  typedef std::shared_ptr<COO> COOPtr;
+  typedef std::shared_ptr<CSR> CSRPtr;
+
+  Bipartite(CSRPtr in_csr, CSRPtr out_csr, COOPtr coo);
+
+  /*! \return Return the in-edge CSR format. Create from other format if not exist. */
+  CSRPtr GetInCSR() const;
+
+  /*! \return Return the out-edge CSR format. Create from other format if not exist. */
+  CSRPtr GetOutCSR() const;
+
+  /*! \return Return the COO format. Create from other format if not exist. */
+  COOPtr GetCOO() const;
+
+  /*! \return Return any existing format. */
+  HeteroGraphPtr GetAny() const;
+
+  // Graph stored in different format. We use an on-demand strategy: the format is
+  // only materialized if the operation that suitable for it is invoked.
+  /*! \brief CSR graph that stores reverse edges */
+  CSRPtr in_csr_;
+  /*! \brief CSR representation */
+  CSRPtr out_csr_;
+  /*! \brief COO representation */
+  COOPtr coo_;
+};
+
+};  // namespace dgl
+
+#endif  // DGL_GRAPH_BIPARTITE_H_

src/graph/graph_apis.cc

@@ -20,36 +20,6 @@ using dgl::runtime::NDArray;
 namespace dgl {
 
 namespace {
-// Convert EdgeArray structure to PackedFunc.
-template<class EdgeArray>
-PackedFunc ConvertEdgeArrayToPackedFunc(const EdgeArray& ea) {
-  auto body = [ea] (DGLArgs args, DGLRetValue* rv) {
-      const int which = args[0];
-      if (which == 0) {
-        *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);
-}
-
-// 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);
-}
 
 // Convert Subgraph structure to PackedFunc.
 PackedFunc ConvertSubgraphToPackedFunc(const Subgraph& sg) {
@@ -256,6 +226,18 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeIds")
     *rv = ConvertEdgeArrayToPackedFunc(g->EdgeIds(src, dst));
   });
 
+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;
+      });
+  });
+
 DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFindEdges")
 .set_body([] (DGLArgs args, DGLRetValue* rv) {
     GraphRef g = args[0];
@@ -347,7 +329,7 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphGetAdj")
     bool transpose = args[1];
     std::string format = args[2];
     auto res = g->GetAdj(transpose, format);
-    *rv = ConvertAdjToPackedFunc(res);
+    *rv = ConvertNDArrayVectorToPackedFunc(res);
   });
 
 DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphContext")

src/graph/heterograph.cc

+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file graph/heterograph.cc
+ * \brief Heterograph implementation
+ */
+#include "./heterograph.h"
+#include <dgl/packed_func_ext.h>
+#include <dgl/runtime/container.h>
+#include "../c_api_common.h"
+#include "./bipartite.h"
+
+using namespace dgl::runtime;
+
+namespace dgl {
+namespace {
+
+HeteroSubgraph EdgeSubgraphPreserveNodes(
+    const HeteroGraph* hg, const std::vector<IdArray>& eids) {
+  CHECK_EQ(eids.size(), hg->NumEdgeTypes())
+    << "Invalid input: the input list size must be the same as the number of edge type.";
+  HeteroSubgraph ret;
+  ret.induced_vertices.resize(hg->NumVertexTypes());
+  ret.induced_edges = eids;
+  // When preserve_nodes is true, simply compute EdgeSubgraph for each bipartite
+  std::vector<HeteroGraphPtr> subrels(hg->NumEdgeTypes());
+  for (dgl_type_t etype = 0; etype < hg->NumEdgeTypes(); ++etype) {
+    auto pair = hg->meta_graph()->FindEdge(etype);
+    const dgl_type_t src_vtype = pair.first;
+    const dgl_type_t dst_vtype = pair.second;
+    const auto& rel_vsg = hg->GetRelationGraph(etype)->EdgeSubgraph(
+        {eids[etype]}, true);
+    subrels[etype] = rel_vsg.graph;
+    ret.induced_vertices[src_vtype] = rel_vsg.induced_vertices[0];
+    ret.induced_vertices[dst_vtype] = rel_vsg.induced_vertices[1];
+  }
+  ret.graph = HeteroGraphPtr(new HeteroGraph(hg->meta_graph(), subrels));
+  return ret;
+}
+
+HeteroSubgraph EdgeSubgraphNoPreserveNodes(
+    const HeteroGraph* hg, const std::vector<IdArray>& eids) {
+  CHECK_EQ(eids.size(), hg->NumEdgeTypes())
+    << "Invalid input: the input list size must be the same as the number of edge type.";
+  HeteroSubgraph ret;
+  ret.induced_vertices.resize(hg->NumVertexTypes());
+  ret.induced_edges = eids;
+  // NOTE(minjie): EdgeSubgraph when preserve_nodes is false is quite complicated in
+  // heterograph. This is because we need to make sure bipartite graphs that incident
+  // on the same vertex type must have the same ID space. For example, suppose we have
+  // following heterograph:
+  //
+  // Meta graph: A -> B -> C
+  // Bipartite graphs:
+  // * A -> B: (0, 0), (0, 1)
+  // * B -> C: (1, 0), (1, 1)
+  //
+  // Suppose for A->B, we only keep edge (0, 0), while for B->C we only keep (1, 0). We need
+  // to make sure that in the result subgraph, node type B still has two nodes. This means
+  // we cannot simply compute EdgeSubgraph for B->C which will relabel node#1 of type B to be
+  // node #0.
+  //
+  // One implementation is as follows:
+  // (1) For each bipartite graph, slice out the edges using the given eids.
+  // (2) Make a dictionary map<vtype, vector<IdArray>>, where the key is the vertex type
+  //     and the value is the incident nodes from the bipartite graphs that has the vertex
+  //     type as either srctype or dsttype.
+  // (3) Then for each vertex type, use aten::Relabel_ on its vector<IdArray>.
+  //     aten::Relabel_ computes the union of the vertex sets and relabel
+  //     the unique elements from zero. The returned mapping array is the final induced
+  //     vertex set for that vertex type.
+  // (4) Use the relabeled edges to construct the bipartite graph.
+  // step (1) & (2)
+  std::vector<EdgeArray> subedges(hg->NumEdgeTypes());
+  std::vector<std::vector<IdArray>> vtype2incnodes(hg->NumVertexTypes());
+  for (dgl_type_t etype = 0; etype < hg->NumEdgeTypes(); ++etype) {
+    auto pair = hg->meta_graph()->FindEdge(etype);
+    const dgl_type_t src_vtype = pair.first;
+    const dgl_type_t dst_vtype = pair.second;
+    auto earray = hg->GetRelationGraph(etype)->FindEdges(0, eids[etype]);
+    vtype2incnodes[src_vtype].push_back(earray.src);
+    vtype2incnodes[dst_vtype].push_back(earray.dst);
+    subedges[etype] = earray;
+  }
+  // step (3)
+  for (dgl_type_t vtype = 0; vtype < hg->NumVertexTypes(); ++vtype) {
+    ret.induced_vertices[vtype] = aten::Relabel_(vtype2incnodes[vtype]);
+  }
+  // step (4)
+  std::vector<HeteroGraphPtr> subrels(hg->NumEdgeTypes());
+  for (dgl_type_t etype = 0; etype < hg->NumEdgeTypes(); ++etype) {
+    auto pair = hg->meta_graph()->FindEdge(etype);
+    const dgl_type_t src_vtype = pair.first;
+    const dgl_type_t dst_vtype = pair.second;
+    subrels[etype] = Bipartite::CreateFromCOO(
+      ret.induced_vertices[src_vtype]->shape[0],
+      ret.induced_vertices[dst_vtype]->shape[0],
+      subedges[etype].src,
+      subedges[etype].dst);
+  }
+  ret.graph = HeteroGraphPtr(new HeteroGraph(hg->meta_graph(), subrels));
+  return ret;
+}
+
+}  // namespace
+
+HeteroGraph::HeteroGraph(GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& rel_graphs)
+  : BaseHeteroGraph(meta_graph), relation_graphs_(rel_graphs) {
+  // Sanity check
+  CHECK_EQ(meta_graph->NumEdges(), rel_graphs.size());
+  CHECK(!rel_graphs.empty()) << "Empty heterograph is not allowed.";
+  // all relation graph must be bipartite graphs
+  for (const auto rg : rel_graphs) {
+    CHECK_EQ(rg->NumVertexTypes(), 2) << "Each relation graph must be a bipartite graph.";
+    CHECK_EQ(rg->NumEdgeTypes(), 1) << "Each relation graph must be a bipartite graph.";
+  }
+  // create num verts per type
+  num_verts_per_type_.resize(meta_graph_->NumVertices(), -1);
+  for (dgl_type_t vtype = 0; vtype < meta_graph_->NumVertices(); ++vtype) {
+    for (dgl_type_t etype : meta_graph->OutEdgeVec(vtype)) {
+      const auto nv = rel_graphs[etype]->NumVertices(Bipartite::kSrcVType);
+      if (num_verts_per_type_[vtype] < 0) {
+        num_verts_per_type_[vtype] = nv;
+      } else {
+        CHECK_EQ(num_verts_per_type_[vtype], nv)
+          << "Mismatch number of vertices for vertex type " << vtype;
+      }
+    }
+  }
+}
+
+bool HeteroGraph::IsMultigraph() const {
+  return const_cast<HeteroGraph*>(this)->is_multigraph_.Get([this] () {
+      for (const auto hg : relation_graphs_) {
+        if (hg->IsMultigraph()) {
+          return true;
+        }
+      }
+      return false;
+    });
+}
+
+BoolArray HeteroGraph::HasVertices(dgl_type_t vtype, IdArray vids) const {
+  CHECK(IsValidIdArray(vids)) << "Invalid id array input";
+  return aten::LT(vids, NumVertices(vtype));
+}
+
+HeteroSubgraph HeteroGraph::VertexSubgraph(const std::vector<IdArray>& vids) const {
+  CHECK_EQ(vids.size(), NumVertexTypes())
+    << "Invalid input: the input list size must be the same as the number of vertex types.";
+  HeteroSubgraph ret;
+  ret.induced_vertices = vids;
+  ret.induced_edges.resize(NumEdgeTypes());
+  std::vector<HeteroGraphPtr> subrels(NumEdgeTypes());
+  for (dgl_type_t etype = 0; etype < NumEdgeTypes(); ++etype) {
+    auto pair = meta_graph_->FindEdge(etype);
+    const dgl_type_t src_vtype = pair.first;
+    const dgl_type_t dst_vtype = pair.second;
+    const auto& rel_vsg = GetRelationGraph(etype)->VertexSubgraph(
+        {vids[src_vtype], vids[dst_vtype]});
+    subrels[etype] = rel_vsg.graph;
+    ret.induced_edges[etype] = rel_vsg.induced_edges[0];
+  }
+  ret.graph = HeteroGraphPtr(new HeteroGraph(meta_graph_, subrels));
+  return ret;
+}
+
+HeteroSubgraph HeteroGraph::EdgeSubgraph(
+    const std::vector<IdArray>& eids, bool preserve_nodes) const {
+  if (preserve_nodes) {
+    return EdgeSubgraphPreserveNodes(this, eids);
+  } else {
+    return EdgeSubgraphNoPreserveNodes(this, eids);
+  }
+}
+
+// creator implementation
+HeteroGraphPtr CreateBipartiteFromCOO(
+    int64_t num_src, int64_t num_dst, IdArray row, IdArray col) {
+  return Bipartite::CreateFromCOO(num_src, num_dst, row, col);
+}
+
+HeteroGraphPtr CreateBipartiteFromCSR(
+    int64_t num_src, int64_t num_dst,
+    IdArray indptr, IdArray indices, IdArray edge_ids) {
+  return Bipartite::CreateFromCSR(num_src, num_dst, indptr, indices, edge_ids);
+}
+
+HeteroGraphPtr CreateHeteroGraph(
+    GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& rel_graphs) {
+  return HeteroGraphPtr(new HeteroGraph(meta_graph, rel_graphs));
+}
+
+///////////////////////// C APIs /////////////////////////
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroCreateBipartiteFromCOO")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    int64_t num_src = args[0];
+    int64_t num_dst = args[1];
+    IdArray row = args[2];
+    IdArray col = args[3];
+    auto hgptr = CreateBipartiteFromCOO(num_src, num_dst, row, col);
+    *rv = HeteroGraphRef(hgptr);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroCreateBipartiteFromCSR")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    int64_t num_src = args[0];
+    int64_t num_dst = args[1];
+    IdArray indptr = args[2];
+    IdArray indices = args[3];
+    IdArray edge_ids = args[4];
+    auto hgptr = CreateBipartiteFromCSR(num_src, num_dst, indptr, indices, edge_ids);
+    *rv = HeteroGraphRef(hgptr);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroCreateHeteroGraph")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    GraphRef meta_graph = args[0];
+    List<HeteroGraphRef> rel_graphs = args[1];
+    std::vector<HeteroGraphPtr> rel_ptrs;
+    rel_ptrs.reserve(rel_graphs.size());
+    for (const auto& ref : rel_graphs) {
+      rel_ptrs.push_back(ref.sptr());
+    }
+    auto hgptr = CreateHeteroGraph(meta_graph.sptr(), rel_ptrs);
+    *rv = HeteroGraphRef(hgptr);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroGetMetaGraph")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    *rv = GraphRef(hg->meta_graph());
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroGetRelationGraph")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    *rv = HeteroGraphRef(hg->GetRelationGraph(etype));
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroAddVertices")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t vtype = args[1];
+    int64_t num = args[2];
+    hg->AddVertices(vtype, num);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroAddEdge")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    dgl_id_t src = args[2];
+    dgl_id_t dst = args[3];
+    hg->AddEdge(etype, src, dst);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroAddEdges")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    IdArray src = args[2];
+    IdArray dst = args[3];
+    hg->AddEdges(etype, src, dst);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroClear")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    hg->Clear();
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroContext")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    *rv = hg->Context();
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroNumBits")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    *rv = hg->NumBits();
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroIsMultigraph")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    *rv = hg->IsMultigraph();
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroIsReadonly")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    *rv = hg->IsReadonly();
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroNumVertices")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t vtype = args[1];
+    *rv = static_cast<int64_t>(hg->NumVertices(vtype));
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroNumEdges")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    *rv = static_cast<int64_t>(hg->NumEdges(etype));
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroHasVertex")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t vtype = args[1];
+    dgl_id_t vid = args[2];
+    *rv = hg->HasVertex(vtype, vid);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroHasVertices")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t vtype = args[1];
+    IdArray vids = args[2];
+    *rv = hg->HasVertices(vtype, vids);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroHasEdgeBetween")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    dgl_id_t src = args[2];
+    dgl_id_t dst = args[3];
+    *rv = hg->HasEdgeBetween(etype, src, dst);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroHasEdgesBetween")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    IdArray src = args[2];
+    IdArray dst = args[3];
+    *rv = hg->HasEdgesBetween(etype, src, dst);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroPredecessors")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    dgl_id_t dst = args[2];
+    *rv = hg->Predecessors(etype, dst);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroSuccessors")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    dgl_id_t src = args[2];
+    *rv = hg->Successors(etype, src);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroEdgeId")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    dgl_id_t src = args[2];
+    dgl_id_t dst = args[3];
+    *rv = hg->EdgeId(etype, src, dst);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroEdgeIds")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    IdArray src = args[2];
+    IdArray dst = args[3];
+    const auto& ret = hg->EdgeIds(etype, src, dst);
+    *rv = ConvertEdgeArrayToPackedFunc(ret);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroFindEdges")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    IdArray eids = args[2];
+    const auto& ret = hg->FindEdges(etype, eids);
+    *rv = ConvertEdgeArrayToPackedFunc(ret);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroInEdges_1")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    dgl_id_t vid = args[2];
+    const auto& ret = hg->InEdges(etype, vid);
+    *rv = ConvertEdgeArrayToPackedFunc(ret);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroInEdges_2")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    IdArray vids = args[2];
+    const auto& ret = hg->InEdges(etype, vids);
+    *rv = ConvertEdgeArrayToPackedFunc(ret);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroOutEdges_1")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    dgl_id_t vid = args[2];
+    const auto& ret = hg->OutEdges(etype, vid);
+    *rv = ConvertEdgeArrayToPackedFunc(ret);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroOutEdges_2")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    IdArray vids = args[2];
+    const auto& ret = hg->OutEdges(etype, vids);
+    *rv = ConvertEdgeArrayToPackedFunc(ret);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroEdges")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    std::string order = args[2];
+    const auto& ret = hg->Edges(etype, order);
+    *rv = ConvertEdgeArrayToPackedFunc(ret);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroInDegree")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    dgl_id_t vid = args[2];
+    *rv = static_cast<int64_t>(hg->InDegree(etype, vid));
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroInDegrees")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    IdArray vids = args[2];
+    *rv = hg->InDegrees(etype, vids);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroOutDegree")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    dgl_id_t vid = args[2];
+    *rv = static_cast<int64_t>(hg->OutDegree(etype, vid));
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroOutDegrees")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    IdArray vids = args[2];
+    *rv = hg->OutDegrees(etype, vids);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroGetAdj")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    dgl_type_t etype = args[1];
+    bool transpose = args[2];
+    std::string fmt = args[3];
+    *rv = ConvertNDArrayVectorToPackedFunc(
+        hg->GetAdj(etype, transpose, fmt));
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroVertexSubgraph")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    List<Value> vids = args[1];
+    std::vector<IdArray> vid_vec;
+    vid_vec.reserve(vids.size());
+    for (Value val : vids) {
+      vid_vec.push_back(val->data);
+    }
+    std::shared_ptr<HeteroSubgraph> subg(
+        new HeteroSubgraph(hg->VertexSubgraph(vid_vec)));
+    *rv = HeteroSubgraphRef(subg);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroEdgeSubgraph")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroGraphRef hg = args[0];
+    List<Value> eids = args[1];
+    bool preserve_nodes = args[2];
+    std::vector<IdArray> eid_vec;
+    eid_vec.reserve(eids.size());
+    for (Value val : eids) {
+      eid_vec.push_back(val->data);
+    }
+    std::shared_ptr<HeteroSubgraph> subg(
+        new HeteroSubgraph(hg->EdgeSubgraph(eid_vec, preserve_nodes)));
+    *rv = HeteroSubgraphRef(subg);
+  });
+
+// HeteroSubgraph C APIs
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroSubgraphGetGraph")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroSubgraphRef subg = args[0];
+    *rv = HeteroGraphRef(subg->graph);
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroSubgraphGetInducedVertices")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroSubgraphRef subg = args[0];
+    List<Value> induced_verts;
+    for (IdArray arr : subg->induced_vertices) {
+      induced_verts.push_back(Value(MakeValue(arr)));
+    }
+    *rv = induced_verts;
+  });
+
+DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLHeteroSubgraphGetInducedEdges")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    HeteroSubgraphRef subg = args[0];
+    List<Value> induced_edges;
+    for (IdArray arr : subg->induced_edges) {
+      induced_edges.push_back(Value(MakeValue(arr)));
+    }
+    *rv = induced_edges;
+  });
+
+}  // namespace dgl

src/graph/heterograph.h

+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file graph/heterograph.h
+ * \brief Heterograph
+ */
+
+#ifndef DGL_GRAPH_HETEROGRAPH_H_
+#define DGL_GRAPH_HETEROGRAPH_H_
+
+#include <dgl/base_heterograph.h>
+#include <dgl/lazy.h>
+#include <utility>
+#include <string>
+#include <vector>
+
+namespace dgl {
+
+/*! \brief Heterograph */
+class HeteroGraph : public BaseHeteroGraph {
+ public:
+  HeteroGraph(GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& rel_graphs);
+
+  uint64_t NumVertexTypes() const override {
+    return meta_graph_->NumVertices();
+  }
+
+  uint64_t NumEdgeTypes() const override {
+    return meta_graph_->NumEdges();
+  }
+
+  HeteroGraphPtr GetRelationGraph(dgl_type_t etype) const override {
+    CHECK_LT(etype, meta_graph_->NumEdges()) << "Invalid edge type: " << etype;
+    return relation_graphs_[etype];
+  }
+
+  void AddVertices(dgl_type_t vtype, uint64_t num_vertices) override {
+    LOG(FATAL) << "Bipartite graph is not mutable.";
+  }
+
+  void AddEdge(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) override {
+    LOG(FATAL) << "Bipartite graph is not mutable.";
+  }
+
+  void AddEdges(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) override {
+    LOG(FATAL) << "Bipartite graph is not mutable.";
+  }
+
+  void Clear() override {
+    LOG(FATAL) << "Bipartite graph is not mutable.";
+  }
+
+  DLContext Context() const override {
+    return relation_graphs_[0]->Context();
+  }
+
+  uint8_t NumBits() const override {
+    return relation_graphs_[0]->NumBits();
+  }
+
+  bool IsMultigraph() const override;
+
+  bool IsReadonly() const override {
+    return true;
+  }
+
+  uint64_t NumVertices(dgl_type_t vtype) const override {
+    CHECK(meta_graph_->HasVertex(vtype)) << "Invalid vertex type: " << vtype;
+    return num_verts_per_type_[vtype];
+  }
+
+  uint64_t NumEdges(dgl_type_t etype) const override {
+    return GetRelationGraph(etype)->NumEdges(0);
+  }
+
+  bool HasVertex(dgl_type_t vtype, dgl_id_t vid) const override {
+    return vid < NumVertices(vtype);
+  }
+
+  BoolArray HasVertices(dgl_type_t vtype, IdArray vids) const override;
+
+  bool HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
+    return GetRelationGraph(etype)->HasEdgeBetween(0, src, dst);
+  }
+
+  BoolArray HasEdgesBetween(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) const override {
+    return GetRelationGraph(etype)->HasEdgesBetween(0, src_ids, dst_ids);
+  }
+
+  IdArray Predecessors(dgl_type_t etype, dgl_id_t dst) const override {
+    return GetRelationGraph(etype)->Predecessors(0, dst);
+  }
+
+  IdArray Successors(dgl_type_t etype, dgl_id_t src) const override {
+    return GetRelationGraph(etype)->Successors(0, src);
+  }
+
+  IdArray EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
+    return GetRelationGraph(etype)->EdgeId(0, src, dst);
+  }
+
+  EdgeArray EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const override {
+    return GetRelationGraph(etype)->EdgeIds(0, src, dst);
+  }
+
+  std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_type_t etype, dgl_id_t eid) const override {
+    return GetRelationGraph(etype)->FindEdge(0, eid);
+  }
+
+  EdgeArray FindEdges(dgl_type_t etype, IdArray eids) const override {
+    return GetRelationGraph(etype)->FindEdges(0, eids);
+  }
+
+  EdgeArray InEdges(dgl_type_t etype, dgl_id_t vid) const override {
+    return GetRelationGraph(etype)->InEdges(0, vid);
+  }
+
+  EdgeArray InEdges(dgl_type_t etype, IdArray vids) const override {
+    return GetRelationGraph(etype)->InEdges(0, vids);
+  }
+
+  EdgeArray OutEdges(dgl_type_t etype, dgl_id_t vid) const override {
+    return GetRelationGraph(etype)->OutEdges(0, vid);
+  }
+
+  EdgeArray OutEdges(dgl_type_t etype, IdArray vids) const override {
+    return GetRelationGraph(etype)->OutEdges(0, vids);
+  }
+
+  EdgeArray Edges(dgl_type_t etype, const std::string &order = "") const override {
+    return GetRelationGraph(etype)->Edges(0, order);
+  }
+
+  uint64_t InDegree(dgl_type_t etype, dgl_id_t vid) const override {
+    return GetRelationGraph(etype)->InDegree(0, vid);
+  }
+
+  DegreeArray InDegrees(dgl_type_t etype, IdArray vids) const override {
+    return GetRelationGraph(etype)->InDegrees(0, vids);
+  }
+
+  uint64_t OutDegree(dgl_type_t etype, dgl_id_t vid) const override {
+    return GetRelationGraph(etype)->OutDegree(0, vid);
+  }
+
+  DegreeArray OutDegrees(dgl_type_t etype, IdArray vids) const override {
+    return GetRelationGraph(etype)->OutDegrees(0, vids);
+  }
+
+  DGLIdIters SuccVec(dgl_type_t etype, dgl_id_t vid) const override {
+    return GetRelationGraph(etype)->SuccVec(0, vid);
+  }
+
+  DGLIdIters OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const override {
+    return GetRelationGraph(etype)->OutEdgeVec(0, vid);
+  }
+
+  DGLIdIters PredVec(dgl_type_t etype, dgl_id_t vid) const override {
+    return GetRelationGraph(etype)->PredVec(0, vid);
+  }
+
+  DGLIdIters InEdgeVec(dgl_type_t etype, dgl_id_t vid) const override {
+    return GetRelationGraph(etype)->InEdgeVec(0, vid);
+  }
+
+  std::vector<IdArray> GetAdj(
+      dgl_type_t etype, bool transpose, const std::string &fmt) const override {
+    return GetRelationGraph(etype)->GetAdj(0, transpose, fmt);
+  }
+
+  HeteroSubgraph VertexSubgraph(const std::vector<IdArray>& vids) const override;
+
+  HeteroSubgraph EdgeSubgraph(
+      const std::vector<IdArray>& eids, bool preserve_nodes = false) const override;
+
+ private:
+  /*! \brief A map from edge type to bipartite graph */
+  std::vector<HeteroGraphPtr> relation_graphs_;
+
+  /*! \brief A map from vert type to the number of verts in the type */
+  std::vector<int64_t> num_verts_per_type_;
+
+  /*! \brief True if the graph is a multigraph */
+  Lazy<bool> is_multigraph_;
+};
+
+}  // namespace dgl
+
+#endif  // DGL_GRAPH_HETEROGRAPH_H_

src/graph/immutable_graph.cc

@@ -449,7 +449,6 @@ Subgraph ImmutableGraph::VertexSubgraph(IdArray vids) const {
 }
 
 Subgraph ImmutableGraph::EdgeSubgraph(IdArray eids, bool preserve_nodes) const {
-  // We prefer to generate a subgraph from out-csr.
   auto sg = GetCOO()->EdgeSubgraph(eids, preserve_nodes);
   COOPtr subcoo = std::dynamic_pointer_cast<COO>(sg.graph);
   return Subgraph{GraphPtr(new ImmutableGraph(subcoo)),

tests/graph_index/test_hetero.py

+import numpy as np
+import dgl
+import dgl.ndarray as nd
+import dgl.graph_index as dgl_gidx
+from dgl.utils import toindex
+import backend as F
+
+"""
+Test with a heterograph of three ntypes and three etypes
+meta graph:
+0 -> 1
+1 -> 2
+2 -> 1
+
+Num nodes per ntype:
+0 : 5
+1 : 2
+2 : 3
+
+rel graph:
+0->1 : [0 -> 0, 1 -> 0, 2 -> 0, 3 -> 0]
+1->2 : [0 -> 0, 1 -> 1, 1 -> 2]
+2->1 : [0 -> 1, 1 -> 1, 2 -> 1]
+"""
+
+def _array_equal(dglidx, l):
+    return list(dglidx) == list(l)
+
+def rel1_from_coo():
+    row = toindex([0, 1, 2, 3])
+    col = toindex([0, 0, 0, 0])
+    return dgl_gidx.create_bipartite_from_coo(5, 2, row, col)
+
+def rel2_from_coo():
+    row = toindex([0, 1, 1])
+    col = toindex([0, 1, 2])
+    return dgl_gidx.create_bipartite_from_coo(2, 3, row, col)
+
+def rel3_from_coo():
+    row = toindex([0, 1, 2])
+    col = toindex([1, 1, 1])
+    return dgl_gidx.create_bipartite_from_coo(3, 2, row, col)
+
+def rel1_from_csr():
+    indptr = toindex([0, 1, 2, 3, 4, 4])
+    indices = toindex([0, 0, 0, 0])
+    edge_ids = toindex([0, 1, 2, 3])
+    return dgl_gidx.create_bipartite_from_csr(5, 2, indptr, indices, edge_ids)
+
+def rel2_from_csr():
+    indptr = toindex([0, 1, 3])
+    indices = toindex([0, 1, 2])
+    edge_ids = toindex([0, 1, 2])
+    return dgl_gidx.create_bipartite_from_csr(2, 3, indptr, indices, edge_ids)
+
+def rel3_from_csr():
+    indptr = toindex([0, 1, 2, 3])
+    indices = toindex([1, 1, 1])
+    edge_ids = toindex([0, 1, 2])
+    return dgl_gidx.create_bipartite_from_csr(3, 2, indptr, indices, edge_ids)
+
+def gen_from_coo():
+    mg = dgl_gidx.from_edge_list([(0, 1), (1, 2), (2, 1)], is_multigraph=False, readonly=True)
+    return dgl_gidx.create_heterograph(mg, [rel1_from_coo(), rel2_from_coo(), rel3_from_coo()])
+
+def gen_from_csr():
+    mg = dgl_gidx.from_edge_list([(0, 1), (1, 2), (2, 1)], is_multigraph=False, readonly=True)
+    return dgl_gidx.create_heterograph(mg, [rel1_from_csr(), rel2_from_csr(), rel3_from_csr()])
+
+def test_query():
+    R1 = 0
+    R2 = 1
+    R3 = 2
+    def _test_g(g):
+        assert g.number_of_ntypes() == 3
+        assert g.number_of_etypes() == 3
+        assert g.meta_graph.number_of_nodes() == 3
+        assert g.meta_graph.number_of_edges() == 3
+        assert g.ctx() == nd.cpu(0)
+        assert g.nbits() == 64
+        assert not g.is_multigraph()
+        assert g.is_readonly()
+        # relation graph 1
+        assert g.number_of_nodes(R1) == 5
+        assert g.number_of_edges(R1) == 4
+        assert g.has_node(0, 0)
+        assert not g.has_node(0, 10)
+        assert _array_equal(g.has_nodes(0, toindex([0, 10])), [1, 0])
+        assert g.has_edge_between(R1, 3, 0)
+        assert not g.has_edge_between(R1, 4, 0)
+        assert _array_equal(g.has_edges_between(R1, toindex([3, 4]), toindex([0, 0])), [1, 0])
+        assert _array_equal(g.predecessors(R1, 0), [0, 1, 2, 3])
+        assert _array_equal(g.predecessors(R1, 1), [])
+        assert _array_equal(g.successors(R1, 3), [0])
+        assert _array_equal(g.successors(R1, 4), [])
+        assert _array_equal(g.edge_id(R1, 0, 0), [0])
+        src, dst, eid = g.edge_ids(R1, toindex([0, 2, 1, 3]), toindex([0, 0, 0, 0]))
+        assert _array_equal(src, [0, 2, 1, 3])
+        assert _array_equal(dst, [0, 0, 0, 0])
+        assert _array_equal(eid, [0, 2, 1, 3])
+        src, dst, eid = g.find_edges(R1, toindex([3, 0]))
+        assert _array_equal(src, [3, 0])
+        assert _array_equal(dst, [0, 0])
+        assert _array_equal(eid, [3, 0])
+        src, dst, eid = g.in_edges(R1, toindex([0, 1]))
+        assert _array_equal(src, [0, 1, 2, 3])
+        assert _array_equal(dst, [0, 0, 0, 0])
+        assert _array_equal(eid, [0, 1, 2, 3])
+        src, dst, eid = g.out_edges(R1, toindex([1, 0, 4]))
+        assert _array_equal(src, [1, 0])
+        assert _array_equal(dst, [0, 0])
+        assert _array_equal(eid, [1, 0])
+        src, dst, eid = g.edges(R1, 'eid')
+        assert _array_equal(src, [0, 1, 2, 3])
+        assert _array_equal(dst, [0, 0, 0, 0])
+        assert _array_equal(eid, [0, 1, 2, 3])
+        assert g.in_degree(R1, 0) == 4
+        assert g.in_degree(R1, 1) == 0
+        assert _array_equal(g.in_degrees(R1, toindex([0, 1])), [4, 0])
+        assert g.out_degree(R1, 2) == 1
+        assert g.out_degree(R1, 4) == 0
+        assert _array_equal(g.out_degrees(R1, toindex([4, 2])), [0, 1])
+        # adjmat
+        adj = g.adjacency_matrix(R1, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[1., 0.],
+                          [1., 0.],
+                          [1., 0.],
+                          [1., 0.],
+                          [0., 0.]]))
+        adj = g.adjacency_matrix(R2, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[1., 0., 0.],
+                          [0., 1., 1.]]))
+        adj = g.adjacency_matrix(R3, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[0., 1.],
+                          [0., 1.],
+                          [0., 1.]]))
+
+    g = gen_from_coo()
+    _test_g(g)
+    g = gen_from_csr()
+    _test_g(g)
+
+def test_subgraph():
+    R1 = 0
+    R2 = 1
+    R3 = 2
+    def _test_g(g):
+        # node subgraph
+        induced_nodes = [toindex([0, 1, 4]), toindex([0]), toindex([0, 2])]
+        sub = g.node_subgraph(induced_nodes)
+        subg = sub.graph
+        assert subg.number_of_ntypes() == 3
+        assert subg.number_of_etypes() == 3
+        assert subg.number_of_nodes(0) == 3
+        assert subg.number_of_nodes(1) == 1
+        assert subg.number_of_nodes(2) == 2
+        assert subg.number_of_edges(R1) == 2
+        assert subg.number_of_edges(R2) == 1
+        assert subg.number_of_edges(R3) == 0
+        adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[1.],
+                          [1.],
+                          [0.]]))
+        adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[1., 0.]]))
+        adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[0.],
+                          [0.]]))
+        assert len(sub.induced_nodes) == 3
+        assert _array_equal(sub.induced_nodes[0], induced_nodes[0])
+        assert _array_equal(sub.induced_nodes[1], induced_nodes[1])
+        assert _array_equal(sub.induced_nodes[2], induced_nodes[2])
+        assert len(sub.induced_edges) == 3
+        assert _array_equal(sub.induced_edges[0], [0, 1])
+        assert _array_equal(sub.induced_edges[1], [0])
+        assert _array_equal(sub.induced_edges[2], [])
+
+        # node subgraph with empty type graph
+        induced_nodes = [toindex([0, 1, 4]), toindex([0]), toindex([])]
+        sub = g.node_subgraph(induced_nodes)
+        subg = sub.graph
+        assert subg.number_of_ntypes() == 3
+        assert subg.number_of_etypes() == 3
+        assert subg.number_of_nodes(0) == 3
+        assert subg.number_of_nodes(1) == 1
+        assert subg.number_of_nodes(2) == 0
+        assert subg.number_of_edges(R1) == 2
+        assert subg.number_of_edges(R2) == 0
+        assert subg.number_of_edges(R3) == 0
+        adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[1.],
+                          [1.],
+                          [0.]]))
+        adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([]))
+        adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([]))
+
+        # edge subgraph (preserve_nodes=False)
+        induced_edges = [toindex([0, 2]), toindex([0]), toindex([0, 1, 2])]
+        sub = g.edge_subgraph(induced_edges, False)
+        subg = sub.graph
+        assert subg.number_of_ntypes() == 3
+        assert subg.number_of_etypes() == 3
+        assert subg.number_of_nodes(0) == 2
+        assert subg.number_of_nodes(1) == 2
+        assert subg.number_of_nodes(2) == 3
+        assert subg.number_of_edges(R1) == 2
+        assert subg.number_of_edges(R2) == 1
+        assert subg.number_of_edges(R3) == 3
+        adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[1., 0.],
+                          [1., 0.]]))
+        adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[1., 0., 0.],
+                          [0., 0., 0.]]))
+        adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[0., 1.],
+                          [0., 1.],
+                          [0., 1.]]))
+        assert len(sub.induced_nodes) == 3
+        assert _array_equal(sub.induced_nodes[0], [0, 2])
+        assert _array_equal(sub.induced_nodes[1], [0, 1])
+        assert _array_equal(sub.induced_nodes[2], [0, 1, 2])
+        assert len(sub.induced_edges) == 3
+        assert _array_equal(sub.induced_edges[0], induced_edges[0])
+        assert _array_equal(sub.induced_edges[1], induced_edges[1])
+        assert _array_equal(sub.induced_edges[2], induced_edges[2])
+
+        # edge subgraph (preserve_nodes=True)
+        induced_edges = [toindex([0, 2]), toindex([0]), toindex([0, 1, 2])]
+        sub = g.edge_subgraph(induced_edges, True)
+        subg = sub.graph
+        assert subg.number_of_ntypes() == 3
+        assert subg.number_of_etypes() == 3
+        assert subg.number_of_nodes(0) == 5
+        assert subg.number_of_nodes(1) == 2
+        assert subg.number_of_nodes(2) == 3
+        assert subg.number_of_edges(R1) == 2
+        assert subg.number_of_edges(R2) == 1
+        assert subg.number_of_edges(R3) == 3
+        adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[1., 0.],
+                          [0., 0.],
+                          [1., 0.],
+                          [0., 0.],
+                          [0., 0.]]))
+        adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[1., 0., 0.],
+                          [0., 0., 0.]]))
+        adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[0., 1.],
+                          [0., 1.],
+                          [0., 1.]]))
+        assert len(sub.induced_nodes) == 3
+        assert _array_equal(sub.induced_nodes[0], [0, 1, 2, 3, 4])
+        assert _array_equal(sub.induced_nodes[1], [0, 1])
+        assert _array_equal(sub.induced_nodes[2], [0, 1, 2])
+        assert len(sub.induced_edges) == 3
+        assert _array_equal(sub.induced_edges[0], induced_edges[0])
+        assert _array_equal(sub.induced_edges[1], induced_edges[1])
+        assert _array_equal(sub.induced_edges[2], induced_edges[2])
+
+        # edge subgraph with empty induced edges (preserve_nodes=False)
+        induced_edges = [toindex([0, 2]), toindex([]), toindex([0, 1, 2])]
+        sub = g.edge_subgraph(induced_edges, False)
+        subg = sub.graph
+        assert subg.number_of_ntypes() == 3
+        assert subg.number_of_etypes() == 3
+        assert subg.number_of_nodes(0) == 2
+        assert subg.number_of_nodes(1) == 2
+        assert subg.number_of_nodes(2) == 3
+        assert subg.number_of_edges(R1) == 2
+        assert subg.number_of_edges(R2) == 0
+        assert subg.number_of_edges(R3) == 3
+        adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[1., 0.],
+                          [1., 0.]]))
+        adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[0., 0., 0.],
+                          [0., 0., 0.]]))
+        adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[0., 1.],
+                          [0., 1.],
+                          [0., 1.]]))
+        assert len(sub.induced_nodes) == 3
+        assert _array_equal(sub.induced_nodes[0], [0, 2])
+        assert _array_equal(sub.induced_nodes[1], [0, 1])
+        assert _array_equal(sub.induced_nodes[2], [0, 1, 2])
+        assert len(sub.induced_edges) == 3
+        assert _array_equal(sub.induced_edges[0], induced_edges[0])
+        assert _array_equal(sub.induced_edges[1], induced_edges[1])
+        assert _array_equal(sub.induced_edges[2], induced_edges[2])
+
+        # edge subgraph with empty induced edges (preserve_nodes=True)
+        induced_edges = [toindex([0, 2]), toindex([]), toindex([0, 1, 2])]
+        sub = g.edge_subgraph(induced_edges, True)
+        subg = sub.graph
+        assert subg.number_of_ntypes() == 3
+        assert subg.number_of_etypes() == 3
+        assert subg.number_of_nodes(0) == 5
+        assert subg.number_of_nodes(1) == 2
+        assert subg.number_of_nodes(2) == 3
+        assert subg.number_of_edges(R1) == 2
+        assert subg.number_of_edges(R2) == 0
+        assert subg.number_of_edges(R3) == 3
+        adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[1., 0.],
+                          [0., 0.],
+                          [1., 0.],
+                          [0., 0.],
+                          [0., 0.]]))
+        adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[0., 0., 0.],
+                          [0., 0., 0.]]))
+        adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
+        assert np.allclose(F.sparse_to_numpy(adj),
+                np.array([[0., 1.],
+                          [0., 1.],
+                          [0., 1.]]))
+        assert len(sub.induced_nodes) == 3
+        assert _array_equal(sub.induced_nodes[0], [0, 1, 2, 3, 4])
+        assert _array_equal(sub.induced_nodes[1], [0, 1])
+        assert _array_equal(sub.induced_nodes[2], [0, 1, 2])
+        assert len(sub.induced_edges) == 3
+        assert _array_equal(sub.induced_edges[0], induced_edges[0])
+        assert _array_equal(sub.induced_edges[1], induced_edges[1])
+        assert _array_equal(sub.induced_edges[2], induced_edges[2])
+
+    g = gen_from_coo()
+    _test_g(g)
+    g = gen_from_csr()
+    _test_g(g)
+
+if __name__ == '__main__':
+    test_query()
+    test_subgraph()