Drop edge features from schema.

dgf/src/learning/ten_lines/node_prediction_test.py

@@ -1712,6 +1712,40 @@ def test_temporal_training_no_edgeset_timestamp_plan_is_temporal(self):
     self.assertTrue(model.data().temporal_sampling)
     self.assertTrue(_sampling_plan(model).temporal_sampling)
 
+  def test_temporal_training_with_edge_timeseries_features(self):
+    # Regression test: Edge timeseries features used to raise a KeyError
+    # during temporal sampling.
+    graph, schema = gen_test_graph.generate_temporal_in_memory_graph(
+        include_e2=False, include_edge_timeseries=True
+    )
+    schema.node_sets["n1"].features["label"] = schema_lib.FeatureSchema(
+        format=schema_lib.FeatureFormat.FLOAT_32,
+        semantic=schema_lib.FeatureSemantic.NUMERICAL,
+    )
+    graph.node_sets["n1"].features["label"] = np.array(
+        [1.0, 2.0, 3.0, 4.0], dtype=np.float32
+    )
+
+    model = node_prediction_lib.train_node_model(
+        graph=graph,
+        valid_graph=graph,
+        schema=schema,
+        target_nodeset="n1",
+        target_column="label",
+        time_aware=True,
+        num_train_steps=2,
+        valid_every_n_steps=1,
+        batch_size=2,
+        sampling_width=2,
+        num_sampling_hops=1,
+        node_embedding_dim=4,
+        num_layers=1,
+        verbose=0,
+    )
+    self.assertTrue(model.data().temporal_sampling)
+    predictions = model.predict(graph, seed_node_idxs=np.array([0, 1, 2, 3]))
+    self.assertEqual(predictions.shape, (4,))
+
 
 class NodePredictionTimeseriesTest(absltest.TestCase):
 

dgf/src/sampling/in_memory_sampler.py

