[Hetero] New syntax (dmlc#824)

* WIP. remove graph arg in NodeBatch and EdgeBatch

* refactor: use graph adapter for scheduler

* WIP: recv

* draft impl

* stuck at bipartite

* bipartite->unitgraph; support dsttype == srctype

* pass test_query

* pass test_query

* pass test_view

* test apply

* pass udf message passing tests

* pass quan's test using builtins

* WIP: wildcard slicing

* new construct methods

* broken

* good

* add stack cross reducer

* fix bug; fix mx

* fix bug in csrmm2 when the CSR is not square

* lint

* removed FlattenedHeteroGraph class

* WIP

* prop nodes, prop edges, filter nodes/edges

* add DGLGraph tests to heterograph. Fix several bugs

* finish nx<->hetero graph conversion

* create bipartite from nx

* more spec on hetero/homo conversion

* silly fixes

* check node and edge types

* repr

* to api

* adj APIs

* inc

* fix some lints and bugs

* fix some lints

* hetero/homo conversion

* fix flatten test

* more spec in hetero_from_homo and test

* flatten using concat names

* WIP: creators

* rewrite hetero_from_homo in a more efficient way

* remove useless variables

* fix lint

* subgraphs and typed subgraphs

* lint & removed heterosubgraph class

* lint x2

* disable heterograph mutation test

* docstring update

* add edge id for nx graph test

* fix mx unittests

* fix bug

* try fix

* fix unittest when cross_reducer is stack

* fix ci

* fix nx bipartite bug; docstring

* fix scipy creation bug

* lint

* fix bug when converting heterograph from homograph

* fix bug in hetero_from_homo about ntype order

* trailing white

* docstring fixes for add_foo and data views

* docstring for relation slice

* to_hetero and to_homo with feature support

* lint

* lint

* DGLGraph compatibility

* incidence matrix & docstring fixes

* example string fixes

* feature in hetero_from_relations

* deduplication of edge types in to_hetero

* fix lint

* fix

include/dgl/base_heterograph.h

@@ -21,7 +21,9 @@ namespace dgl {
 
 // Forward declaration
 class BaseHeteroGraph;
+class FlattenedHeteroGraph;
 typedef std::shared_ptr<BaseHeteroGraph> HeteroGraphPtr;
+typedef std::shared_ptr<FlattenedHeteroGraph> FlattenedHeteroGraphPtr;
 struct HeteroSubgraph;
 
 /*!
@@ -46,10 +48,14 @@ class BaseHeteroGraph : public runtime::Object {
   ////////////////////////// query/operations on meta graph ////////////////////////
 
   /*! \return the number of vertex types */
-  virtual uint64_t NumVertexTypes() const = 0;
+  virtual uint64_t NumVertexTypes() const {
+    return meta_graph_->NumVertices();
+  }
 
   /*! \return the number of edge types */
-  virtual uint64_t NumEdgeTypes() const = 0;
+  virtual uint64_t NumEdgeTypes() const {
+    return meta_graph_->NumEdges();
+  }
 
   /*! \return the meta graph */
   virtual GraphPtr meta_graph() const {
@@ -351,6 +357,17 @@ class BaseHeteroGraph : public runtime::Object {
   virtual HeteroSubgraph EdgeSubgraph(
       const std::vector<IdArray>& eids, bool preserve_nodes = false) const = 0;
 
+  /*!
+   * \brief Convert the list of requested unitgraph graphs into a single unitgraph graph.
+   *
+   * \param etypes The list of edge type IDs.
+   * \return The flattened graph, with induced source/edge/destination types/IDs.
+   */
+  virtual FlattenedHeteroGraphPtr Flatten(const std::vector<dgl_type_t>& etypes) const {
+    LOG(FATAL) << "Flatten operation unsupported";
+    return nullptr;
+  }
+
   static constexpr const char* _type_key = "graph.HeteroGraph";
   DGL_DECLARE_OBJECT_TYPE_INFO(BaseHeteroGraph, runtime::Object);
 
@@ -381,6 +398,62 @@ struct HeteroSubgraph : public runtime::Object {
   DGL_DECLARE_OBJECT_TYPE_INFO(HeteroSubgraph, runtime::Object);
 };
 
+/*! \brief The flattened heterograph */
+struct FlattenedHeteroGraph : public runtime::Object {
+  /*! \brief The graph */
+  HeteroGraphRef graph;
+  /*!
+   * \brief Mapping from source node ID to node type in parent graph
+   * \note The induced type array guarantees that the same type always appear contiguously.
+   */
+  IdArray induced_srctype;
+  /*!
+   * \brief The set of node types in parent graph appearing in source nodes.
+   */
+  IdArray induced_srctype_set;
+  /*! \brief Mapping from source node ID to local node ID in parent graph */
+  IdArray induced_srcid;
+  /*!
+   * \brief Mapping from edge ID to edge type in parent graph
+   * \note The induced type array guarantees that the same type always appear contiguously.
+   */
+  IdArray induced_etype;
+  /*!
+   * \brief The set of edge types in parent graph appearing in edges.
+   */
+  IdArray induced_etype_set;
+  /*! \brief Mapping from edge ID to local edge ID in parent graph */
+  IdArray induced_eid;
+  /*!
+   * \brief Mapping from destination node ID to node type in parent graph
+   * \note The induced type array guarantees that the same type always appear contiguously.
+   */
+  IdArray induced_dsttype;
+  /*!
+   * \brief The set of node types in parent graph appearing in destination nodes.
+   */
+  IdArray induced_dsttype_set;
+  /*! \brief Mapping from destination node ID to local node ID in parent graph */
+  IdArray induced_dstid;
+
+  void VisitAttrs(runtime::AttrVisitor *v) final {
+    v->Visit("graph", &graph);
+    v->Visit("induced_srctype", &induced_srctype);
+    v->Visit("induced_srctype_set", &induced_srctype_set);
+    v->Visit("induced_srcid", &induced_srcid);
+    v->Visit("induced_etype", &induced_etype);
+    v->Visit("induced_etype_set", &induced_etype_set);
+    v->Visit("induced_eid", &induced_eid);
+    v->Visit("induced_dsttype", &induced_dsttype);
+    v->Visit("induced_dsttype_set", &induced_dsttype_set);
+    v->Visit("induced_dstid", &induced_dstid);
+  }
+
+  static constexpr const char* _type_key = "graph.FlattenedHeteroGraph";
+  DGL_DECLARE_OBJECT_TYPE_INFO(FlattenedHeteroGraph, runtime::Object);
+};
+DGL_DEFINE_OBJECT_REF(FlattenedHeteroGraphRef, FlattenedHeteroGraph);
+
 // Define HeteroSubgraphRef
 DGL_DEFINE_OBJECT_REF(HeteroSubgraphRef, HeteroSubgraph);
 

include/dgl/runtime/object.h

@@ -18,6 +18,7 @@ namespace runtime {
 // forward declaration
 class Object;
 class ObjectRef;
+class NDArray;
 
 /*!
  * \brief Visitor class to each object attribute.
@@ -33,6 +34,7 @@ class AttrVisitor {
   virtual void Visit(const char* key, bool* value) = 0;
   virtual void Visit(const char* key, std::string* value) = 0;
   virtual void Visit(const char* key, ObjectRef* value) = 0;
+  virtual void Visit(const char* key, NDArray* value) = 0;
   template<typename ENum,
            typename = typename std::enable_if<std::is_enum<ENum>::value>::type>
   void Visit(const char* key, ENum* ptr) {

python/dgl/__init__.py

@@ -13,9 +13,10 @@
 from ._ffi.function import register_func, get_global_func, list_global_func_names, extract_ext_funcs
 from ._ffi.base import DGLError, __version__
 
-from .base import ALL
+from .base import ALL, NTYPE, NID, ETYPE, EID
 from .backend import load_backend
 from .batched_graph import *
+from .convert import *
 from .graph import DGLGraph
 from .heterograph import DGLHeteroGraph
 from .nodeflow import *

python/dgl/base.py

@@ -8,6 +8,13 @@
 
 # A special symbol for selecting all nodes or edges.
 ALL = "__ALL__"
+# An alias for [:]
+SLICE_FULL = slice(None, None, None)
+# Reserved column names for storing parent node/edge types and IDs in flattened heterographs
+NTYPE = '_TYPE'
+NID = '_ID'
+ETYPE = '_TYPE'
+EID = '_ID'
 
 def is_all(arg):
     """Return true if the argument is a special symbol for all nodes or edges."""