[GraphBolt] enrich node types for input/output nodes of sampled subgraph (#6348)

Author: Rhett-Ying

python/dgl/graphbolt/impl/csc_sampling_graph.py

@@ -242,8 +242,6 @@ def _convert_to_sampled_subgraph(
                 src_ntype_id = self.metadata.node_type_to_id[src_ntype]
                 dst_ntype_id = self.metadata.node_type_to_id[dst_ntype]
                 mask = type_per_edge == etype_id
-                if mask.count_nonzero() == 0:
-                    continue
                 hetero_row = row[mask] - self.node_type_offset[src_ntype_id]
                 hetero_column = (
                     column[mask] - self.node_type_offset[dst_ntype_id]

python/dgl/graphbolt/impl/neighbor_sampler.py

@@ -78,6 +78,7 @@ def __init__(
         3
         """
         super().__init__(datapipe)
+        self.graph = graph
         # Convert fanouts to a list of tensors.
         self.fanouts = []
         for fanout in fanouts:
@@ -91,6 +92,13 @@ def __init__(
     def _sample_subgraphs(self, seeds):
         subgraphs = []
         num_layers = len(self.fanouts)
+        # Enrich seeds with all node types.
+        if isinstance(seeds, dict):
+            ntypes = list(self.graph.metadata.node_type_to_id.keys())
+            seeds = {
+                ntype: seeds.get(ntype, torch.LongTensor([]))
+                for ntype in ntypes
+            }
         for hop in range(num_layers):
             subgraph = self.sampler(
                 seeds,

tests/python/pytorch/graphbolt/impl/test_csc_sampling_graph.py

@@ -584,8 +584,12 @@ def test_sample_neighbors_hetero(labor):
             torch.LongTensor([0, 2]),
             torch.LongTensor([0, 0]),
         ),
+        "n1:e1:n2": (
+            torch.LongTensor([]),
+            torch.LongTensor([]),
+        ),
     }
-    assert len(subgraph.node_pairs) == 1
+    assert len(subgraph.node_pairs) == 2
     for etype, pairs in expected_node_pairs.items():
         assert torch.equal(subgraph.node_pairs[etype][0], pairs[0])
         assert torch.equal(subgraph.node_pairs[etype][1], pairs[1])

tests/python/pytorch/graphbolt/test_subgraph_sampler.py

@@ -129,6 +129,23 @@ def get_hetero_graph():
     )
 
 
+@pytest.mark.parametrize("labor", [False, True])
+def test_SubgraphSampler_Node_Hetero(labor):
+    graph = get_hetero_graph()
+    itemset = gb.ItemSetDict(
+        {"n2": gb.ItemSet(torch.arange(3), names="seed_nodes")}
+    )
+    item_sampler = gb.ItemSampler(itemset, batch_size=2)
+    num_layer = 2
+    fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
+    Sampler = gb.LayerNeighborSampler if labor else gb.NeighborSampler
+    sampler_dp = Sampler(item_sampler, graph, fanouts)
+    assert len(list(sampler_dp)) == 2
+    for minibatch in sampler_dp:
+        blocks = minibatch.to_dgl_blocks()
+        assert len(blocks) == num_layer
+
+
 @pytest.mark.parametrize("labor", [False, True])
 def test_SubgraphSampler_Link_Hetero(labor):
     graph = get_hetero_graph()