[GraphBolt] Add NOT deduplication func to homo graph. (#6436)

Co-authored-by: Ubuntu <ubuntu@ip-172-31-0-133.us-west-2.compute.internal>

python/dgl/graphbolt/impl/csc_sampling_graph.py

@@ -308,12 +308,13 @@ class CSCSamplingGraph(SamplingGraph):
         # TODO: change the result to 'SampledSubgraphImpl'.
         return self._c_csc_graph.in_subgraph(nodes)
 
-    def _convert_to_sampled_subgraph(
+    def _convert_to_sampled_subgraph_node_pairs_and_eids(
         self,
         C_sampled_subgraph: torch.ScriptObject,
     ):
-        """An internal function used to convert a fused homogeneous sampled
-        subgraph to general struct 'SampledSubgraphImpl'."""
+        """An internal function used to generate node pairs and eids of
+        structure `SampledSubgraphImpl` from a fused homogeneous sampled
+        subgraph."""
         column_num = (
             C_sampled_subgraph.indptr[1:] - C_sampled_subgraph.indptr[:-1]
         )
@@ -357,6 +358,48 @@ class CSCSamplingGraph(SamplingGraph):
             node_pairs=node_pairs, original_edge_ids=original_edge_ids
         )
 
+    def _convert_to_sampled_subgraph(
+        self,
+        C_sampled_subgraph: torch.ScriptObject,
+    ):
+        """An internal function used to convert a fused homogeneous sampled
+        subgraph to general struct 'SampledSubgraphImpl'."""
+        type_per_edge = C_sampled_subgraph.type_per_edge
+        if type_per_edge is None:
+            column_num = (
+                C_sampled_subgraph.indptr[1:] - C_sampled_subgraph.indptr[:-1]
+            )
+            column = torch.arange(len(column_num)).repeat_interleave(column_num)
+            row = torch.arange(len(C_sampled_subgraph.indices)) + len(
+                column_num
+            )
+            original_edge_ids = C_sampled_subgraph.original_edge_ids
+            has_original_eids = (
+                self.edge_attributes is not None
+                and ORIGINAL_EDGE_ID in self.edge_attributes
+            )
+            if has_original_eids and original_edge_ids is not None:
+                original_edge_ids = self.edge_attributes[ORIGINAL_EDGE_ID][
+                    original_edge_ids
+                ]
+            # The sampled graph is already a homogeneous graph.
+            node_pairs = (row, column)
+            return SampledSubgraphImpl(
+                node_pairs=node_pairs,
+                original_edge_ids=original_edge_ids,
+                original_column_node_ids=(
+                    C_sampled_subgraph.original_column_node_ids
+                ),
+                original_row_node_ids=torch.cat(
+                    (
+                        C_sampled_subgraph.original_column_node_ids,
+                        C_sampled_subgraph.indices,
+                    )
+                ),
+            )
+        else:
+            raise RuntimeError("Not implemented yet.")
+
     def _convert_to_homogeneous_nodes(self, nodes):
         homogeneous_nodes = []
         for ntype, ids in nodes.items():
@@ -370,6 +413,7 @@ class CSCSamplingGraph(SamplingGraph):
         fanouts: torch.Tensor,
         replace: bool = False,
         probs_name: Optional[str] = None,
+        deduplicate=True,
     ) -> SampledSubgraphImpl:
         """Sample neighboring edges of the given nodes and return the induced
         subgraph.
@@ -409,6 +453,10 @@ class CSCSamplingGraph(SamplingGraph):
             corresponding to each neighboring edge of a node. It must be a 1D
             floating-point or boolean tensor, with the number of elements
             equalling the total number of edges.
+        deduplicate: bool
+            Boolean indicating whether seeds between hops will be deduplicated.
+            If True, the same elements in seeds will be removed to only one.
+            Otherwise, the same elements will be remained.
         Returns
         -------
         SampledSubgraphImpl
@@ -440,7 +488,11 @@ class CSCSamplingGraph(SamplingGraph):
         C_sampled_subgraph = self._sample_neighbors(
             nodes, fanouts, replace, probs_name
         )
-
+        if deduplicate:
+            return self._convert_to_sampled_subgraph_node_pairs_and_eids(
+                C_sampled_subgraph
+            )
+        else:
             return self._convert_to_sampled_subgraph(C_sampled_subgraph)
 
     def _check_sampler_arguments(self, nodes, fanouts, probs_name):
@@ -548,6 +600,7 @@ class CSCSamplingGraph(SamplingGraph):
         fanouts: torch.Tensor,
         replace: bool = False,
         probs_name: Optional[str] = None,
+        deduplicate=True,
     ) -> SampledSubgraphImpl:
         """Sample neighboring edges of the given nodes and return the induced
         subgraph via layer-neighbor sampling from the NeurIPS 2023 paper
