[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)

python/dgl/graph_index.py

@@ -89,7 +89,7 @@ class GraphIndex(ObjectBase):
         v : int
             The dst node.
         """
-        _CAPI_DGLGraphAddEdge(self, u, v)
+        _CAPI_DGLGraphAddEdge(self, int(u), int(v))
         self.clear_cache()
 
     def add_edges(self, u, v):
@@ -323,6 +323,24 @@ class GraphIndex(ObjectBase):
 
         return src, dst, eid
 
+    def find_edge(self, eid):
+        """Return the edge tuple of the given id.
+
+        Parameters
+        ----------
+        eid : int
+            The edge id.
+
+        Returns
+        -------
+        int
+            src node id
+        int
+            dst node id
+        """
+        ret = _CAPI_DGLGraphFindEdge(self, int(eid))
+        return ret(0), ret(1)
+
     def find_edges(self, eid):
         """Return a triplet of arrays that contains the edge IDs.
 
@@ -1235,4 +1253,737 @@ def create_graph_index(graph_data, multigraph, readonly):
                            % type(graph_data))
         return gidx
 
+#############################################################
+# Hetero graph
+#############################################################
+
+@register_object('graph.HeteroGraph')
+class HeteroGraphIndex(ObjectBase):
+    """HeteroGraph index object.
+
+    Note
+    ----
+    Do not create GraphIndex directly.
+    """
+    def __new__(cls):
+        obj = ObjectBase.__new__(cls)
+        return obj
+
+    def __getstate__(self):
+        # TODO
+        return
+
+    def __setstate__(self, state):
+        # TODO
+        pass
+
+    @property
+    def meta_graph(self):
+        """Meta graph
+
+        Returns
+        -------
+        GraphIndex
+            The meta graph.
+        """
+        return _CAPI_DGLHeteroGetMetaGraph(self)
+
+    def number_of_ntypes(self):
+        """Return number of node types."""
+        return self.meta_graph.number_of_nodes()
+
+    def number_of_etypes(self):
+        """Return number of edge types."""
+        return self.meta_graph.number_of_edges()
+
+    def get_relation_graph(self, etype):
+        """Get the bipartite graph of the given edge/relation type.
+
+        Parameters
+        ----------
+        etype : int
+            The edge/relation type.
+
+        Returns
+        -------
+        HeteroGraphIndex
+            The bipartite graph.
+        """
+        return _CAPI_DGLHeteroGetRelationGraph(self, int(etype))
+
+    def add_nodes(self, ntype, num):
+        """Add nodes.
+
+        Parameters
+        ----------
+        ntype : int
+            Node type
+        num : int
+            Number of nodes to be added.
+        """
+        _CAPI_DGLHetero(self, int(ntype), int(num))
+
+    def add_edge(self, etype, u, v):
+        """Add one edge.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        u : int
+            The src node.
+        v : int
+            The dst node.
+        """
+        _CAPI_DGLHeteroAddEdge(self, int(etype), int(u), int(v))
+
+    def add_edges(self, etype, u, v):
+        """Add many edges.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        u : utils.Index
+            The src nodes.
+        v : utils.Index
+            The dst nodes.
+        """
+        _CAPI_DGLHeteroAddEdges(self, int(etype), u.todgltensor(), v.todgltensor())
+
+    def clear(self):
+        """Clear the graph."""
+        _CAPI_DGLHeteroClear(self)
+
+    def ctx(self):
+        """Return the context of this graph index.
+
+        Returns
+        -------
+        DGLContext
+            The context of the graph.
+        """
+        return _CAPI_DGLHeteroContext(self)
+
+    def nbits(self):
+        """Return the number of integer bits used in the storage (32 or 64).
+
+        Returns
+        -------
+        int
+            The number of bits.
+        """
+        return _CAPI_DGLHeteroNumBits(self)
+
+    def is_multigraph(self):
+        """Return whether the graph is a multigraph
+
+        Returns
+        -------
+        bool
+            True if it is a multigraph, False otherwise.
+        """
+        return bool(_CAPI_DGLHeteroIsMultigraph(self))
+
+    def is_readonly(self):
+        """Return whether the graph index is read-only.
+
+        Returns
+        -------
+        bool
+            True if it is a read-only graph, False otherwise.
+        """
+        return bool(_CAPI_DGLHeteroIsReadonly(self))
+
+    def number_of_nodes(self, ntype):
+        """Return the number of nodes.
+
+        Parameters
+        ----------
+        ntype : int
+            Node type
+
+        Returns
+        -------
+        int
+            The number of nodes
+        """
+        return _CAPI_DGLHeteroNumVertices(self, int(ntype))
+
+    def number_of_edges(self, etype):
+        """Return the number of edges.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+
+        Returns
+        -------
+        int
+            The number of edges
+        """
+        return _CAPI_DGLHeteroNumEdges(self, int(etype))
+
+    def has_node(self, ntype, vid):
+        """Return true if the node exists.
+
+        Parameters
+        ----------
+        ntype : int
+            Node type
+        vid : int
+            The nodes
+
+        Returns
+        -------
+        bool
+            True if the node exists, False otherwise.
+        """
+        return bool(_CAPI_DGLHeteroHasVertex(self, int(ntype), int(vid)))
+
+    def has_nodes(self, ntype, vids):
+        """Return true if the nodes exist.
+
+        Parameters
+        ----------
+        ntype : int
+            Node type
+        vid : utils.Index
+            The nodes
+
+        Returns
+        -------
+        utils.Index
+            0-1 array indicating existence
+        """
+        vid_array = vids.todgltensor()
+        return utils.toindex(_CAPI_DGLHeteroHasVertices(self, int(ntype), vid_array))
+
+    def has_edge_between(self, etype, u, v):
+        """Return true if the edge exists.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        u : int
+            The src node.
+        v : int
+            The dst node.
+
+        Returns
+        -------
+        bool
+            True if the edge exists, False otherwise
+        """
+        return bool(_CAPI_DGLHeteroHasEdgeBetween(self, int(etype), int(u), int(v)))
+
+    def has_edges_between(self, etype, u, v):
+        """Return true if the edge exists.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        u : utils.Index
+            The src nodes.
+        v : utils.Index
+            The dst nodes.
+
+        Returns
+        -------
+        utils.Index
+            0-1 array indicating existence
+        """
+        u_array = u.todgltensor()
+        v_array = v.todgltensor()
+        return utils.toindex(_CAPI_DGLHeteroHasEdgesBetween(
+            self, int(etype), u_array, v_array))
+
+    def predecessors(self, etype, v):
+        """Return the predecessors of the node.
+
+        Assume that node_type(v) == dst_type(etype). Thus, the ntype argument is omitted.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        v : int
+            The node.
+
+        Returns
+        -------
+        utils.Index
+            Array of predecessors
+        """
+        return utils.toindex(_CAPI_DGLHeteroPredecessors(
+            self, int(etype), int(v)))
+
+    def successors(self, etype, v):
+        """Return the successors of the node.
+
+        Assume that node_type(v) == src_type(etype). Thus, the ntype argument is omitted.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        v : int
+            The node.
+
+        Returns
+        -------
+        utils.Index
+            Array of successors
+        """
+        return utils.toindex(_CAPI_DGLHeteroSuccessors(
+            self, int(etype), int(v)))
+
+    def edge_id(self, etype, u, v):
+        """Return the id array of all edges between u and v.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        u : int
+            The src node.
+        v : int
+            The dst node.
+
+        Returns
+        -------
+        utils.Index
+            The edge id array.
+        """
+        return utils.toindex(_CAPI_DGLHeteroEdgeId(
+            self, int(etype), int(u), int(v)))
+
+    def edge_ids(self, etype, u, v):
+        """Return a triplet of arrays that contains the edge IDs.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        u : utils.Index
+            The src nodes.
+        v : utils.Index
+            The dst nodes.
+
+        Returns
+        -------
+        utils.Index
+            The src nodes.
+        utils.Index
+            The dst nodes.
+        utils.Index
+            The edge ids.
+        """
+        u_array = u.todgltensor()
+        v_array = v.todgltensor()
+        edge_array = _CAPI_DGLHeteroEdgeIds(self, int(etype), u_array, v_array)
+
+        src = utils.toindex(edge_array(0))
+        dst = utils.toindex(edge_array(1))
+        eid = utils.toindex(edge_array(2))
+
+        return src, dst, eid
+
+    def find_edges(self, etype, eid):
+        """Return a triplet of arrays that contains the edge IDs.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        eid : utils.Index
+            The edge ids.
+
+        Returns
+        -------
+        utils.Index
+            The src nodes.
+        utils.Index
+            The dst nodes.
+        utils.Index
+            The edge ids.
+        """
+        eid_array = eid.todgltensor()
+        edge_array = _CAPI_DGLHeteroFindEdges(self, int(etype), eid_array)
+
+        src = utils.toindex(edge_array(0))
+        dst = utils.toindex(edge_array(1))
+        eid = utils.toindex(edge_array(2))
+
+        return src, dst, eid
+
+    def in_edges(self, etype, v):
+        """Return the in edges of the node(s).
+
+        Assume that node_type(v) == dst_type(etype). Thus, the ntype argument is omitted.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        v : utils.Index
+            The node(s).
+
+        Returns
+        -------
+        utils.Index
+            The src nodes.
+        utils.Index
+            The dst nodes.
+        utils.Index
+            The edge ids.
+        """
+        if len(v) == 1:
+            edge_array = _CAPI_DGLHeteroInEdges_1(self, int(etype), int(v[0]))
+        else:
+            v_array = v.todgltensor()
+            edge_array = _CAPI_DGLHeteroInEdges_2(self, int(etype), v_array)
+        src = utils.toindex(edge_array(0))
+        dst = utils.toindex(edge_array(1))
+        eid = utils.toindex(edge_array(2))
+        return src, dst, eid
+
+    def out_edges(self, etype, v):
+        """Return the out edges of the node(s).
+
+        Assume that node_type(v) == src_type(etype). Thus, the ntype argument is omitted.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        v : utils.Index
+            The node(s).
+
+        Returns
+        -------
+        utils.Index
+            The src nodes.
+        utils.Index
+            The dst nodes.
+        utils.Index
+            The edge ids.
+        """
+        if len(v) == 1:
+            edge_array = _CAPI_DGLHeteroOutEdges_1(self, int(etype), int(v[0]))
+        else:
+            v_array = v.todgltensor()
+            edge_array = _CAPI_DGLHeteroOutEdges_2(self, int(etype), v_array)
+        src = utils.toindex(edge_array(0))
+        dst = utils.toindex(edge_array(1))
+        eid = utils.toindex(edge_array(2))
+        return src, dst, eid
+
+    def edges(self, etype, order=None):
+        """Return all the edges
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        order : string
+            The order of the returned edges. Currently support:
+
+            - 'srcdst' : sorted by their src and dst ids.
+            - 'eid'    : sorted by edge Ids.
+            - None     : the arbitrary order.
+
+        Returns
+        -------
+        utils.Index
+            The src nodes.
+        utils.Index
+            The dst nodes.
+        utils.Index
+            The edge ids.
+        """
+        if order is None:
+            order = ""
+        edge_array = _CAPI_DGLHeteroEdges(self, int(etype), order)
+        src = edge_array(0)
+        dst = edge_array(1)
+        eid = edge_array(2)
+        src = utils.toindex(src)
+        dst = utils.toindex(dst)
+        eid = utils.toindex(eid)
+        return src, dst, eid
+
+    def in_degree(self, etype, v):
+        """Return the in degree of the node.
+
+        Assume that node_type(v) == dst_type(etype). Thus, the ntype argument is omitted.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        v : int
+            The node.
+
+        Returns
+        -------
+        int
+            The in degree.
+        """
+        return _CAPI_DGLHeteroInDegree(self, int(etype), int(v))
+
+    def in_degrees(self, etype, v):
+        """Return the in degrees of the nodes.
+
+        Assume that node_type(v) == dst_type(etype). Thus, the ntype argument is omitted.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        v : utils.Index
+            The nodes.
+
+        Returns
+        -------
+        int
+            The in degree array.
+        """
+        v_array = v.todgltensor()
+        return utils.toindex(_CAPI_DGLHeteroInDegrees(self, int(etype), v_array))
+
+    def out_degree(self, etype, v):
+        """Return the out degree of the node.
+
+        Assume that node_type(v) == src_type(etype). Thus, the ntype argument is omitted.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        v : int
+            The node.
+
+        Returns
+        -------
+        int
+            The out degree.
+        """
+        return _CAPI_DGLHeteroOutDegree(self, int(etype), int(v))
+
+    def out_degrees(self, etype, v):
+        """Return the out degrees of the nodes.
+
+        Assume that node_type(v) == src_type(etype). Thus, the ntype argument is omitted.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        v : utils.Index
+            The nodes.
+
+        Returns
+        -------
+        int
+            The out degree array.
+        """
+        v_array = v.todgltensor()
+        return utils.toindex(_CAPI_DGLHeteroOutDegrees(self, int(etype), v_array))
+
+    def adjacency_matrix(self, etype, transpose, ctx):
+        """Return the adjacency matrix representation of this graph.
+
+        By default, a row of returned adjacency matrix represents the destination
+        of an edge and the column represents the source.
+
+        When transpose is True, a row represents the source and a column represents
+        a destination.
+
+        Parameters
+        ----------
+        etype : int
+            Edge type
+        transpose : bool
+            A flag to transpose the returned adjacency matrix.
+        ctx : context
+            The context of the returned matrix.
+
+        Returns
+        -------
+        SparseTensor
+            The adjacency matrix.
+        utils.Index
+            A index for data shuffling due to sparse format change. Return None
+            if shuffle is not required.
+        """
+        if not isinstance(transpose, bool):
+            raise DGLError('Expect bool value for "transpose" arg,'
+                           ' but got %s.' % (type(transpose)))
+        fmt = F.get_preferred_sparse_format()
+        rst = _CAPI_DGLHeteroGetAdj(self, int(etype), transpose, fmt)
+        # convert to framework-specific sparse matrix
+        srctype, dsttype = self.meta_graph.find_edge(etype)
+        nrows = self.number_of_nodes(srctype) if transpose else self.number_of_nodes(dsttype)
+        ncols = self.number_of_nodes(dsttype) if transpose else self.number_of_nodes(srctype)
+        nnz = self.number_of_edges(etype)
+        if fmt == "csr":
+            indptr = F.copy_to(utils.toindex(rst(0)).tousertensor(), ctx)
+            indices = F.copy_to(utils.toindex(rst(1)).tousertensor(), ctx)
+            shuffle = utils.toindex(rst(2))
+            dat = F.ones(nnz, dtype=F.float32, ctx=ctx)  # FIXME(minjie): data type
+            spmat = F.sparse_matrix(dat, ('csr', indices, indptr), (nrows, ncols))[0]
+            return spmat, shuffle
+        elif fmt == "coo":
+            idx = F.copy_to(utils.toindex(rst(0)).tousertensor(), ctx)
+            idx = F.reshape(idx, (2, nnz))
+            dat = F.ones((nnz,), dtype=F.float32, ctx=ctx)
+            adj, shuffle_idx = F.sparse_matrix(dat, ('coo', idx), (nrows, ncols))
+            shuffle_idx = utils.toindex(shuffle_idx) if shuffle_idx is not None else None
+            return adj, shuffle_idx
+        else:
+            raise Exception("unknown format")
+
+    def node_subgraph(self, induced_nodes):
+        """Return the induced node subgraph.
+
+        Parameters
+        ----------
+        induced_nodes : list of utils.Index
+            Induced nodes. The length should be equal to the number of
+            node types in this heterograph.
+
+        Returns
+        -------
+        SubgraphIndex
+            The subgraph index.
+        """
+        vids = [nodes.todgltensor() for nodes in induced_nodes]
+        return _CAPI_DGLHeteroVertexSubgraph(self, vids)
+
+    def edge_subgraph(self, induced_edges, preserve_nodes):
+        """Return the induced edge subgraph.
+
+        Parameters
+        ----------
+        induced_edges : list of utils.Index
+            Induced edges. The length should be equal to the number of
+            edge types in this heterograph.
+        preserve_nodes : bool
+            Indicates whether to preserve all nodes or not.
+            If true, keep the nodes which have no edge connected in the subgraph;
+            If false, all nodes without edge connected to it would be removed.
+
+        Returns
+        -------
+        SubgraphIndex
+            The subgraph index.
+        """
+        eids = [edges.todgltensor() for edges in induced_edges]
+        return _CAPI_DGLHeteroEdgeSubgraph(self, eids, preserve_nodes)
+
+@register_object('graph.HeteroSubgraph')
+class HeteroSubgraphIndex(ObjectBase):
+    """Hetero-subgraph data structure"""
+    @property
+    def graph(self):
+        """The subgraph structure
+
+        Returns
+        -------
+        HeteroGraphIndex
+            The subgraph
+        """
+        return _CAPI_DGLHeteroSubgraphGetGraph(self)
+
+    @property
+    def induced_nodes(self):
+        """Induced nodes for each node type. The return list
+        length should be equal to the number of node types.
+
+        Returns
+        -------
+        list of utils.Index
+            Induced nodes
+        """
+        ret = _CAPI_DGLHeteroSubgraphGetInducedVertices(self)
+        return [utils.toindex(v.data) for v in ret]
+
+    @property
+    def induced_edges(self):
+        """Induced edges for each edge type. The return list
+        length should be equal to the number of edge types.
+
+        Returns
+        -------
+        list of utils.Index
+            Induced edges
+        """
+        ret = _CAPI_DGLHeteroSubgraphGetInducedEdges(self)
+        return [utils.toindex(v.data) for v in ret]
+
+def create_bipartite_from_coo(num_src, num_dst, row, col):
+    """Create a bipartite graph index from COO format
+
+    Parameters
+    ----------
+    num_src : int
+        Number of nodes in the src type.
+    num_dst : int
+        Number of nodes in the dst type.
+    row : utils.Index
+        Row index.
+    col : utils.Index
+        Col index.
+
+    Returns
+    -------
+    HeteroGraphIndex
+    """
+    return _CAPI_DGLHeteroCreateBipartiteFromCOO(
+        int(num_src), int(num_dst), row.todgltensor(), col.todgltensor())
+
+def create_bipartite_from_csr(num_src, num_dst, indptr, indices, edge_ids):
+    """Create a bipartite graph index from CSR format
+
+    Parameters
+    ----------
+    num_src : int
+        Number of nodes in the src type.
+    num_dst : int
+        Number of nodes in the dst type.
+    indptr : utils.Index
+        CSR indptr.
+    indices : utils.Index
+        CSR indices.
+    edge_ids : utils.Index
+        Edge shuffle id.
+
+    Returns
+    -------
+    HeteroGraphIndex
+    """
+    return _CAPI_DGLHeteroCreateBipartiteFromCSR(
+        int(num_src), int(num_dst),
+        indptr.todgltensor(), indices.todgltensor(), edge_ids.todgltensor())
+
+def create_heterograph(meta_graph, rel_graphs):
+    """Create a heterograph from metagraph and graphs of every relation.
+
+    Parameters
+    ----------
+    meta_graph : GraphIndex
+        Meta-graph.
+    rel_graphs : list of HeteroGraphIndex
+        Bipartite graph of each relation.
+
+    Returns
+    -------
+    HeteroGraphIndex
+    """
+    return _CAPI_DGLHeteroCreateHeteroGraph(meta_graph, rel_graphs)
+
 _init_api("dgl.graph_index")

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.cc

