Add past-clipping to sampler.

PiperOrigin-RevId: 962757853

Author: mesimon

dgf/src/api/BUILD

@@ -144,6 +144,7 @@ py_library(
         "//dgf/src/sampling:beam_semi_distributed_sampler_v2",
         "//dgf/src/sampling:config",
         "//dgf/src/sampling:in_memory_sampler",
+        "//dgf/src/sampling:temporal",
         "//dgf/src/sampling/gcp:spanner_graph_sampler",
     ],
 )

dgf/src/learning/ten_lines/dataset_test.py

@@ -276,6 +276,7 @@ def test_per_sample_transformations(self):
             num_hops=1,
             hop_width=2,
             temporal_sampling=True,
+            max_timeseries_len=5,
         ),
         temporal=True,
         timeseries_pad_and_cap=pad_and_cap_config,
@@ -350,6 +351,7 @@ def test_default_no_per_sample_transforms(self):
             num_hops=1,
             hop_width=2,
             temporal_sampling=True,
+            max_timeseries_len=3,
         ),
         temporal=True,
         drop_remainder=False,
@@ -404,6 +406,7 @@ def test_sampler_returns_node_idxs_only_with_transforms_raises(self):
               num_hops=1,
               hop_width=2,
               temporal_sampling=True,
+              max_timeseries_len=5,
           ),
           temporal=True,
           timeseries_pad_and_cap=timeseries_transform.PadAndCapTimeseriesConfig(),
@@ -445,6 +448,7 @@ def test_timedelta_extraction_with_temporal_false(self):
             seed_nodeset="alerts",
             num_hops=1,
             hop_width=2,
+            max_timeseries_len=5,
         ),
         timeseries_pad_and_cap=timeseries_transform.PadAndCapTimeseriesConfig(
             sequence_length=5
@@ -482,6 +486,7 @@ def test_dynamic_set_sampler_returns_node_idxs_only_raises(self):
             num_hops=1,
             hop_width=2,
             temporal_sampling=True,
+            max_timeseries_len=5,
         ),
         temporal=True,
         timeseries_pad_and_cap=timeseries_transform.PadAndCapTimeseriesConfig(),
@@ -517,6 +522,7 @@ def test_timedelta_extraction_without_pad_and_cap_dynamic_ts_raises(self):
               num_hops=1,
               hop_width=2,
               temporal_sampling=True,
+              max_timeseries_len=5,
           ),
           temporal=True,
           timedelta_extraction=timeseries_transform.TimestampFeatureExtractorConfig(),

dgf/src/sampling/BUILD

@@ -31,6 +31,7 @@ py_library(
     srcs = ["temporal.py"],
     deps = [
         "//dgf/src/data:in_memory_graph",
+        "//dgf/src/data:schema",
         "//dgf/src/util:temporal",
         # numpy dep,
     ],

dgf/src/sampling/config.py

@@ -61,6 +61,8 @@ class SimpleSamplingConfig:
       grpah.
     temporal_sampling: If True, temporal sampling is enabled and causal
       timestamps are inferred from the schema.
+    max_timeseries_len: The maximum number of historical causal sequence steps
+      retained for each timeseries feature.
   """
 
   seed_nodeset: str
@@ -69,6 +71,7 @@ class SimpleSamplingConfig:
   reverse: bool = True
   with_replacement: bool = False
   temporal_sampling: bool = False
+  max_timeseries_len: int = 32
 
 
 @dataclasses.dataclass
@@ -112,6 +115,8 @@ class SamplingPlan:
       timestamps are inferred from the schema.
     edgeset_timestamp_features: Mapping from edgeset name to its timestamp
       feature name for causal filtering.
+    max_timeseries_len: The maximum number of historical causal sequence steps
+      retained for each timeseries feature.
   """
 
   root: PlanNode
@@ -120,6 +125,7 @@ class SamplingPlan:
   edgeset_timestamp_features: Dict[str, str] = dataclasses.field(
       default_factory=dict
   )
+  max_timeseries_len: int = 32
 
 
 def simple_sampling_config_to_sampling_plan(
@@ -177,5 +183,6 @@ def rec_build(nodeset: str, depth: int) -> PlanNode:
       with_replacement=src.with_replacement,
       temporal_sampling=src.temporal_sampling,
       edgeset_timestamp_features=edgeset_ts_features,
+      max_timeseries_len=src.max_timeseries_len,
   )
 

dgf/src/sampling/config_test.py

@@ -120,6 +120,29 @@ def test_simple_sampling_config_to_sampling_config(self):
     )
     self.assertEqual(sampling_config, expected_sampling_config)
 
