test(distributed): with_edge derivation across weighted, PPR, and graph-store modes

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

Author: kmontemayor

tests/unit/distributed/dist_ppr_sampler_test.py

@@ -316,6 +316,13 @@ def _run_ppr_loader_correctness_check(
 
     dataset = create_homogeneous_dataset(edge_index=_TEST_EDGE_INDEX, edge_dir=edge_dir)
 
+    # This dataset has no edge features, so with_edge is derived as False. PPR
+    # outputs come from metadata (ppr_edge_index/ppr_edge_attr), not sampled edge
+    # ids, so with_edge must not affect the PPR results verified below.
+    assert dataset.edge_features is None, (
+        "PPR regression fixture must be edge-feature-free so with_edge is False."
+    )
+
     loader = DistNeighborLoader(
         dataset=dataset,
         num_neighbors=[],  # Unused by PPR sampler; required by interface

tests/unit/distributed/distributed_weighted_sampling_test.py

@@ -327,6 +327,12 @@ def _run_weighted_sampling_correctness_homogeneous(
     for datum in loader:
         assert isinstance(datum, Data), f"Expected Data, got {type(datum)}"
         assert datum.x is not None, "Node features missing from sampled subgraph"
+        # The graph has no edge features, so with_edge is derived as False and GLT
+        # must not attach sampled edge ids -- yet weighted sampling still works.
+        assert getattr(datum, "edge", None) is None, (
+            "Featureless weighted dataset should sample with with_edge=False, "
+            "but sampled edge ids were attached to the batch."
+        )
         # Bad nodes have feature 0.0; hub and good nodes have feature > 0.
         bad_mask = datum.x[:, 0] == 0.0
         assert not bad_mask.any(), (

tests/unit/distributed/graph_store/remote_dist_dataset_test.py

@@ -200,6 +200,36 @@ def test_graph_metadata_getters_homogeneous(self, mock_request):
         self.assertIsNone(remote_dataset.fetch_edge_types())
         self.assertIsNone(remote_dataset.fetch_node_types())
 
+    def test_fetch_edge_feature_info_populated_from_server(self, mock_request):
+        """Edge feature info is fetched from server rank 0 for the compute-side schema.
+
+        The ``with_edge`` sampling flag is derived on the compute side as
+        ``dataset_schema.edge_feature_info is not None`` (see
+        ``BaseDistLoader.create_sampling_config``). This test confirms the field
+        is populated from the storage cluster when the graph has edge features,
+        so that derivation yields ``with_edge=True`` off-node.
+        """
+        global _test_server
+        num_edges = DEFAULT_HOMOGENEOUS_EDGE_INDEX.shape[1]
+        dataset_with_edge_features = create_homogeneous_dataset(
+            edge_index=DEFAULT_HOMOGENEOUS_EDGE_INDEX,
+            node_features=torch.zeros(10, 3),
+            edge_features=torch.zeros(num_edges, 4),
+        )
+        _test_server = DistServer(dataset_with_edge_features)
+        dist_server_module._dist_server = _test_server
+
+        cluster_info = _create_mock_graph_store_info()
+        remote_dataset = RemoteDistDataset(cluster_info=cluster_info, local_rank=0)
+
+        edge_feature_info = remote_dataset.fetch_edge_feature_info()
+        self.assertEqual(
+            edge_feature_info,
+            FeatureInfo(dim=4, dtype=torch.float32),
+        )
+        # The derivation used off-node by create_sampling_config.
+        self.assertTrue(edge_feature_info is not None)
+
     def test_cluster_info_property(self, mock_request):
         cluster_info = _create_mock_graph_store_info(
             num_storage_nodes=3, num_compute_nodes=2