@@ -106,8 +106,12 @@ def __init__(
     self._slice_timeseries_by_seed = slice_timeseries_by_seed
     self._max_timeseries_len = max_timeseries_len
     self._has_temporal_edgesets = has_temporal_edgesets
+    # The sampled graphs only contain node features (see
+    # `add_features_to_samples`): Edge timeseries features are not sliced.
     self._timeseries_schema_cache = (
-        temporal_util.extract_timeseries_schema_cache(self._schema)
+        temporal_util.extract_timeseries_schema_cache(
+            _drop_edge_features(self._schema)
+        )
     )
 
     if self._max_timeseries_len <= 0:
@@ -494,7 +498,7 @@ def _configure_merging(self):
       # clipping to `max_timeseries_len`, and timeseries padding) in the c++
       # merging, and remove `_python_merger`.
       self._python_merger = merge_lib.GraphMerger(
-          schema=self._output_schema(), padding=self._padding
+          schema=self.output_schema(), padding=self._padding
       )
       return
 
@@ -614,8 +618,23 @@ def _convert_sample_inputs(
         )
     return seed_node_idxs, seed_timestamps, masked_edge_idxs
 
-  def _output_schema(self) -> schema_lib.GraphSchema:
-    """Schema of the sampled graphs given the `return_*` options."""
+  def output_schema(self) -> schema_lib.GraphSchema:
+    """Returns the schema of the sampled graphs.
+
+    The output schema is the sampler schema reduced to what the sampled graphs
+    actually contain: Node features are only present with `return_features`,
+    the "#idx" node feature is only present with `return_node_idxs`, and edge
+    features are never present.
+
+    Usage example:
+    ```python
+    sampler = dgf.sampling.create_sampler(graph, config, schema)
+    output_schema = sampler.output_schema()
+    ```
+
+    Returns:
+      The schema of the graphs returned by the sampling methods.
+    """
     node_sets = {}
     for node_set_name, node_set_schema in self._schema.node_sets.items():
       features = {}
@@ -630,7 +649,9 @@ def _output_schema(self) -> schema_lib.GraphSchema:
       node_sets[node_set_name] = dataclasses.replace(
           node_set_schema, features=features
       )
-    return dataclasses.replace(self._schema, node_sets=node_sets)
+    return dataclasses.replace(
+        _drop_edge_features(self._schema), node_sets=node_sets
+    )
 
   def subgraph(
       self,
@@ -840,6 +861,26 @@ def _add_finalize_graphs(
         )
 
 
+def _drop_edge_features(
+    schema: schema_lib.GraphSchema,
+) -> schema_lib.GraphSchema:
+  """Removes the edge features of a schema.
+
+  The sampled graphs do not contain edge features.
+
+  Args:
+    schema: The schema to reduce.
+
+  Returns:
+    A copy of `schema` without edge features.
+  """
+  edge_sets = {
+      edge_set_name: dataclasses.replace(edge_set_schema, features={})
+      for edge_set_name, edge_set_schema in schema.edge_sets.items()
+  }
+  return dataclasses.replace(schema, edge_sets=edge_sets)
+
+
 def _normalized_random_seeds(
     random_seeds: int | list[int] | np.ndarray | None, num_seeds: int
 ) -> np.ndarray | None:

dgf/src/sampling/in_memory_sampler_test.py

@@ -1369,6 +1369,68 @@ def test_sample_numpy_array_input(self):
     )
     self.assertLen(samples, 2)
 
+  def test_temporal_sampling_with_edge_timeseries_features(self):
+    # Regression test: Edge timeseries features used to raise a KeyError as the
+    # sampler tried to slice them while it does not return edge features.
+    graph, schema = gen_test_graph.generate_temporal_in_memory_graph(
+        include_e2=False, include_edge_timeseries=True
+    )
+    sampler = in_memory_sampler_lib.create_sampler(
+        graph,
+        config_lib.SimpleSamplingConfig(
+            seed_nodeset="n1",
+            num_hops=2,
+            hop_width=2,
+            reverse=False,
+            temporal_sampling=True,
+        ),
+        schema,
+        batch_size=5,
+    )
+    samples = sampler.sample(
+        np.array([0, 1]),
+        seed_timestamps=np.array([30, 40]),
+    )
+    self.assertLen(samples, 2)
+    for sample in samples:
+      self.assertIn("feat", sample.node_sets["n1"].features)
+      self.assertEmpty(sample.edge_sets["e1"].features)
+
+    # The output schema only contains the returned features.
+    self.assertIn("deletion_time", schema.edge_sets["e1"].features)
+    output_schema = sampler.output_schema()
+    self.assertEqual(
+        output_schema.node_sets["n1"].features,
+        schema.node_sets["n1"].features,
+    )
+    self.assertEmpty(output_schema.edge_sets["e1"].features)
+
+  def test_output_schema(self):
+    sampler = in_memory_sampler_lib.create_sampler(
+        self.graph,
+        config_lib.SimpleSamplingConfig(seed_nodeset="n1", num_hops=1),
+        self.schema,
+        batch_size=5,
+        return_features=True,
+        return_node_idxs=True,
+    )
+    output_schema = sampler.output_schema()
+    self.assertEqual(
+        output_schema.node_sets["n1"].features,
+        self.schema.node_sets["n1"].features
+        | {
+            "#idx": schema_lib.FeatureSchema(
+                format=schema_lib.FeatureFormat.INTEGER_64
+            )
+        },
+    )
+    for edge_set_schema in output_schema.edge_sets.values():
+      self.assertEmpty(edge_set_schema.features)
+
+    sampler.set_return_options(return_features=False, return_node_idxs=False)
+    for node_set_schema in sampler.output_schema().node_sets.values():
+      self.assertEmpty(node_set_schema.features)
+
   def test_edgeset_name_to_idx(self):
     sampling_config = config_lib.SimpleSamplingConfig(
         seed_nodeset="n1", num_hops=1

dgf/src/util/gen_test_graph.py

@@ -2192,8 +2192,18 @@ def gen_toy_regression_dataset(
 
 def generate_temporal_in_memory_graph(
     include_e2: bool,
+    include_edge_timeseries: bool = False,
 ) -> tuple[in_memory_graph_lib.InMemoryGraph, schema_lib.GraphSchema]:
-  """Generates a temporal InMemoryGraph for testing."""
+  """Generates a temporal InMemoryGraph for testing.
+
+  Args:
+    include_e2: If true, adds a non-temporal edgeset "e2".
+    include_edge_timeseries: If true, adds a timestamp timeseries feature
+      "deletion_time" (with its own group) to the edgeset "e1".
+
+  Returns:
+    The graph and its schema.
+  """
 
   edge_sets_schema = {
       "e1": schema_lib.EdgeSchema(
@@ -2216,6 +2226,21 @@ def generate_temporal_in_memory_graph(
       )
   }
 
+  if include_edge_timeseries:
+    edge_sets_schema["e1"].features["deletion_time"] = schema_lib.FeatureSchema(
+        format=schema_lib.FeatureFormat.INTEGER_64,
+        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+        shape=(None,),
+        is_timeseries=True,
+        is_creation_time=True,
+        group="deletion_time",
+    )
+    deletion_time = np.empty(3, dtype=object)
+    deletion_time[0] = np.array([50], dtype=np.int64)
+    deletion_time[1] = np.array([], dtype=np.int64)
+    deletion_time[2] = np.array([45, 55], dtype=np.int64)
+    edge_sets_data["e1"].features["deletion_time"] = deletion_time
+
   if include_e2:
     edge_sets_schema["e2"] = schema_lib.EdgeSchema(
         source="n1",