+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file graph/bipartite.cc
+ * \brief Bipartite graph implementation
+ */
+#include <dgl/array.h>
+#include <dgl/lazy.h>
+#include <dgl/immutable_graph.h>
+
+#include "./bipartite.h"
+#include "../c_api_common.h"
+
+namespace dgl {
+namespace {
+inline GraphPtr CreateBipartiteMetaGraph() {
+  std::vector<int64_t> row_vec(1, Bipartite::kSrcVType);
+  std::vector<int64_t> col_vec(1, Bipartite::kDstVType);
+  IdArray row = aten::VecToIdArray(row_vec);
+  IdArray col = aten::VecToIdArray(col_vec);
+  GraphPtr g = ImmutableGraph::CreateFromCOO(2, row, col);
+  return g;
+}
+static const GraphPtr kBipartiteMetaGraph = CreateBipartiteMetaGraph();
+}  // namespace
+
+//////////////////////////////////////////////////////////
+//
+// COO graph implementation
+//
+//////////////////////////////////////////////////////////
+
+/*! \brief COO graph */
+class Bipartite::COO : public BaseHeteroGraph {
+ public:
+  COO(int64_t num_src, int64_t num_dst,
+               IdArray src, IdArray dst)
+    : BaseHeteroGraph(kBipartiteMetaGraph) {
+    adj_ = aten::COOMatrix{num_src, num_dst, src, dst};
+  }
+  COO(int64_t num_src, int64_t num_dst,
+               IdArray src, IdArray dst, bool is_multigraph)
+    : BaseHeteroGraph(kBipartiteMetaGraph),
+      is_multigraph_(is_multigraph) {
+    adj_ = aten::COOMatrix{num_src, num_dst, src, dst};
+  }
+  explicit COO(const aten::COOMatrix& coo)
+    : BaseHeteroGraph(kBipartiteMetaGraph), adj_(coo) {}
+
+  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 {
+    return adj_.row->ctx;
+  }
+
+  uint8_t NumBits() const override {
+    return adj_.row->dtype.bits;
+  }
+
+  bool IsMultigraph() const override {
+    return const_cast<COO*>(this)->is_multigraph_.Get([this] () {
+        return aten::COOHasDuplicate(adj_);
+      });
+  }
+
+  bool IsReadonly() const override {
+    return true;
+  }
+
+  uint64_t NumVertices(dgl_type_t vtype) const override {
+    if (vtype == Bipartite::kSrcVType) {
+      return adj_.num_rows;
+    } else if (vtype == Bipartite::kDstVType) {
+      return adj_.num_cols;
+    } else {
+      LOG(FATAL) << "Invalid vertex type: " << vtype;
+      return 0;
+    }
+  }
+
+  uint64_t NumEdges(dgl_type_t etype) const override {
+    return adj_.row->shape[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 {
+    LOG(FATAL) << "Not enabled for COO graph";
+    return {};
+  }
+
+  bool HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  BoolArray HasEdgesBetween(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  IdArray Predecessors(dgl_type_t etype, dgl_id_t dst) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  IdArray Successors(dgl_type_t etype, dgl_id_t src) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  IdArray EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  EdgeArray EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_type_t etype, dgl_id_t eid) const override {
+    CHECK(eid < NumEdges(etype)) << "Invalid edge id: " << eid;
+    const auto src = aten::IndexSelect(adj_.row, eid);
+    const auto dst = aten::IndexSelect(adj_.col, eid);
+    return std::pair<dgl_id_t, dgl_id_t>(src, dst);
+  }
+
+  EdgeArray FindEdges(dgl_type_t etype, IdArray eids) const override {
+    CHECK(IsValidIdArray(eids)) << "Invalid edge id array";
+    return EdgeArray{aten::IndexSelect(adj_.row, eids),
+                     aten::IndexSelect(adj_.col, eids),
+                     eids};
+  }
+
+  EdgeArray InEdges(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  EdgeArray InEdges(dgl_type_t etype, IdArray vids) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  EdgeArray OutEdges(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  EdgeArray OutEdges(dgl_type_t etype, IdArray vids) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  EdgeArray Edges(dgl_type_t etype, const std::string &order = "") const override {
+    CHECK(order.empty() || order == std::string("eid"))
+      << "COO only support Edges of order \"eid\", but got \""
+      << order << "\".";
+    IdArray rst_eid = aten::Range(0, NumEdges(etype), NumBits(), Context());
+    return EdgeArray{adj_.row, adj_.col, rst_eid};
+  }
+
+  uint64_t InDegree(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  DegreeArray InDegrees(dgl_type_t etype, IdArray vids) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  uint64_t OutDegree(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  DegreeArray OutDegrees(dgl_type_t etype, IdArray vids) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  DGLIdIters SuccVec(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  DGLIdIters OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  DGLIdIters PredVec(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  DGLIdIters InEdgeVec(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  std::vector<IdArray> GetAdj(
+      dgl_type_t etype, bool transpose, const std::string &fmt) const override {
+    CHECK(fmt == "coo") << "Not valid adj format request.";
+    if (transpose) {
+      return {aten::HStack(adj_.col, adj_.row)};
+    } else {
+      return {aten::HStack(adj_.row, adj_.col)};
+    }
+  }
+
+  HeteroSubgraph VertexSubgraph(const std::vector<IdArray>& vids) const override {
+    LOG(INFO) << "Not enabled for COO graph.";
+    return {};
+  }
+
+  HeteroSubgraph EdgeSubgraph(
+      const std::vector<IdArray>& eids, bool preserve_nodes = false) const override {
+    CHECK_EQ(eids.size(), 1) << "Edge type number mismatch.";
+    HeteroSubgraph subg;
+    if (!preserve_nodes) {
+      IdArray new_src = aten::IndexSelect(adj_.row, eids[0]);
+      IdArray new_dst = aten::IndexSelect(adj_.col, eids[0]);
+      subg.induced_vertices.emplace_back(aten::Relabel_({new_src}));
+      subg.induced_vertices.emplace_back(aten::Relabel_({new_dst}));
+      const auto new_nsrc = subg.induced_vertices[0]->shape[0];
+      const auto new_ndst = subg.induced_vertices[1]->shape[0];
+      subg.graph = std::make_shared<COO>(
+          new_nsrc, new_ndst, new_src, new_dst);
+      subg.induced_edges = eids;
+    } else {
+      IdArray new_src = aten::IndexSelect(adj_.row, eids[0]);
+      IdArray new_dst = aten::IndexSelect(adj_.col, eids[0]);
+      subg.induced_vertices.emplace_back(aten::Range(0, NumVertices(0), NumBits(), Context()));
+      subg.induced_vertices.emplace_back(aten::Range(0, NumVertices(1), NumBits(), Context()));
+      subg.graph = std::make_shared<COO>(
+          NumVertices(0), NumVertices(1), new_src, new_dst);
+      subg.induced_edges = eids;
+    }
+    return subg;
+  }
+
+  aten::COOMatrix adj() const {
+    return adj_;
+  }
+
+ private:
+  /*! \brief internal adjacency matrix. Data array is empty */
+  aten::COOMatrix adj_;
+
+  /*! \brief multi-graph flag */
+  Lazy<bool> is_multigraph_;
+};
+
+//////////////////////////////////////////////////////////
+//
+// CSR graph implementation
+//
+//////////////////////////////////////////////////////////
+
+
+/*! \brief CSR graph */
+class Bipartite::CSR : public BaseHeteroGraph {
+ public:
+  CSR(int64_t num_src, int64_t num_dst,
+      IdArray indptr, IdArray indices, IdArray edge_ids)
+    : BaseHeteroGraph(kBipartiteMetaGraph) {
+    adj_ = aten::CSRMatrix{num_src, num_dst, indptr, indices, edge_ids};
+  }
+
+  CSR(int64_t num_src, int64_t num_dst,
+      IdArray indptr, IdArray indices, IdArray edge_ids, bool is_multigraph)
+    : BaseHeteroGraph(kBipartiteMetaGraph),
+      is_multigraph_(is_multigraph) {
+    adj_ = aten::CSRMatrix{num_src, num_dst, indptr, indices, edge_ids};
+  }
+
+  explicit CSR(const aten::CSRMatrix& csr)
+    : BaseHeteroGraph(kBipartiteMetaGraph), adj_(csr) {}
+
+  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 {
+    return adj_.indices->ctx;
+  }
+
+  uint8_t NumBits() const override {
+    return adj_.indices->dtype.bits;
+  }
+
+  bool IsMultigraph() const override {
+    return const_cast<CSR*>(this)->is_multigraph_.Get([this] () {
+        return aten::CSRHasDuplicate(adj_);
+      });
+  }
+
+  bool IsReadonly() const override {
+    return true;
+  }
+
+  uint64_t NumVertices(dgl_type_t vtype) const override {
+    if (vtype == Bipartite::kSrcVType) {
+      return adj_.num_rows;
+    } else if (vtype == Bipartite::kDstVType) {
+      return adj_.num_cols;
+    } else {
+      LOG(FATAL) << "Invalid vertex type: " << vtype;
+      return 0;
+    }
+  }
+
+  uint64_t NumEdges(dgl_type_t etype) const override {
+    return adj_.indices->shape[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 {
+    LOG(FATAL) << "Not enabled for COO graph";
+    return {};
+  }
+
+  bool HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
+    CHECK(HasVertex(0, src)) << "Invalid src vertex id: " << src;
+    CHECK(HasVertex(1, dst)) << "Invalid dst vertex id: " << dst;
+    return aten::CSRIsNonZero(adj_, src, dst);
+  }
+
+  BoolArray HasEdgesBetween(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) const override {
+    CHECK(IsValidIdArray(src_ids)) << "Invalid vertex id array.";
+    CHECK(IsValidIdArray(dst_ids)) << "Invalid vertex id array.";
+    return aten::CSRIsNonZero(adj_, src_ids, dst_ids);
+  }
+
+  IdArray Predecessors(dgl_type_t etype, dgl_id_t dst) const override {
+    LOG(INFO) << "Not enabled for CSR graph.";
+    return {};
+  }
+
+  IdArray Successors(dgl_type_t etype, dgl_id_t src) const override {
+    CHECK(HasVertex(0, src)) << "Invalid src vertex id: " << src;
+    return aten::CSRGetRowColumnIndices(adj_, src);
+  }
+
+  IdArray EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
+    CHECK(HasVertex(0, src)) << "Invalid src vertex id: " << src;
+    CHECK(HasVertex(1, dst)) << "Invalid dst vertex id: " << dst;
+    return aten::CSRGetData(adj_, src, dst);
+  }
+
+  EdgeArray EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const override {
+    CHECK(IsValidIdArray(src)) << "Invalid vertex id array.";
+    CHECK(IsValidIdArray(dst)) << "Invalid vertex id array.";
+    const auto& arrs = aten::CSRGetDataAndIndices(adj_, src, dst);
+    return EdgeArray{arrs[0], arrs[1], arrs[2]};
+  }
+
+  std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_type_t etype, dgl_id_t eid) const override {
+    LOG(INFO) << "Not enabled for CSR graph.";
+    return {};
+  }
+
+  EdgeArray FindEdges(dgl_type_t etype, IdArray eids) const override {
+    LOG(INFO) << "Not enabled for CSR graph.";
+    return {};
+  }
+
+  EdgeArray InEdges(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for CSR graph.";
+    return {};
+  }
+
+  EdgeArray InEdges(dgl_type_t etype, IdArray vids) const override {
+    LOG(INFO) << "Not enabled for CSR graph.";
+    return {};
+  }
+
+  EdgeArray OutEdges(dgl_type_t etype, dgl_id_t vid) const override {
+    CHECK(HasVertex(0, vid)) << "Invalid src vertex id: " << vid;
+    IdArray ret_dst = aten::CSRGetRowColumnIndices(adj_, vid);
+    IdArray ret_eid = aten::CSRGetRowData(adj_, vid);
+    IdArray ret_src = aten::Full(vid, ret_dst->shape[0], NumBits(), ret_dst->ctx);
+    return EdgeArray{ret_src, ret_dst, ret_eid};
+  }
+
+  EdgeArray OutEdges(dgl_type_t etype, IdArray vids) const override {
+    CHECK(IsValidIdArray(vids)) << "Invalid vertex id array.";
+    auto csrsubmat = aten::CSRSliceRows(adj_, vids);
+    auto coosubmat = aten::CSRToCOO(csrsubmat, false);
+    // Note that the row id in the csr submat is relabled, so
+    // we need to recover it using an index select.
+    auto row = aten::IndexSelect(vids, coosubmat.row);
+    return EdgeArray{row, coosubmat.col, coosubmat.data};
+  }
+
+  EdgeArray Edges(dgl_type_t etype, const std::string &order = "") const override {
+    CHECK(order.empty() || order == std::string("srcdst"))
+      << "CSR only support Edges of order \"srcdst\","
+      << " but got \"" << order << "\".";
+    const auto& coo = aten::CSRToCOO(adj_, false);
+    return EdgeArray{coo.row, coo.col, coo.data};
+  }
+
+  uint64_t InDegree(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for CSR graph.";
+    return {};
+  }
+
+  DegreeArray InDegrees(dgl_type_t etype, IdArray vids) const override {
+    LOG(INFO) << "Not enabled for CSR graph.";
+    return {};
+  }
+
+  uint64_t OutDegree(dgl_type_t etype, dgl_id_t vid) const override {
+    CHECK(HasVertex(0, vid)) << "Invalid src vertex id: " << vid;
+    return aten::CSRGetRowNNZ(adj_, vid);
+  }
+
+  DegreeArray OutDegrees(dgl_type_t etype, IdArray vids) const override {
+    CHECK(IsValidIdArray(vids)) << "Invalid vertex id array.";
+    return aten::CSRGetRowNNZ(adj_, vids);
+  }
+
+  DGLIdIters SuccVec(dgl_type_t etype, dgl_id_t vid) const override {
+    // TODO(minjie): This still assumes the data type and device context
+    //   of this graph. Should fix later.
+    const dgl_id_t* indptr_data = static_cast<dgl_id_t*>(adj_.indptr->data);
+    const dgl_id_t* indices_data = static_cast<dgl_id_t*>(adj_.indices->data);
+    const dgl_id_t start = indptr_data[vid];
+    const dgl_id_t end = indptr_data[vid + 1];
+    return DGLIdIters(indices_data + start, indices_data + end);
+  }
+
+  DGLIdIters OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const override {
+    // TODO(minjie): This still assumes the data type and device context
+    //   of this graph. Should fix later.
+    const dgl_id_t* indptr_data = static_cast<dgl_id_t*>(adj_.indptr->data);
+    const dgl_id_t* eid_data = static_cast<dgl_id_t*>(adj_.data->data);
+    const dgl_id_t start = indptr_data[vid];
+    const dgl_id_t end = indptr_data[vid + 1];
+    return DGLIdIters(eid_data + start, eid_data + end);
+  }
+
+  DGLIdIters PredVec(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for CSR graph.";
+    return {};
+  }
+
+  DGLIdIters InEdgeVec(dgl_type_t etype, dgl_id_t vid) const override {
+    LOG(INFO) << "Not enabled for CSR graph.";
+    return {};
+  }
+
+  std::vector<IdArray> GetAdj(
+      dgl_type_t etype, bool transpose, const std::string &fmt) const override {
+    CHECK(!transpose && fmt == "csr") << "Not valid adj format request.";
+    return {adj_.indptr, adj_.indices, adj_.data};
+  }
+
+  HeteroSubgraph VertexSubgraph(const std::vector<IdArray>& vids) const override {
+    CHECK_EQ(vids.size(), 2) << "Number of vertex types mismatch";
+    CHECK(IsValidIdArray(vids[0])) << "Invalid vertex id array.";
+    CHECK(IsValidIdArray(vids[1])) << "Invalid vertex id array.";
+    HeteroSubgraph subg;
+    const auto& submat = aten::CSRSliceMatrix(adj_, vids[0], vids[1]);
+    IdArray sub_eids = aten::Range(0, submat.data->shape[0], NumBits(), Context());
+    subg.graph = std::make_shared<CSR>(submat.num_rows, submat.num_cols,
+        submat.indptr, submat.indices, sub_eids);
+    subg.induced_vertices = vids;
+    subg.induced_edges.emplace_back(submat.data);
+    return subg;
+  }
+
+  HeteroSubgraph EdgeSubgraph(
+      const std::vector<IdArray>& eids, bool preserve_nodes = false) const override {
+    LOG(INFO) << "Not enabled for CSR graph.";
+    return {};
+  }
+
+  aten::CSRMatrix adj() const {
+    return adj_;
+  }
+
+ private:
+  /*! \brief internal adjacency matrix. Data array stores edge ids */
+  aten::CSRMatrix adj_;
+
+  /*! \brief multi-graph flag */
+  Lazy<bool> is_multigraph_;
+};
+
+//////////////////////////////////////////////////////////
+//
+// bipartite graph implementation
+//
+//////////////////////////////////////////////////////////
+
+DLContext Bipartite::Context() const {
+  return GetAny()->Context();
+}
+
+uint8_t Bipartite::NumBits() const {
+  return GetAny()->NumBits();
+}
+
+bool Bipartite::IsMultigraph() const {
+  return GetAny()->IsMultigraph();
+}
+
+uint64_t Bipartite::NumVertices(dgl_type_t vtype) const {
+  return GetAny()->NumVertices(vtype);
+}
+
+uint64_t Bipartite::NumEdges(dgl_type_t etype) const {
+  return GetAny()->NumEdges(etype);
+}
+
+bool Bipartite::HasVertex(dgl_type_t vtype, dgl_id_t vid) const {
+  return GetAny()->HasVertex(vtype, vid);
+}
+
+BoolArray Bipartite::HasVertices(dgl_type_t vtype, IdArray vids) const {
+  CHECK(IsValidIdArray(vids)) << "Invalid id array input";
+  return aten::LT(vids, NumVertices(vtype));
+}
+
+bool Bipartite::HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
+  if (in_csr_) {
+    return in_csr_->HasEdgeBetween(etype, dst, src);
+  } else {
+    return GetOutCSR()->HasEdgeBetween(etype, src, dst);
+  }
+}
+
+BoolArray Bipartite::HasEdgesBetween(
+    dgl_type_t etype, IdArray src, IdArray dst) const {
+  if (in_csr_) {
+    return in_csr_->HasEdgesBetween(etype, dst, src);
+  } else {
+    return GetOutCSR()->HasEdgesBetween(etype, src, dst);
+  }
+}
+
+IdArray Bipartite::Predecessors(dgl_type_t etype, dgl_id_t dst) const {
+  return GetInCSR()->Successors(etype, dst);
+}
+
+IdArray Bipartite::Successors(dgl_type_t etype, dgl_id_t src) const {
+  return GetOutCSR()->Successors(etype, src);
+}
+
+IdArray Bipartite::EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
+  if (in_csr_) {
+    return in_csr_->EdgeId(etype, dst, src);
+  } else {
+    return GetOutCSR()->EdgeId(etype, src, dst);
+  }
+}
+
+EdgeArray Bipartite::EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const {
+  if (in_csr_) {
+    EdgeArray edges = in_csr_->EdgeIds(etype, dst, src);
+    return EdgeArray{edges.dst, edges.src, edges.id};
+  } else {
+    return GetOutCSR()->EdgeIds(etype, src, dst);
+  }
+}
+
+std::pair<dgl_id_t, dgl_id_t> Bipartite::FindEdge(dgl_type_t etype, dgl_id_t eid) const {
+  return GetCOO()->FindEdge(etype, eid);
+}
+
+EdgeArray Bipartite::FindEdges(dgl_type_t etype, IdArray eids) const {
+  return GetCOO()->FindEdges(etype, eids);
+}
+
+EdgeArray Bipartite::InEdges(dgl_type_t etype, dgl_id_t vid) const {
+  const EdgeArray& ret = GetInCSR()->OutEdges(etype, vid);
+  return {ret.dst, ret.src, ret.id};
+}
+
+EdgeArray Bipartite::InEdges(dgl_type_t etype, IdArray vids) const {
+  const EdgeArray& ret = GetInCSR()->OutEdges(etype, vids);
+  return {ret.dst, ret.src, ret.id};
+}
+
+EdgeArray Bipartite::OutEdges(dgl_type_t etype, dgl_id_t vid) const {
+  return GetOutCSR()->OutEdges(etype, vid);
+}
+
+EdgeArray Bipartite::OutEdges(dgl_type_t etype, IdArray vids) const {
+  return GetOutCSR()->OutEdges(etype, vids);
+}
+
+EdgeArray Bipartite::Edges(dgl_type_t etype, const std::string &order) const {
+  if (order.empty()) {
+    // arbitrary order
+    if (in_csr_) {
+      // transpose
+      const auto& edges = in_csr_->Edges(etype, order);
+      return EdgeArray{edges.dst, edges.src, edges.id};
+    } else {
+      return GetAny()->Edges(etype, order);
+    }
+  } else if (order == std::string("srcdst")) {
+    // TODO(minjie): CSR only guarantees "src" to be sorted.
+    //   Maybe we should relax this requirement?
+    return GetOutCSR()->Edges(etype, order);
+  } else if (order == std::string("eid")) {
+    return GetCOO()->Edges(etype, order);
+  } else {
+    LOG(FATAL) << "Unsupported order request: " << order;
+  }
+  return {};
+}
+
+uint64_t Bipartite::InDegree(dgl_type_t etype, dgl_id_t vid) const {
+  return GetInCSR()->OutDegree(etype, vid);
+}
+
+DegreeArray Bipartite::InDegrees(dgl_type_t etype, IdArray vids) const {
+  return GetInCSR()->OutDegrees(etype, vids);
+}
+
+uint64_t Bipartite::OutDegree(dgl_type_t etype, dgl_id_t vid) const {
+  return GetOutCSR()->OutDegree(etype, vid);
+}
+
+DegreeArray Bipartite::OutDegrees(dgl_type_t etype, IdArray vids) const {
+  return GetOutCSR()->OutDegrees(etype, vids);
+}
+
+DGLIdIters Bipartite::SuccVec(dgl_type_t etype, dgl_id_t vid) const {
+  return GetOutCSR()->SuccVec(etype, vid);
+}
+
+DGLIdIters Bipartite::OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
+  return GetOutCSR()->OutEdgeVec(etype, vid);
+}
+
+DGLIdIters Bipartite::PredVec(dgl_type_t etype, dgl_id_t vid) const {
+  return GetInCSR()->SuccVec(etype, vid);
+}
+
+DGLIdIters Bipartite::InEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
+  return GetInCSR()->OutEdgeVec(etype, vid);
+}
+
+std::vector<IdArray> Bipartite::GetAdj(
+    dgl_type_t etype, bool transpose, const std::string &fmt) const {
+  // TODO(minjie): Our current semantics of adjacency matrix is row for dst nodes and col for
+  //   src nodes. Therefore, we need to flip the transpose flag. For example, transpose=False
+  //   is equal to in edge CSR.
+  //   We have this behavior because previously we use framework's SPMM and we don't cache
+  //   reverse adj. This is not intuitive and also not consistent with networkx's
+  //   to_scipy_sparse_matrix. With the upcoming custom kernel change, we should change the
+  //   behavior and make row for src and col for dst.
+  if (fmt == std::string("csr")) {
+    return transpose? GetOutCSR()->GetAdj(etype, false, "csr")
+      : GetInCSR()->GetAdj(etype, false, "csr");
+  } else if (fmt == std::string("coo")) {
+    return GetCOO()->GetAdj(etype, !transpose, fmt);
+  } else {
+    LOG(FATAL) << "unsupported adjacency matrix format: " << fmt;
+    return {};
+  }
+}
+
+HeteroSubgraph Bipartite::VertexSubgraph(const std::vector<IdArray>& vids) const {
+  // We prefer to generate a subgraph from out-csr.
+  auto sg = GetOutCSR()->VertexSubgraph(vids);
+  CSRPtr subcsr = std::dynamic_pointer_cast<CSR>(sg.graph);
+  HeteroSubgraph ret;
+  ret.graph = HeteroGraphPtr(new Bipartite(nullptr, subcsr, nullptr));
+  ret.induced_vertices = std::move(sg.induced_vertices);
+  ret.induced_edges = std::move(sg.induced_edges);
+  return ret;
+}
+
+HeteroSubgraph Bipartite::EdgeSubgraph(
+    const std::vector<IdArray>& eids, bool preserve_nodes) const {
+  auto sg = GetCOO()->EdgeSubgraph(eids, preserve_nodes);
+  COOPtr subcoo = std::dynamic_pointer_cast<COO>(sg.graph);
+  HeteroSubgraph ret;
+  ret.graph = HeteroGraphPtr(new Bipartite(nullptr, nullptr, subcoo));
+  ret.induced_vertices = std::move(sg.induced_vertices);
+  ret.induced_edges = std::move(sg.induced_edges);
+  return ret;
+}
+
+HeteroGraphPtr Bipartite::CreateFromCOO(
+    int64_t num_src, int64_t num_dst,
+    IdArray row, IdArray col) {
+  COOPtr coo(new COO(num_src, num_dst, row, col));
+  return HeteroGraphPtr(new Bipartite(nullptr, nullptr, coo));
+}
+
+HeteroGraphPtr Bipartite::CreateFromCSR(
+    int64_t num_src, int64_t num_dst,
+    IdArray indptr, IdArray indices, IdArray edge_ids) {
+  CSRPtr csr(new CSR(num_src, num_dst, indptr, indices, edge_ids));
+  return HeteroGraphPtr(new Bipartite(nullptr, csr, nullptr));
+}
+
+Bipartite::Bipartite(CSRPtr in_csr, CSRPtr out_csr, COOPtr coo)
+  : BaseHeteroGraph(kBipartiteMetaGraph), in_csr_(in_csr), out_csr_(out_csr), coo_(coo) {
+  CHECK(GetAny()) << "At least one graph structure should exist.";
+}
+
+Bipartite::CSRPtr Bipartite::GetInCSR() const {
+  if (!in_csr_) {
+    if (out_csr_) {
+      const auto& newadj = aten::CSRTranspose(out_csr_->adj());
+      const_cast<Bipartite*>(this)->in_csr_ = std::make_shared<CSR>(newadj);
+    } else {
+      CHECK(coo_) << "None of CSR, COO exist";
+      const auto& adj = coo_->adj();
+      const auto& newadj = aten::COOToCSR(
+          aten::COOMatrix{adj.num_cols, adj.num_rows, adj.col, adj.row});
+      const_cast<Bipartite*>(this)->in_csr_ = std::make_shared<CSR>(newadj);
+    }
+  }
+  return in_csr_;
+}
+
+/* !\brief Return out csr. If not exist, transpose the other one.*/
+Bipartite::CSRPtr Bipartite::GetOutCSR() const {
+  if (!out_csr_) {
+    if (in_csr_) {
+      const auto& newadj = aten::CSRTranspose(in_csr_->adj());
+      const_cast<Bipartite*>(this)->out_csr_ = std::make_shared<CSR>(newadj);
+    } else {
+      CHECK(coo_) << "None of CSR, COO exist";
+      const auto& newadj = aten::COOToCSR(coo_->adj());
+      const_cast<Bipartite*>(this)->out_csr_ = std::make_shared<CSR>(newadj);
+    }
+  }
+  return out_csr_;
+}
+
+/* !\brief Return coo. If not exist, create from csr.*/
+Bipartite::COOPtr Bipartite::GetCOO() const {
+  if (!coo_) {
+    if (in_csr_) {
+      const auto& newadj = aten::CSRToCOO(in_csr_->adj(), true);
+      const_cast<Bipartite*>(this)->coo_ = std::make_shared<COO>(
+          aten::COOMatrix{newadj.num_cols, newadj.num_rows, newadj.col, newadj.row});
+    } else {
+      CHECK(out_csr_) << "Both CSR are missing.";
+      const auto& newadj = aten::CSRToCOO(out_csr_->adj(), true);
+      const_cast<Bipartite*>(this)->coo_ = std::make_shared<COO>(newadj);
+    }
+  }
+  return coo_;
+}
+
+HeteroGraphPtr Bipartite::GetAny() const {
+  if (in_csr_) {
+    return in_csr_;
+  } else if (out_csr_) {
+    return out_csr_;
+  } else {
+    return coo_;
+  }
+}
+
+}  // namespace dgl

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()