@@ -589,6 +642,10 @@ class CSCSamplingGraph(SamplingGraph):
             corresponding to each neighboring edge of a node. It must be a 1D
             floating-point or boolean tensor, with the number of elements
             equalling the total number of edges.
+        deduplicate: bool
+            Boolean indicating whether seeds between hops will be deduplicated.
+            If True, the same elements in seeds will be deleted to only one.
+            Otherwise, the same elements will be remained.
         Returns
         -------
         SampledSubgraphImpl
@@ -630,7 +687,11 @@ class CSCSamplingGraph(SamplingGraph):
             has_original_eids,
             probs_name,
         )
-
+        if deduplicate:
+            return self._convert_to_sampled_subgraph_node_pairs_and_eids(
+                C_sampled_subgraph
+            )
+        else:
             return self._convert_to_sampled_subgraph(C_sampled_subgraph)
 
     def sample_negative_edges_uniform(

python/dgl/graphbolt/impl/neighbor_sampler.py

@@ -116,14 +116,13 @@ class NeighborSampler(SubgraphSampler):
                 self.fanouts[hop],
                 self.replace,
                 self.prob_name,
+                self.deduplicate,
             )
             if self.deduplicate:
                 (
                     original_row_node_ids,
                     compacted_node_pairs,
                 ) = unique_and_compact_node_pairs(subgraph.node_pairs, seeds)
-            else:
-                raise RuntimeError("Not implemented yet.")
                 subgraph = SampledSubgraphImpl(
                     node_pairs=compacted_node_pairs,
                     original_column_node_ids=seeds,
@@ -131,7 +130,7 @@ class NeighborSampler(SubgraphSampler):
                     original_edge_ids=subgraph.original_edge_ids,
                 )
             subgraphs.insert(0, subgraph)
-            seeds = original_row_node_ids
+            seeds = subgraph.original_row_node_ids
         return seeds, subgraphs
 
 

tests/python/pytorch/graphbolt/test_subgraph_sampler.py

+import dgl
 import dgl.graphbolt as gb
 import gb_test_utils
 import pytest
@@ -192,3 +193,36 @@ def test_SubgraphSampler_Link_Hetero_With_Negative(labor):
     Sampler = gb.LayerNeighborSampler if labor else gb.NeighborSampler
     neighbor_dp = Sampler(negative_dp, graph, fanouts)
     assert len(list(neighbor_dp)) == 5
+
+
+@pytest.mark.parametrize("labor", [False, True])
+def test_test_SubgraphSampler_without_dedpulication(labor):
+    graph = dgl.graph(
+        ([5, 0, 1, 5, 6, 7, 2, 2, 4], [0, 1, 2, 2, 2, 2, 3, 4, 4])
+    )
+    graph = gb.from_dglgraph(graph, True)
+    seed_nodes = torch.LongTensor([0, 3, 4])
+
+    itemset = gb.ItemSet(seed_nodes, names="seed_nodes")
+    item_sampler = gb.ItemSampler(itemset, batch_size=len(seed_nodes))
+    num_layer = 2
+    fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
+
+    Sampler = gb.LayerNeighborSampler if labor else gb.NeighborSampler
+    datapipe = Sampler(item_sampler, graph, fanouts, deduplicate=False)
+
+    length = [17, 7]
+    compacted_dst = [
+        torch.tensor([0, 1, 2, 2, 4, 4, 5, 5, 6, 6]),
+        torch.tensor([0, 1, 2, 2]),
+    ]
+    seeds = [torch.tensor([0, 3, 4, 5, 2, 2, 4]), torch.tensor([0, 3, 4])]
+    for data in datapipe:
+        for step, sampled_subgraph in enumerate(data.sampled_subgraphs):
+            assert len(sampled_subgraph.original_row_node_ids) == length[step]
+            assert torch.equal(
+                sampled_subgraph.node_pairs[1], compacted_dst[step]
+            )
+            assert torch.equal(
+                sampled_subgraph.original_column_node_ids, seeds[step]
+            )