From bfeec499ad830d6468c68d066d0cac9bcd9efc0a Mon Sep 17 00:00:00 2001 From: Simon Meierhans Date: Thu, 3 Sep 2026 03:55:06 -0700 Subject: [PATCH] Set timeseries padding parameter to `max_timeseries_len` PiperOrigin-RevId: 975643936 --- dgf/src/analyse/BUILD | 4 + dgf/src/analyse/padding.py | 32 ++++- dgf/src/analyse/padding_test.py | 123 ++++++++++++++++++ .../ten_lines/node_prediction_dataset.py | 6 +- 4 files changed, 161 insertions(+), 4 deletions(-) diff --git a/dgf/src/analyse/BUILD b/dgf/src/analyse/BUILD index 05158aa..915183b 100644 --- a/dgf/src/analyse/BUILD +++ b/dgf/src/analyse/BUILD @@ -57,6 +57,7 @@ py_library( "//dgf/src/data:in_memory_graph", "//dgf/src/data:padding", "//dgf/src/data:schema", + # numpy dep, ], ) @@ -183,9 +184,12 @@ py_test( deps = [ ":padding", # absl/testing:absltest dep, + "//dgf/src/data:in_memory_graph", "//dgf/src/data:padding", + "//dgf/src/data:schema", "//dgf/src/util:gen_test_graph", "//dgf/src/util:test_util", + # numpy dep, ], ) diff --git a/dgf/src/analyse/padding.py b/dgf/src/analyse/padding.py index 1d3f950..f69031e 100644 --- a/dgf/src/analyse/padding.py +++ b/dgf/src/analyse/padding.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import Iterator, Optional +from typing import Dict, Iterator, Optional from dgf.src.data import in_memory_graph as in_memory_graph_lib from dgf.src.data import padding as padding_lib from dgf.src.data import schema as schema_lib @@ -21,11 +21,27 @@ # TODO(gbm): Implement Beam version. +def _feature_padding_from_schema( + features_schema: schema_lib.FeatureSetSchema, + max_timeseries_len: Optional[int] = None, +) -> Dict[str, padding_lib.FeaturePadding]: + """Creates feature padding for timeseries features from a feature schema.""" + features = {} + if max_timeseries_len is not None: + for feat_name, feat_schema in features_schema.items(): + if feat_schema.is_timeseries: + features[feat_name] = padding_lib.FeaturePadding( + max_timeseries_len=max_timeseries_len + ) + return features + + def padding_from_graph_generator( schema: schema_lib.GraphSchema, graphs: Iterator[in_memory_graph_lib.InMemoryGraph], relative_margin: float = 0.1, absolute_margin: int = 1, + max_timeseries_len: Optional[int] = None, ) -> padding_lib.Padding: """Creates a padding configuration from a set of in-memory graphs. @@ -55,6 +71,8 @@ def padding_from_graph_generator( graphs: An iterator over in-memory heterogeneous graphs. relative_margin: A relative margin. absolute_margin: An absolute margin. + max_timeseries_len: Optional maximum sequence length for timeseries + features. Returns: Padding configuration compatible with the graphs. @@ -86,16 +104,24 @@ def padding_from_graph_generator( padded_node_sets = {} for node_set_name, max_n in max_nodes.items(): padded_n = math.ceil((max_n + absolute_margin) * (1.0 + relative_margin)) + node_set_schema = schema.node_sets[node_set_name] padded_node_sets[node_set_name] = padding_lib.NodeSetPadding( - num_nodes=padded_n + num_nodes=padded_n, + features=_feature_padding_from_schema( + node_set_schema.features, max_timeseries_len + ), ) # Create the padding with a margin. padded_edge_sets = {} for edge_set_name, max_e in max_edges.items(): padded_e = math.ceil((max_e + absolute_margin) * (1.0 + relative_margin)) + edge_set_schema = schema.edge_sets[edge_set_name] padded_edge_sets[edge_set_name] = padding_lib.EdgeSetPadding( - num_edges=padded_e + num_edges=padded_e, + features=_feature_padding_from_schema( + edge_set_schema.features, max_timeseries_len + ), ) return padding_lib.Padding( diff --git a/dgf/src/analyse/padding_test.py b/dgf/src/analyse/padding_test.py index ed688b0..ef92bfe 100644 --- a/dgf/src/analyse/padding_test.py +++ b/dgf/src/analyse/padding_test.py @@ -14,9 +14,12 @@ from absl.testing import absltest from dgf.src.analyse import padding as padding_lib +from dgf.src.data import in_memory_graph as in_memory_graph_lib from dgf.src.data import padding as padding_data_lib +from dgf.src.data import schema as schema_lib from dgf.src.util import gen_test_graph from dgf.src.util import test_util +import numpy as np test_util.disable_diff_truncation() @@ -43,6 +46,110 @@ def test_basic(self): ) test_util.assert_are_equal(self, padding, expected_padding) + def test_timeseries_padding(self): + # 3 features: 1 static timeseries, 2 variable timeseries sharing a group. + schema = schema_lib.GraphSchema( + node_sets={ + "n": schema_lib.NodeSchema( + features={ + "static_ts": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + shape=(5,), + is_timeseries=True, + ), + "var_ts1": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + shape=(None,), + is_timeseries=True, + group="g1", + ), + "var_ts2": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(None,), + is_timeseries=True, + is_creation_time=True, + group="g1", + ), + } + ) + }, + edge_sets={ + "e": schema_lib.EdgeSchema( + source="n", + target="n", + features={ + "edge_ts": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + shape=(None,), + is_timeseries=True, + ), + }, + ) + }, + ) + g1 = in_memory_graph_lib.InMemoryGraph( + node_sets={ + "n": in_memory_graph_lib.InMemoryNodeSet( + num_nodes=2, + features={ + "static_ts": np.zeros((2, 5), dtype=np.float32), + "var_ts1": np.array( + [np.array([1.0, 2.0]), np.array([3.0, 4.0, 5.0])], + dtype=object, + ), + "var_ts2": np.array( + [ + np.array([10, 20], dtype=np.int64), + np.array([20, 30, 40], dtype=np.int64), + ], + dtype=object, + ), + }, + ) + }, + edge_sets={ + "e": in_memory_graph_lib.InMemoryEdgeSet( + adjacency=np.array([[0], [1]], dtype=np.int32), + features={ + "edge_ts": np.array([np.array([1.0, 2.0])], dtype=object), + }, + ) + }, + ) + padding = padding_lib.padding_from_graph_generator( + schema, iter([g1]), max_timeseries_len=4 + ) + expected_padding = padding_data_lib.Padding( + node_sets={ + "n": padding_data_lib.NodeSetPadding( + num_nodes=4, + features={ + "static_ts": padding_data_lib.FeaturePadding( + max_timeseries_len=4 + ), + "var_ts1": padding_data_lib.FeaturePadding( + max_timeseries_len=4 + ), + "var_ts2": padding_data_lib.FeaturePadding( + max_timeseries_len=4 + ), + }, + ) + }, + edge_sets={ + "e": padding_data_lib.EdgeSetPadding( + num_edges=3, + features={ + "edge_ts": padding_data_lib.FeaturePadding( + max_timeseries_len=4 + ), + }, + ) + }, + ) + test_util.assert_are_equal(self, padding, expected_padding) + def test_print_padding(self): padding = padding_data_lib.Padding( node_sets={ @@ -64,6 +171,22 @@ def test_print_padding(self): e1: 100 edges""" self.assertEqual(output, expected_output) + def test_feature_padding_from_schema(self): + features_schema = { + "f": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + is_timeseries=True, + ), + } + feature_padding = padding_lib._feature_padding_from_schema( + features_schema, max_timeseries_len=5 + ) + expected_padding = { + "f": padding_data_lib.FeaturePadding(max_timeseries_len=5), + } + self.assertEqual(feature_padding, expected_padding) + if __name__ == "__main__": absltest.main() + diff --git a/dgf/src/learning/ten_lines/node_prediction_dataset.py b/dgf/src/learning/ten_lines/node_prediction_dataset.py index 13610f5..d947eba 100644 --- a/dgf/src/learning/ten_lines/node_prediction_dataset.py +++ b/dgf/src/learning/ten_lines/node_prediction_dataset.py @@ -340,7 +340,11 @@ def gen_normalized_samples(): if self.verbose_preparation: log.info("Compute graph statistics for padding") padding = padding_lib.padding_from_graph_generator( - self.schema, gen_normalized_samples_iter + self.schema, + gen_normalized_samples_iter, + max_timeseries_len=getattr( + sample_generator.sampling_config, "max_timeseries_len", None + ), ) if self.verbose_preparation: log.info(