+  def test_simple_sampling_config_to_sampling_plan_with_max_timeseries_len(
+      self,
+  ):
+    schema = schema_lib.GraphSchema(
+        node_sets={
+            "n1": schema_lib.NodeSchema(features={}),
+        },
+        edge_sets={},
+    )
+    simple_config = config_lib.SimpleSamplingConfig(
+        seed_nodeset="n1", num_hops=0, max_timeseries_len=10
+    )
+    plan = config_lib.simple_sampling_config_to_sampling_plan(
+        simple_config, schema
+    )
+    self.assertEqual(plan.max_timeseries_len, 10)
+
+  def test_default_max_timeseries_len(self):
+    simple_config = config_lib.SimpleSamplingConfig(seed_nodeset="n1")
+    self.assertEqual(simple_config.max_timeseries_len, 32)
+    plan = config_lib.SamplingPlan(root=config_lib.PlanNode(nodeset="n1"))
+    self.assertEqual(plan.max_timeseries_len, 32)
+
 
 if __name__ == "__main__":
   absltest.main()

dgf/src/sampling/in_memory_sampler.py

@@ -36,7 +36,7 @@ def __init__(
       return_node_idxs: bool,
       schema: Optional[schema_lib.GraphSchema] = None,
       slice_timeseries_by_seed: bool = True,
-      max_timeseries_len: Optional[int] = None,
+      max_timeseries_len: int = 32,
       has_temporal_edgesets: bool = False,
   ):
     self._cc_sampler = cc_sampler
