perf(distributed): derive with_edge from edge_feature_info instead of hardcoding True

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: kmontemayor

gigl/distributed/base_dist_loader.py

@@ -416,7 +416,7 @@ def create_sampling_config(
             batch_size=batch_size,
             shuffle=shuffle,
             drop_last=drop_last,
-            with_edge=True,
+            with_edge=dataset_schema.edge_feature_info is not None,
             collect_features=True,
             with_neg=False,
             with_weight=with_weight,

tests/unit/distributed/distributed_neighborloader_test.py

@@ -1,5 +1,5 @@
 import unittest
-from collections.abc import Mapping
+from collections.abc import Mapping, MutableMapping
 
 import torch
 import torch.multiprocessing as mp
@@ -398,6 +398,31 @@ def _run_cora_supervised_node_classification(
     shutdown_rpc()
 
 
+def _run_featureless_edge_ids_absent(
+    _,
+    dataset: DistDataset,
+    holder: MutableMapping,
+):
+    # A featureless dataset should sample with with_edge=False, so GLT never
+    # attaches sampled edge ids (``data.edge``) to the produced batches.
+    create_test_process_group()
+    loader = DistNeighborLoader(
+        dataset=dataset,
+        num_neighbors=[2, 2],
+        pin_memory_device=torch.device("cpu"),
+    )
+    count = 0
+    edge_ids_absent = True
+    for datum in loader:
+        assert isinstance(datum, Data)
+        if getattr(datum, "edge", None) is not None:
+            edge_ids_absent = False
+        count += 1
+    holder["edge_ids_absent"] = edge_ids_absent
+    holder["count"] = count
+    shutdown_rpc()
+
+
 class DistributedNeighborLoaderTest(TestCase):
     def setUp(self):
         super().setUp()
@@ -861,5 +886,93 @@ def test_distributed_neighbor_loader_invalid_inputs_colocated(
             DistNeighborLoader(**kwargs)
 
 
+class WithEdgeDerivationTest(TestCase):
+    """Regression coverage for deriving ``with_edge`` from edge-feature presence."""
+
+    def test_config_with_edge_false_when_no_edge_features(self) -> None:
+        from gigl.distributed.base_dist_loader import BaseDistLoader
+        from gigl.distributed.utils.neighborloader import DatasetSchema
+
+        schema = DatasetSchema(
+            is_homogeneous_with_labeled_edge_type=False,
+            edge_types=None,
+            node_feature_info=None,
+            edge_feature_info=None,  # no edge features
+            edge_dir="out",
+        )
+        config = BaseDistLoader.create_sampling_config(
+            num_neighbors=[2, 2], dataset_schema=schema
+        )
+        self.assertFalse(config.with_edge)
+
+    def test_config_with_edge_true_when_edge_features_present(self) -> None:
+        from gigl.distributed.base_dist_loader import BaseDistLoader
+        from gigl.distributed.utils.neighborloader import DatasetSchema
+        from gigl.types.graph import FeatureInfo
+
+        schema = DatasetSchema(
+            is_homogeneous_with_labeled_edge_type=False,
+            edge_types=None,
+            node_feature_info=None,
+            edge_feature_info=FeatureInfo(dim=4, dtype=torch.float32),
+            edge_dir="out",
+        )
+        config = BaseDistLoader.create_sampling_config(
+            num_neighbors=[2, 2], dataset_schema=schema
+        )
+        self.assertTrue(config.with_edge)
+
+    def test_config_with_edge_true_for_heterogeneous_edge_feature_dict(self) -> None:
+        from gigl.distributed.base_dist_loader import BaseDistLoader
+        from gigl.distributed.utils.neighborloader import DatasetSchema
+        from gigl.types.graph import FeatureInfo
+
+        schema = DatasetSchema(
+            is_homogeneous_with_labeled_edge_type=False,
+            edge_types=[_USER_TO_STORY, _STORY_TO_USER],
+            node_feature_info=None,
+            edge_feature_info={
+                _USER_TO_STORY: FeatureInfo(dim=4, dtype=torch.float32),
+            },
+            edge_dir="out",
+        )
+        config = BaseDistLoader.create_sampling_config(
+            num_neighbors=[2, 2], dataset_schema=schema
+        )
+        self.assertTrue(config.with_edge)
+
+    def test_featureless_dataset_produces_batches_without_edge_ids(self) -> None:
+        # Homogeneous dataset with node features but no edge features: the loader
+        # must still yield batches, and none carry sampled edge ids.
+        partition_output = PartitionOutput(
+            node_partition_book=torch.zeros(5),
+            edge_partition_book=torch.zeros(5),
+            partitioned_edge_index=GraphPartitionData(
+                edge_index=torch.tensor([[0, 1, 2, 3, 4], [1, 2, 3, 4, 0]]),
+                edge_ids=None,
+            ),
+            partitioned_node_features=FeaturePartitionData(
+                feats=torch.zeros(5, 2), ids=torch.arange(5)
+            ),
+            partitioned_edge_features=None,
+            partitioned_positive_labels=None,
+            partitioned_negative_labels=None,
+            partitioned_node_labels=None,
+        )
+        dataset = DistDataset(rank=0, world_size=1, edge_dir="out")
+        dataset.build(partition_output=partition_output)
+
+        manager = mp.Manager()
+        holder: MutableMapping = manager.dict()
+        proc = mp.spawn(
+            fn=_run_featureless_edge_ids_absent,
+            args=(dataset, holder),
+            join=False,
+        )
+        proc.join(timeout=120)
+        self.assertGreater(holder["count"], 0)
+        self.assertTrue(holder["edge_ids_absent"])
+
+
 if __name__ == "__main__":
     absltest.main()