[BugFix] fix in_degree/out_degree computation logic (#3477)

* [BugFix] fix in/out degree computation

* add unit tests

src/graph/unit_graph.cc

@@ -1000,30 +1000,38 @@ EdgeArray UnitGraph::Edges(dgl_type_t etype, const std::string &order) const {
 uint64_t UnitGraph::InDegree(dgl_type_t etype, dgl_id_t vid) const {
   SparseFormat fmt = SelectFormat(CSC_CODE);
   const auto ptr = GetFormat(fmt);
-  if (fmt == SparseFormat::kCSC)
-    return ptr->OutDegree(etype, vid);
-  else
-    return ptr->InDegree(etype, vid);
+  CHECK(fmt == SparseFormat::kCSC || fmt == SparseFormat::kCOO)
+      << "In degree cannot be computed as neither CSC nor COO format is "
+         "allowed for this graph. Please enable one of them at least.";
+  return fmt == SparseFormat::kCSC ? ptr->OutDegree(etype, vid)
+                                   : ptr->InDegree(etype, vid);
 }
 
 DegreeArray UnitGraph::InDegrees(dgl_type_t etype, IdArray vids) const {
   SparseFormat fmt = SelectFormat(CSC_CODE);
   const auto ptr = GetFormat(fmt);
-  if (fmt == SparseFormat::kCSC)
-    return ptr->OutDegrees(etype, vids);
-  else
-    return ptr->InDegrees(etype, vids);
+  CHECK(fmt == SparseFormat::kCSC || fmt == SparseFormat::kCOO)
+      << "In degree cannot be computed as neither CSC nor COO format is "
+         "allowed for this graph. Please enable one of them at least.";
+  return fmt == SparseFormat::kCSC ? ptr->OutDegrees(etype, vids)
+                                   : ptr->InDegrees(etype, vids);
 }
 
 uint64_t UnitGraph::OutDegree(dgl_type_t etype, dgl_id_t vid) const {
   SparseFormat fmt = SelectFormat(CSR_CODE);
   const auto ptr = GetFormat(fmt);
+  CHECK(fmt == SparseFormat::kCSR || fmt == SparseFormat::kCOO)
+      << "Out degree cannot be computed as neither CSR nor COO format is "
+         "allowed for this graph. Please enable one of them at least.";
   return ptr->OutDegree(etype, vid);
 }
 
 DegreeArray UnitGraph::OutDegrees(dgl_type_t etype, IdArray vids) const {
   SparseFormat fmt = SelectFormat(CSR_CODE);
   const auto ptr = GetFormat(fmt);
+  CHECK(fmt == SparseFormat::kCSR || fmt == SparseFormat::kCOO)
+      << "Out degree cannot be computed as neither CSR nor COO format is "
+         "allowed for this graph. Please enable one of them at least.";
   return ptr->OutDegrees(etype, vids);
 }
 

tests/compute/test_graph.py

@@ -376,6 +376,40 @@ def test_default_types():
     assert dg.ntypes == g.ntypes
     assert dg.etypes == g.etypes
 
+
+def test_formats():
+    g = dgl.rand_graph(10, 20)
+    # in_degrees works if coo or csc available
+    # out_degrees works if coo or csr available
+    try:
+        g.in_degrees()
+        g.out_degrees()
+        g.formats('coo').in_degrees()
+        g.formats('coo').out_degrees()
+        g.formats('csc').in_degrees()
+        g.formats('csr').out_degrees()
+        fail = False
+    except DGLError:
+        fail = True
+    finally:
+        assert not fail
+    # in_degrees NOT works if csc available only
+    try:
+        g.formats('csc').out_degrees()
+        fail = True
+    except DGLError:
+        fail = False
+    finally:
+        assert not fail
+    # out_degrees NOT works if csr available only
+    try:
+        g.formats('csr').in_degrees()
+        fail = True
+    except DGLError:
+        fail = False
+    finally:
+        assert not fail
+    
 if __name__ == '__main__':
     test_query()
     test_mutation()
@@ -385,3 +419,4 @@ if __name__ == '__main__':
     test_hypersparse_query()
     test_is_sorted()
     test_default_types()
+    test_formats()

tests/cpp/test_unit_graph.cc

@@ -67,6 +67,52 @@ aten::COOMatrix COO1(DLContext ctx) {
 template aten::COOMatrix COO1<int32_t>(DLContext ctx);
 template aten::COOMatrix COO1<int64_t>(DLContext ctx);
 
+template <typename IdType> void _TestUnitGraph_InOutDegrees(DLContext ctx) {
+  /*
+  InDegree(s) is available only if COO or CSC formats permitted.
+  OutDegree(s) is available only if COO or CSR formats permitted.
+  */
+
+  // COO
+  {
+    const aten::COOMatrix &coo = COO1<IdType>(ctx);
+    auto &&g = CreateFromCOO(2, coo, COO_CODE);
+    ASSERT_EQ(g->InDegree(0, 0), 1);
+    auto &&nids = aten::Range(0, g->NumVertices(0), g->NumBits(), g->Context());
+    ASSERT_TRUE(ArrayEQ<IdType>(
+        g->InDegrees(0, nids),
+        aten::VecToIdArray<IdType>({1, 2}, g->NumBits(), g->Context())));
+    ASSERT_EQ(g->OutDegree(0, 0), 2);
+    ASSERT_TRUE(ArrayEQ<IdType>(
+        g->OutDegrees(0, nids),
+        aten::VecToIdArray<IdType>({2, 1}, g->NumBits(), g->Context())));
+  }
+  // CSC
+  {
+    const aten::CSRMatrix &csr = CSR1<IdType>(ctx);
+    auto &&g = CreateFromCSC(2, csr, CSC_CODE);
+    ASSERT_EQ(g->InDegree(0, 0), 1);
+    auto &&nids = aten::Range(0, g->NumVertices(0), g->NumBits(), g->Context());
+    ASSERT_TRUE(ArrayEQ<IdType>(
+        g->InDegrees(0, nids),
+        aten::VecToIdArray<IdType>({1, 2, 1}, g->NumBits(), g->Context())));
+    EXPECT_ANY_THROW(g->OutDegree(0, 0));
+    EXPECT_ANY_THROW(g->OutDegrees(0, nids));
+  }
+  // CSR
+  {
+    const aten::CSRMatrix &csr = CSR1<IdType>(ctx);
+    auto &&g = CreateFromCSR(2, csr, CSR_CODE);
+    ASSERT_EQ(g->OutDegree(0, 0), 1);
+    auto &&nids = aten::Range(0, g->NumVertices(0), g->NumBits(), g->Context());
+    ASSERT_TRUE(ArrayEQ<IdType>(
+        g->OutDegrees(0, nids),
+        aten::VecToIdArray<IdType>({1, 2, 1, 2}, g->NumBits(), g->Context())));
+    EXPECT_ANY_THROW(g->InDegree(0, 0));
+    EXPECT_ANY_THROW(g->InDegrees(0, nids));
+  }
+}
+
 template <typename IdType>
 void _TestUnitGraph(DLContext ctx) {
   const aten::CSRMatrix &csr = CSR1<IdType>(ctx);
@@ -340,6 +386,15 @@ TEST(UniGraphTest, TestUnitGraph_CopyTo) {
 #endif
 }
 
+TEST(UniGraphTest, TestUnitGraph_InOutDegrees) {
+  _TestUnitGraph_InOutDegrees<int32_t>(CPU);
+  _TestUnitGraph_InOutDegrees<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestUnitGraph_InOutDegrees<int32_t>(GPU);
+  _TestUnitGraph_InOutDegrees<int64_t>(GPU);
+#endif
+}
+
 TEST(UniGraphTest, TestUnitGraph_Create) {
   _TestUnitGraph<int32_t>(CPU);
   _TestUnitGraph<int64_t>(CPU);