@@ -49,10 +49,18 @@ def __init__(
     self._has_temporal_edgesets = has_temporal_edgesets
     self._timeseries_schema_cache = (
         temporal_util.extract_timeseries_schema_cache(self._schema)
-        if self._schema is not None and self._slice_timeseries_by_seed
+        if self._schema is not None
         else None
     )
 
+    if self._slice_timeseries_by_seed and self._schema is None:
+      raise ValueError(
+          "schema must be provided when `slice_timeseries_by_seed=True`."
+      )
+
+    if self._max_timeseries_len <= 0:
+      raise ValueError("max_timeseries_len must be positive")
+
   def set_return_options(self, return_features: bool, return_node_idxs: bool):
     """Sets whether to return features and node indices in sampled graphs.
 
@@ -117,6 +125,17 @@ def sample(
     """
 
     # Check and convert the user input into what the c++ sampler expects.
+    if (
+        self._return_features
+        and self._slice_timeseries_by_seed
+        and self._has_timeseries_features()
+    ):
+      assert seed_timestamps is not None, (
+          "`seed_timestamps` must be provided when"
+          " `slice_timeseries_by_seed=True` and the schema contains"
+          " `is_timeseries=True` features."
+      )
+
     return_single_graph = False
     if isinstance(seed_node_idxs, int):
       return_single_graph = True
@@ -277,49 +296,51 @@ def _add_finalize_graphs(
 
     If `_return_features` is True, full feature values are added.
     If `_return_node_idxs` is False, the "#idx" feature is removed.
-    If `seed_timestamps` and `_schema` are available and
-    `_slice_timeseries_by_seed` is True, any `is_timeseries=True` features are
-    causally filtered by the seed node timestamp.
+    If `_return_features` is True and the schema has timeseries features:
+    - If `_slice_timeseries_by_seed` is True, timeseries features are causally
+      filtered by the seed node timestamp and clipped to `max_timeseries_len`.
+    - Otherwise, timeseries features are clipped to `max_timeseries_len`.
 
     Args:
       graphs: A list of `InMemoryGraph` objects to be finalized.
       seed_timestamps: Optional timestamps for causal timeseries filtering.
     """
     add_features_to_samples(
-        self._full_graph, graphs, self._return_features, self._return_node_idxs
+        full_graph=self._full_graph,
+        samples=graphs,
+        return_features=self._return_features,
+        return_node_idxs=self._return_node_idxs,
+        schema=self._schema,
+        seed_timestamps=seed_timestamps,
+        slice_timeseries_by_seed=self._slice_timeseries_by_seed,
+        max_timeseries_len=self._max_timeseries_len,
+        timeseries_schema_cache=self._timeseries_schema_cache,
     )
-    if self._slice_timeseries_by_seed and self._return_features:
-      if seed_timestamps is None and self._has_timeseries_features():
-        raise ValueError(
-            "`seed_timestamps` must be provided when"
-            " `slice_timeseries_by_seed=True` and the schema contains"
-            " `is_timeseries=True` features."
-        )
-      if seed_timestamps is not None:
-        if self._schema is None:
-          raise ValueError(
-              "schema must be provided when `slice_timeseries_by_seed=True` and"
-              " `seed_timestamps` are passed."
-          )
-        for i, sample in enumerate(graphs):
-          sampling_temporal_lib.filter_timeseries_by_timestamp(
-              graph=sample,
-              schema_cache=self._timeseries_schema_cache,  # pyrefly: ignore[bad-argument-type]
-              target_timestamp=int(seed_timestamps[i]),
-              max_timeseries_len=self._max_timeseries_len,
-          )
 
 
 def add_features_to_samples(
     full_graph: in_memory_graph_lib.InMemoryGraph,
     samples: List[in_memory_graph_lib.InMemoryGraph],
     return_features: bool,
     return_node_idxs: bool,
+    schema: Optional[schema_lib.GraphSchema] = None,
+    seed_timestamps: Optional[np.ndarray] = None,
+    slice_timeseries_by_seed: bool = False,
+    max_timeseries_len: int = 32,
+    timeseries_schema_cache: Optional[
+        temporal_util.TimeseriesSchemaCache
+    ] = None,
 ):
   """Adds features and optionally removes temporary node indices from sampled graphs.
 
   If `return_features` is True, full feature values are copied from the
-  `full_graph` to the corresponding nodes in each graph within `graphs`.
+  `full_graph` to the corresponding nodes in each graph within `samples`.
+  If `schema` or `timeseries_schema_cache` is provided and contains timeseries
+  features, extraction is fused with temporal filtering or clipping:
+  - If `slice_timeseries_by_seed` is True, timeseries features are causally
+    sliced up to each sample's `seed_timestamps` and capped to
+    `max_timeseries_len`.
+  - Otherwise, timeseries features are clipped to `max_timeseries_len`.
   If `return_node_idxs` is False, the "#idx" feature, which contains the
   original node indices, is removed from each node set in the sampled graphs.
 
@@ -331,19 +352,59 @@ def add_features_to_samples(
       values from `full_graph`.
     return_node_idxs: Whether to keep the "#idx" feature in the sampled graphs.
       If False, this feature is removed.
+    schema: Optional graph schema.
+    seed_timestamps: Optional timestamps corresponding to each sample in
+      `samples` for causal timeseries filtering.
+    slice_timeseries_by_seed: Whether to causally slice timeseries features by
+      the seed timestamp.
+    max_timeseries_len: Optional cap on the maximum sequence length to retain.
+    timeseries_schema_cache: Optional precomputed `TimeseriesSchemaCache`.
   """
-  if return_features or not return_node_idxs:
-    # Extract feature values.
-    # TODO(gbm): Do this in C++.
-    for sample in samples:
-      for node_set_name, node_set in sample.node_sets.items():
-        node_idxs = node_set.features["#idx"]
-        if not return_node_idxs:
+  if timeseries_schema_cache is None and schema is not None:
+    timeseries_schema_cache = temporal_util.extract_timeseries_schema_cache(
+        schema
+    )
+
+  has_ts = (
+      timeseries_schema_cache is not None
+      and timeseries_schema_cache.has_timeseries
+  )
+
+  for i, sample in enumerate(samples):
+    target_ts = None
+    if has_ts and slice_timeseries_by_seed and seed_timestamps is not None:
+      target_ts = int(seed_timestamps[i])
+
+    if return_features:
+      if has_ts and timeseries_schema_cache is not None:
+        for node_set_name, node_set in sample.node_sets.items():
+          pk_name = timeseries_schema_cache.node_primary_keys.get(node_set_name)
+          if (
+              pk_name is not None
+              and "#idx" in node_set.features
+              and pk_name not in node_set.features
+          ):
+            node_set.features[pk_name] = node_set.features["#idx"]
+        sampling_temporal_lib.extract_features_timeseries(
+            graph=sample,
+            source_graph=full_graph,
+            schema_cache=timeseries_schema_cache,
+            max_timeseries_len=max_timeseries_len,
+            target_timestamp=target_ts,
+        )
+      else:
+        for node_set_name, node_set in sample.node_sets.items():
+          if "#idx" in node_set.features:
+            node_idxs = node_set.features["#idx"]
+            for feature_name, feature_values in full_graph.node_sets[
+                node_set_name
+            ].features.items():
+              node_set.features[feature_name] = feature_values[node_idxs]
+
+    if not return_node_idxs:
+      for node_set in sample.node_sets.values():
+        if "#idx" in node_set.features:
           del node_set.features["#idx"]
-        if return_features:
-          features = full_graph.node_sets[node_set_name].features
-          for feature_name, full_feature_value in features.items():
-            node_set.features[feature_name] = full_feature_value[node_idxs]
 
 
 def create_sampler(
@@ -359,7 +420,6 @@ def create_sampler(
     seed: Optional[int] = None,
     edgeset_to_mask: Optional[str] = None,
     slice_timeseries_by_seed: Optional[bool] = None,
-    max_timeseries_len: Optional[int] = None,
 ) -> Sampler:
   """Creates an in-memory sampler.
 
@@ -387,8 +447,6 @@ def create_sampler(
     slice_timeseries_by_seed: Whether to causally slice `is_timeseries=True`
       sequence features by the seed node timestamp. Defaults to
       `plan.temporal_sampling`.
-    max_timeseries_len: Optional cap on the number of historical causal sequence
-      steps retained for each timeseries feature.
 
   TODO(gbm): Should we remove the compilation variations (e.g., change in random
     number generator, change in hashmaps).
@@ -440,6 +498,6 @@ def create_sampler(
       return_node_idxs=return_node_idxs,
       schema=schema,
       slice_timeseries_by_seed=slice_timeseries_by_seed,
-      max_timeseries_len=max_timeseries_len,
+      max_timeseries_len=plan.max_timeseries_len,
       has_temporal_edgesets=bool(edgeset_timestamp_features),
   )