Add past-clipping to sampler.

dgf/src/api/BUILD

@@ -145,6 +145,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/BUILD

@@ -271,9 +271,10 @@ py_library(
         "//dgf/src/io:jax",
         "//dgf/src/learning/jax:common",
         "//dgf/src/sampling:config",
-        "//dgf/src/sampling:in_memory_sampler",
+        "//dgf/src/sampling:temporal",
         "//dgf/src/transform:normalize",
         "//dgf/src/util:log",
+        "//dgf/src/util:temporal",
         "//dgf/src/util:util_py",
         # jax dep,
         # numpy dep,

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/learning/ten_lines/link_prediction_dataset.py

@@ -803,19 +803,26 @@ def _generate_one(
             "Cached normalized graph is not available in numpy format. If using"
             " device cache, use generate_jax()."
         )
+      sampling_schema = live.sampling_schema or self._sampling_schema()
       pos_src_graph = (
           node_prediction_dataset_lib.attach_features_from_numpy_graph(
-              live.normalized_source_graph, raw.positive_source_graph
+              live.normalized_source_graph,
+              raw.positive_source_graph,
+              sampling_schema,
           )
       )
       pos_trg_graph = (
           node_prediction_dataset_lib.attach_features_from_numpy_graph(
-              live.normalized_target_graph, raw.positive_target_graph
+              live.normalized_target_graph,
+              raw.positive_target_graph,
+              sampling_schema,
           )
       )
       neg_trg_graph = (
           node_prediction_dataset_lib.attach_features_from_numpy_graph(
-              live.normalized_target_graph, raw.negative_target_graph
+              live.normalized_target_graph,
+              raw.negative_target_graph,
+              sampling_schema,
           )
       )
     else:

dgf/src/learning/ten_lines/node_prediction_dataset.py

@@ -33,9 +33,10 @@
 from dgf.src.learning.ten_lines import common
 from dgf.src.learning.ten_lines import dataset
 from dgf.src.sampling import config as sampling_config_lib
-from dgf.src.sampling import in_memory_sampler as in_memory_sampler_lib
+from dgf.src.sampling import temporal as sampling_temporal_lib
 from dgf.src.transform import normalize as normalize_lib
 from dgf.src.util import log
+from dgf.src.util import temporal as temporal_util
 from dgf.src.util import util
 import jax
 import jax.numpy as jnp
@@ -55,6 +56,7 @@ class LiveData:
   sample_generator: dataset.SampleGeneratorFromAnything
   normalized_graph: Optional[in_memory_graph_lib.InMemoryGraph] = None
   normalized_jax_graph: Optional[jax_in_memory_graph.JaxInMemoryGraph] = None
+  timeseries_schema_cache: Optional[temporal_util.TimeseriesSchemaCache] = None
 
 
 @dataclasses.dataclass(kw_only=True)
@@ -235,6 +237,9 @@ def prepare_from_existing_one(self, other: "GNNDatasetPreparator"):
         sampling_plan=sample_generator.sampling_config,  # pyrefly: ignore[bad-argument-type]
         num_nodes_in_seed_nodeset=sample_generator.num_seed_nodes,
         sample_generator=sample_generator,
+        timeseries_schema_cache=temporal_util.extract_timeseries_schema_cache(
+            self.schema
+        ),
     )
 
     if self.cache_normalized_features:
@@ -351,6 +356,9 @@ def gen_normalized_samples():
         sampling_plan=sample_generator.sampling_config,  # pyrefly: ignore[bad-argument-type]
         num_nodes_in_seed_nodeset=sample_generator.num_seed_nodes,
         sample_generator=sample_generator,
+        timeseries_schema_cache=temporal_util.extract_timeseries_schema_cache(
+            self.schema
+        ),
     )
 
     if self.cache_normalized_features:
@@ -392,7 +400,9 @@ def generate(
               "prepared in `prepare()`."
           )
         normalized_sample = attach_features_from_numpy_graph(
-            live.normalized_graph, sample
+            live.normalized_graph,
+            sample,
+            live.timeseries_schema_cache or self.schema,
         )
       else:
         normalized_sample = live.normalizer.normalize_numpy(sample)
@@ -437,11 +447,23 @@ def generate_jax(
 def attach_features_from_numpy_graph(
     graph: in_memory_graph_lib.InMemoryGraph,
     sample: in_memory_graph_lib.InMemoryGraph,
+    schema_or_cache: Union[
+        schema_lib.GraphSchema, temporal_util.TimeseriesSchemaCache
+    ],
 ) -> in_memory_graph_lib.InMemoryGraph:
   """Attaches the numpy features from `graph` to the `sample`."""
-  in_memory_sampler_lib.add_features_to_samples(
-      graph, [sample], return_features=True, return_node_idxs=False
+  if isinstance(schema_or_cache, schema_lib.GraphSchema):
+    cache = temporal_util.extract_timeseries_schema_cache(schema_or_cache)
+  else:
+    cache = schema_or_cache
+  sampling_temporal_lib.extract_features_timeseries(
+      graph=sample,
+      source_graph=graph,
+      timeseries_schema_cache=cache,
   )
+  for node_set in sample.node_sets.values():
+    if "#idx" in node_set.features:
+      del node_set.features["#idx"]
   return sample
 
 

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,
     ],
@@ -183,6 +184,7 @@ py_test(
         "//dgf/src/data:schema",
         "//dgf/src/transform:temporal",
         "//dgf/src/util:gen_test_graph",
+        "//dgf/src/util:temporal",
         "//dgf/src/util:test_util",
         # 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
@@ -70,6 +72,7 @@ class SimpleSamplingConfig:
   with_replacement: bool = False
   temporal_sampling: bool = False
   multi_visit: bool = True
+  max_timeseries_len: int = 32
 
 
 @dataclasses.dataclass
@@ -113,6 +116,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
@@ -122,6 +127,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(
@@ -178,4 +184,5 @@ def rec_build(nodeset: str, depth: int) -> PlanNode:
       temporal_sampling=src.temporal_sampling,
       multi_visit=src.multi_visit,
       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()