diff --git a/dgf/src/api/BUILD b/dgf/src/api/BUILD index 2a92df2..7f11e67 100644 --- a/dgf/src/api/BUILD +++ b/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", ], ) diff --git a/dgf/src/learning/ten_lines/BUILD b/dgf/src/learning/ten_lines/BUILD index dcbca37..045d02f 100644 --- a/dgf/src/learning/ten_lines/BUILD +++ b/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, diff --git a/dgf/src/learning/ten_lines/dataset_test.py b/dgf/src/learning/ten_lines/dataset_test.py index ab59bb5..452dd95 100644 --- a/dgf/src/learning/ten_lines/dataset_test.py +++ b/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(), diff --git a/dgf/src/learning/ten_lines/link_prediction_dataset.py b/dgf/src/learning/ten_lines/link_prediction_dataset.py index 7fe2066..509b03c 100644 --- a/dgf/src/learning/ten_lines/link_prediction_dataset.py +++ b/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: diff --git a/dgf/src/learning/ten_lines/node_prediction_dataset.py b/dgf/src/learning/ten_lines/node_prediction_dataset.py index bc5edaa..5eb0ac6 100644 --- a/dgf/src/learning/ten_lines/node_prediction_dataset.py +++ b/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 diff --git a/dgf/src/sampling/BUILD b/dgf/src/sampling/BUILD index 70b4b93..dfae975 100644 --- a/dgf/src/sampling/BUILD +++ b/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, ], diff --git a/dgf/src/sampling/config.py b/dgf/src/sampling/config.py index e5d75d0..ab1d832 100644 --- a/dgf/src/sampling/config.py +++ b/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, ) diff --git a/dgf/src/sampling/config_test.py b/dgf/src/sampling/config_test.py index f2322bf..a856694 100644 --- a/dgf/src/sampling/config_test.py +++ b/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() diff --git a/dgf/src/sampling/in_memory_sampler.py b/dgf/src/sampling/in_memory_sampler.py index 5c9f72f..ec6a283 100644 --- a/dgf/src/sampling/in_memory_sampler.py +++ b/dgf/src/sampling/in_memory_sampler.py @@ -32,27 +32,28 @@ def __init__( self, cc_sampler, full_graph: in_memory_graph_lib.InMemoryGraph, + schema: schema_lib.GraphSchema, return_features: bool, return_node_idxs: bool, - schema: Optional[schema_lib.GraphSchema] = None, - slice_timeseries_by_seed: bool = True, - max_timeseries_len: Optional[int] = None, - has_temporal_edgesets: bool = False, + slice_timeseries_by_seed: bool, + max_timeseries_len: int, + has_temporal_edgesets: bool, ): self._cc_sampler = cc_sampler self._full_graph = full_graph + self._schema = schema self._return_features = return_features self._return_node_idxs = return_node_idxs - self._schema = schema self._slice_timeseries_by_seed = slice_timeseries_by_seed self._max_timeseries_len = max_timeseries_len 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 - else None ) + 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 +118,19 @@ def sample( """ # Check and convert the user input into what the c++ sampler expects. + assert ( + not ( + self._return_features + and self._slice_timeseries_by_seed + and self._has_timeseries_features() + ) + or 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,9 +291,10 @@ 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. @@ -288,26 +303,25 @@ def _add_finalize_graphs( add_features_to_samples( self._full_graph, graphs, self._return_features, self._return_node_idxs ) - if self._slice_timeseries_by_seed and self._return_features: - if seed_timestamps is None and self._has_timeseries_features(): + if self._has_timeseries_features() and self._return_features: + if self._slice_timeseries_by_seed and seed_timestamps is None: 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, - ) + for i, sample in enumerate(graphs): + target_timestamp = ( + int(seed_timestamps[i]) + if (self._slice_timeseries_by_seed and seed_timestamps is not None) + else None + ) + sampling_temporal_lib.extract_features_timeseries( + graph=sample, + timeseries_schema_cache=self._timeseries_schema_cache, # pyrefly: ignore[bad-argument-type] + target_timestamp=target_timestamp, + max_timeseries_len=self._max_timeseries_len, + ) def add_features_to_samples( @@ -335,18 +349,15 @@ def add_features_to_samples( if return_features or not return_node_idxs: # Extract feature values. # TODO(gbm): Do this in C++. - full_node_sets = full_graph.node_sets 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: del node_set.features["#idx"] if return_features: - features = full_node_sets[node_set_name].features + features = full_graph.node_sets[node_set_name].features for feature_name, full_feature_value in features.items(): - node_set.features[feature_name] = np.take( - full_feature_value, node_idxs, axis=0 - ) + node_set.features[feature_name] = full_feature_value[node_idxs] def create_sampler( @@ -362,7 +373,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. @@ -390,8 +400,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). @@ -443,6 +451,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), ) diff --git a/dgf/src/sampling/in_memory_sampler_test.py b/dgf/src/sampling/in_memory_sampler_test.py index d75c343..0eeda0b 100644 --- a/dgf/src/sampling/in_memory_sampler_test.py +++ b/dgf/src/sampling/in_memory_sampler_test.py @@ -1281,11 +1281,15 @@ def _create_causal_timeseries_test_graph(self): node_sets={ "alerts": in_memory_graph_lib.InMemoryNodeSet( num_nodes=2, - features={"#creation_time": np.array([30, 60], dtype=np.int64)}, + features={ + "#id": np.array([0, 1], dtype=np.int64), + "#creation_time": np.array([30, 60], dtype=np.int64), + }, ), "hardware": in_memory_graph_lib.InMemoryNodeSet( num_nodes=1, features={ + "#id": np.array([0], dtype=np.int64), "time": np.array( [ np.array([10, 20, 30, 40, 50, 60, 70]), @@ -1315,15 +1319,23 @@ def _create_causal_timeseries_test_graph(self): node_sets={ "alerts": schema_lib.NodeSchema( features={ + "#id": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.PRIMARY_ID, + ), "#creation_time": schema_lib.FeatureSchema( format=schema_lib.FeatureFormat.INTEGER_64, semantic=schema_lib.FeatureSemantic.TIMESTAMP, is_creation_time=True, - ) + ), } ), "hardware": schema_lib.NodeSchema( features={ + "#id": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.PRIMARY_ID, + ), "time": schema_lib.FeatureSchema( format=schema_lib.FeatureFormat.INTEGER_64, semantic=schema_lib.FeatureSemantic.TIMESTAMP, @@ -1357,6 +1369,7 @@ def test_causal_timeseries_filtering_in_sample(self): hop_width=10, reverse=True, temporal_sampling=True, + max_timeseries_len=3, ) sampler = in_memory_sampler_lib.create_sampler( graph, @@ -1365,7 +1378,6 @@ def test_causal_timeseries_filtering_in_sample(self): batch_size=2, return_features=True, slice_timeseries_by_seed=True, - max_timeseries_len=3, ) samples = sampler.sample( [0, 1], seed_timestamps=np.array([30, 60], dtype=np.int64) @@ -1407,8 +1419,8 @@ def test_causal_timeseries_filtering_in_sample(self): temporal_sampling=True, slice_timeseries_by_seed=False, expected_slice=False, - expected_time=[10, 20, 30, 40, 50, 60, 70], - expected_signal=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], + expected_time=[50, 60, 70], + expected_signal=[5.0, 6.0, 7.0], ), dict( testcase_name="override_plan_to_true", @@ -1434,6 +1446,7 @@ def test_slice_timeseries_by_seed( hop_width=10, reverse=True, temporal_sampling=temporal_sampling, + max_timeseries_len=3, ) kwargs = {} if slice_timeseries_by_seed is not None: @@ -1445,9 +1458,9 @@ def test_slice_timeseries_by_seed( schema, batch_size=2, return_features=True, - max_timeseries_len=3, **kwargs, ) + self.assertEqual(sampler._slice_timeseries_by_seed, expected_slice) if expected_time is not None: samples = sampler.sample( @@ -1464,32 +1477,12 @@ def test_slice_timeseries_by_seed( def test_slice_timeseries_by_seed_validation_errors(self): empty_graph = in_memory_graph_lib.InMemoryGraph(node_sets={}, edge_sets={}) - # 1. Missing schema when slice_timeseries_by_seed=True and seed_timestamps - # passed - sampler_no_schema = in_memory_sampler_lib.Sampler( - cc_sampler=None, - full_graph=empty_graph, - return_features=True, - return_node_idxs=False, - schema=None, - slice_timeseries_by_seed=True, - ) - with self.assertRaisesRegex( - ValueError, - "schema must be provided when `slice_timeseries_by_seed=True` and" - " `seed_timestamps` are passed", - ): - sampler_no_schema._add_finalize_graphs( - [empty_graph], seed_timestamps=np.array([123], dtype=np.int64) - ) - - # 2. Missing seed_timestamps when schema contains is_timeseries=True - # featuresh + # 2. Invalid max_timeseries_len (<= 0) ts_schema = schema_lib.GraphSchema( node_sets={ "n1": schema_lib.NodeSchema( features={ - "time": schema_lib.FeatureSchema( + "f1": schema_lib.FeatureSchema( format=schema_lib.FeatureFormat.INTEGER_64, semantic=schema_lib.FeatureSemantic.TIMESTAMP, is_timeseries=True, @@ -1499,24 +1492,149 @@ def test_slice_timeseries_by_seed_validation_errors(self): }, edge_sets={}, ) - sampler_with_schema = in_memory_sampler_lib.Sampler( + with self.assertRaisesRegex( + ValueError, + "max_timeseries_len must be positive", + ): + in_memory_sampler_lib.Sampler( + cc_sampler=None, + full_graph=empty_graph, + return_features=True, + return_node_idxs=False, + schema=ts_schema, + slice_timeseries_by_seed=True, + max_timeseries_len=0, + has_temporal_edgesets=False, + ) + + # 3. Missing seed_timestamps when slice_timeseries_by_seed=True and schema + # contains is_timeseries=True features + sampler_missing_timestamps = in_memory_sampler_lib.Sampler( cc_sampler=None, full_graph=empty_graph, return_features=True, return_node_idxs=False, schema=ts_schema, slice_timeseries_by_seed=True, + max_timeseries_len=3, + has_temporal_edgesets=False, ) with self.assertRaisesRegex( - ValueError, + AssertionError, "`seed_timestamps` must be provided when" " `slice_timeseries_by_seed=True` and the schema contains" " `is_timeseries=True` features.", ): - sampler_with_schema._add_finalize_graphs( - [empty_graph], seed_timestamps=None + sampler_missing_timestamps.sample( + seed_node_idxs=0, seed_timestamps=None ) + def test_timeseries_subgraph_and_multisubgraph_clipping(self): + graph, schema = self._create_causal_timeseries_test_graph() + plan = config_lib.SimpleSamplingConfig( + seed_nodeset="alerts", + num_hops=1, + hop_width=10, + max_timeseries_len=3, + ) + sampler = in_memory_sampler_lib.create_sampler( + graph, + plan, + schema, + batch_size=2, + return_features=True, + ) + + # Test subgraph + sub = sampler.subgraph([0]) + # Original hardware: [10, 20, 30, 40, 50, 60, 70] -> last 3: [50, 60, 70] + np.testing.assert_array_equal( + sub.node_sets["hardware"].features["time"][0], [50, 60, 70] + ) + np.testing.assert_array_equal( + sub.node_sets["hardware"].features["signal"][0], [5.0, 6.0, 7.0] + ) + + # Test multisubgraph + multisubs = sampler.multisubgraph([0, 1]) + self.assertLen(multisubs, 2) + np.testing.assert_array_equal( + multisubs[0].node_sets["hardware"].features["time"][0], [50, 60, 70] + ) + np.testing.assert_array_equal( + multisubs[1].node_sets["hardware"].features["time"][0], [50, 60, 70] + ) + + def test_timeseries_unseeded_sample_clipping(self): + graph, schema = self._create_causal_timeseries_test_graph() + plan = config_lib.SimpleSamplingConfig( + seed_nodeset="alerts", + num_hops=1, + hop_width=10, + temporal_sampling=False, + max_timeseries_len=4, + ) + sampler = in_memory_sampler_lib.create_sampler( + graph, + plan, + schema, + batch_size=1, + return_features=True, + debug_sampling=True, + ) + sample = sampler.sample(0) + # Original hardware: 7 items -> last 4: [40, 50, 60, 70] + np.testing.assert_array_equal( + sample.node_sets["hardware"].features["time"][0], [40, 50, 60, 70] + ) + np.testing.assert_array_equal( + sample.node_sets["hardware"].features["signal"][0], + [4.0, 5.0, 6.0, 7.0], + ) + + def test_add_features_to_samples_fused_direct(self): + graph, schema = self._create_causal_timeseries_test_graph() + # Create sample graph containing only #idx + sample = in_memory_graph_lib.InMemoryGraph( + node_sets={ + "hardware": in_memory_graph_lib.InMemoryNodeSet( + num_nodes=1, + features={"#idx": np.array([0], dtype=np.int64)}, + ), + "alerts": in_memory_graph_lib.InMemoryNodeSet( + num_nodes=1, + features={"#idx": np.array([0], dtype=np.int64)}, + ), + }, + edge_sets={ + "hardware_to_alert": in_memory_graph_lib.InMemoryEdgeSet( + adjacency=np.empty((2, 0), dtype=np.int64), + features={}, + ), + }, + ) + sampler = in_memory_sampler_lib.Sampler( + cc_sampler=None, + full_graph=graph, + schema=schema, + return_features=True, + return_node_idxs=False, + slice_timeseries_by_seed=True, + max_timeseries_len=2, + has_temporal_edgesets=False, + ) + sampler._add_finalize_graphs( + [sample], seed_timestamps=np.array([30], dtype=np.int64) + ) + self.assertNotIn("#idx", sample.node_sets["hardware"].features) + # Original times <= 30: [10, 20, 30], max_len 2 -> [20, 30] + np.testing.assert_array_equal( + sample.node_sets["hardware"].features["time"][0], [20, 30] + ) + np.testing.assert_array_equal( + sample.node_sets["hardware"].features["signal"][0], [2.0, 3.0] + ) + class InMemorySamplerDiamon(parameterized.TestCase): diff --git a/dgf/src/sampling/temporal.py b/dgf/src/sampling/temporal.py index 9e3a4ae..46e40af 100644 --- a/dgf/src/sampling/temporal.py +++ b/dgf/src/sampling/temporal.py @@ -14,117 +14,192 @@ """Temporal sampling utilities for filtering and slicing timeseries features.""" -from typing import List, Optional, Union +from typing import List, Optional, Tuple, Union from dgf.src.data import in_memory_graph from dgf.src.util import temporal as temporal_util import numpy as np +# TODO(simonmeierhans): Improve performance ofhandling of static shaped +# timeseries. -def _filter_entity_set_timeseries( - entity_val: Union[ + +def _compute_group_slices( + timestamp_values: np.ndarray, + node_idxs: np.ndarray, + max_timeseries_len: int, + target_timestamp: Optional[int] = None, +) -> Tuple[np.ndarray, np.ndarray]: + """Computes (start_indices, end_indices) for a sequence group on selected nodes. + + Args: + timestamp_values: Timestamp sequence feature array indexed by `node_idxs` + (array of integer timestamps or object array of 1D arrays). + node_idxs: 1D array of selected node/entity indices. + max_timeseries_len: Positive integer cap on sequence steps to retain. + target_timestamp: Optional causal cutoff timestamp (`int`). If None, slices + up to the end of each sequence. + + Returns: + A tuple `(start_indices, end_indices)` of 1D int64 arrays. + """ + if max_timeseries_len <= 0: + raise ValueError("max_timeseries_len must be positive") + + num_entities = len(node_idxs) + if num_entities == 0: + return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64) + + start_indices = np.empty(num_entities, dtype=np.int64) + end_indices = np.empty(num_entities, dtype=np.int64) + + for i, idx in enumerate(node_idxs): + times = timestamp_values[idx] + if target_timestamp is not None: + end_idx = int(np.searchsorted(times, target_timestamp, side="right")) + else: + end_idx = len(times) + start_idx = max(0, end_idx - max_timeseries_len) + start_indices[i] = start_idx + end_indices[i] = end_idx + + return start_indices, end_indices + + +def _crop_timeseries( + feature_values: np.ndarray, + node_idxs: np.ndarray, + start_indices: np.ndarray, + end_indices: np.ndarray, +) -> np.ndarray: + """Extracts and crops timeseries features directly using start and end indices.""" + num_entities = len(node_idxs) + if num_entities == 0: + return np.empty(0, dtype=object) + + target_arr = np.empty(num_entities, dtype=object) + for i, idx in enumerate(node_idxs): + target_arr[i] = feature_values[idx][start_indices[i] : end_indices[i]] + return target_arr + + +def _process_entity_set_timeseries( + target_val: Union[ in_memory_graph.InMemoryNodeSet, in_memory_graph.InMemoryEdgeSet ], + source_val: Union[ + in_memory_graph.InMemoryNodeSet, in_memory_graph.InMemoryEdgeSet + ], + node_idxs: np.ndarray, ts_specs: List[temporal_util.TimeseriesGroupSpec], - num_entities: int, - target_timestamp: int, - max_timeseries_len: Optional[int], + max_timeseries_len: int, + target_timestamp: Optional[int] = None, ) -> None: - """Causally slices pre-scanned timeseries features for a node set or edge set in place.""" + """Causally filters and/or clips timeseries features in place or from source into target.""" + timeseries_features = set() for group in ts_specs: - # The features is a timeseries, but does not have associated timestamps - # this can be the case for regularly sampled features, but prevents - # filtering to a specific timestamp to prevent lookahead. - if group.timestamp_feature_name is None: - continue - - ts_val = entity_val.features[group.timestamp_feature_name] - feature_names = group.feature_names - feature_arrays = [entity_val.features[fname] for fname in feature_names] - sliced_target = { - fname: np.empty(num_entities, dtype=np.object_) - for fname in feature_names - } - target_lists = [sliced_target[fname] for fname in feature_names] - - # TODO(mesimon): Move into C++ for performance. - for idx in range(num_entities): - times = ts_val[idx] - end_idx = np.searchsorted(times, target_timestamp, side="right") - start_idx = ( - max(0, end_idx - max_timeseries_len) - if max_timeseries_len is not None - else 0 + if group.timestamp_feature_name is not None: + ts_val = source_val.features[group.timestamp_feature_name] + start_indices, end_indices = _compute_group_slices( + timestamp_values=ts_val, + node_idxs=node_idxs, + max_timeseries_len=max_timeseries_len, + target_timestamp=target_timestamp, ) - slc = slice(start_idx, end_idx) - - for feat_arr, target_arr in zip(feature_arrays, target_lists): - target_arr[idx] = feat_arr[idx][slc] - - for fname, arr in sliced_target.items(): - entity_val.features[fname] = arr - - -def filter_timeseries_by_timestamp( + for fname in group.feature_names: + timeseries_features.add(fname) + target_val.features[fname] = _crop_timeseries( + feature_values=source_val.features[fname], + node_idxs=node_idxs, + start_indices=start_indices, + end_indices=end_indices, + ) + else: + for fname in group.feature_names: + timeseries_features.add(fname) + val = source_val.features[fname] + if val.dtype == object: + target_val.features[fname] = np.array( + [elem[-max_timeseries_len:] for elem in val[node_idxs]], + dtype=object, + ) + else: + target_val.features[fname] = val[node_idxs, -max_timeseries_len:] + + for fname, full_val in source_val.features.items(): + if fname not in timeseries_features: + target_val.features[fname] = full_val[node_idxs] + + +def extract_features_timeseries( graph: in_memory_graph.InMemoryGraph, - schema_cache: temporal_util.TimeseriesSchemaCache, - target_timestamp: Union[int, np.integer], - max_timeseries_len: Optional[int] = None, + timeseries_schema_cache: temporal_util.TimeseriesSchemaCache, + max_timeseries_len: int = 32, + target_timestamp: Optional[int] = None, + source_graph: Optional[in_memory_graph.InMemoryGraph] = None, ) -> None: - """In-place filters and caps `is_timeseries=True` features by a causal cutoff. + """In-place filters and/or extracts `is_timeseries=True` features for `graph`. - For each node or edge with `is_timeseries=True` features (such as `time` and - `signal`), this function identifies valid causal sequence indices where the - timestamp is less than or equal to `target_timestamp`. Any future sequence - entries (`time > target_timestamp`) are stripped in place from all timeseries - feature arrays. Optionally caps the sequence to the most recent - `max_timeseries_len` observations. + - If `source_graph` is provided: extracts features directly from + `source_graph` for each node set based on the `"#idx"` row indices in a + single pass. `graph.node_sets` must contain `"#idx"` with 0-based node row + indices into `source_graph.node_sets`. + - If `source_graph` is None: modifies `graph` in place. Note: This function modifies `graph` in place. - Usage example: - - ```python - cache = dgf.util.temporal.extract_timeseries_schema_cache(schema) - dgf.sampling.temporal.filter_timeseries_by_timestamp( - graph=subgraph, - schema_cache=cache, - target_timestamp=1680000000, - max_timeseries_len=30, - ) - ``` - Args: - graph: The input in-memory graph (or sampled subgraph) to be modified in - place. - schema_cache: A pre-computed `TimeseriesSchemaCache` derived from the graph - schema via `extract_timeseries_schema_cache`. - target_timestamp: The causal cutoff timestamp (`int`). - max_timeseries_len: Optional integer cap on the number of causal sequence - steps to retain. If a node or edge has more than `max_timeseries_len` - causal steps, only the most recent `max_timeseries_len` steps are kept. + graph: The in-memory graph to be modified in place. When `source_graph` is + provided, each node set in `graph` must contain `"#idx"`. + timeseries_schema_cache: A pre-computed `TimeseriesSchemaCache`. + max_timeseries_len: Positive integer cap on sequence steps to retain. + target_timestamp: Optional causal cutoff timestamp (`int`). + source_graph: Optional source graph to extract features from. """ - - # Process Node Sets - for ns_name, ts_specs in schema_cache.node_sets.items(): - ns_val = graph.node_sets.get(ns_name) - if ns_val and ns_val.num_nodes: - _filter_entity_set_timeseries( - ns_val, - ts_specs, - ns_val.num_nodes, - target_timestamp=target_timestamp, # pyrefly: ignore[bad-argument-type] - max_timeseries_len=max_timeseries_len, - ) - - # Process Edge Sets - for es_name, ts_specs in schema_cache.edge_sets.items(): - es_val = graph.edge_sets.get(es_name) - if es_val and es_val.num_edges(): - _filter_entity_set_timeseries( - es_val, - ts_specs, - es_val.num_edges(), - target_timestamp=target_timestamp, # pyrefly: ignore[bad-argument-type] + if max_timeseries_len <= 0: + raise ValueError("max_timeseries_len must be positive") + + if source_graph is not None: + for ns_name, target_ns in graph.node_sets.items(): + source_ns = source_graph.node_sets[ns_name] + if "#idx" not in target_ns.features: + raise ValueError( + f"NodeSet '{ns_name}' in sampled graph is missing required '#idx' " + "feature mapping back to source_graph." + ) + node_idxs = target_ns.features["#idx"] + ts_specs = timeseries_schema_cache.node_sets.get(ns_name, []) + _process_entity_set_timeseries( + target_val=target_ns, + source_val=source_ns, + node_idxs=node_idxs, + ts_specs=ts_specs, max_timeseries_len=max_timeseries_len, + target_timestamp=target_timestamp, ) + else: + # Process Node Sets + for ns_name, ts_specs in timeseries_schema_cache.node_sets.items(): + ns_val = graph.node_sets[ns_name] + if ns_val.num_nodes: + _process_entity_set_timeseries( + target_val=ns_val, + source_val=ns_val, + node_idxs=np.arange(ns_val.num_nodes), + ts_specs=ts_specs, + max_timeseries_len=max_timeseries_len, + target_timestamp=target_timestamp, + ) + + # Process Edge Sets + for es_name, ts_specs in timeseries_schema_cache.edge_sets.items(): + es_val = graph.edge_sets[es_name] + if es_val.num_edges(): + _process_entity_set_timeseries( + target_val=es_val, + source_val=es_val, + node_idxs=np.arange(es_val.num_edges()), + ts_specs=ts_specs, + max_timeseries_len=max_timeseries_len, + target_timestamp=target_timestamp, + ) diff --git a/dgf/src/sampling/temporal_test.py b/dgf/src/sampling/temporal_test.py index f3b4b0d..bc3a8e0 100644 --- a/dgf/src/sampling/temporal_test.py +++ b/dgf/src/sampling/temporal_test.py @@ -117,9 +117,9 @@ def test_filter_timeseries_by_timestamp_scalar(self): edge_sets={}, ) cache = temporal_util.extract_timeseries_schema_cache(schema) - temporal.filter_timeseries_by_timestamp( + temporal.extract_features_timeseries( graph=graph, - schema_cache=cache, + timeseries_schema_cache=cache, target_timestamp=28, max_timeseries_len=2, ) @@ -179,8 +179,11 @@ def test_filter_timeseries_skips_non_timestamp_series(self): edge_sets={}, ) cache = temporal_util.extract_timeseries_schema_cache(schema) - temporal.filter_timeseries_by_timestamp( - graph=graph, schema_cache=cache, target_timestamp=25 + temporal.extract_features_timeseries( + graph=graph, + timeseries_schema_cache=cache, + target_timestamp=25, + max_timeseries_len=32, ) hw_set = graph.node_sets["hardware"] np.testing.assert_array_equal(hw_set.features["time"][0], [10, 20]) @@ -242,16 +245,17 @@ def test_filter_timeseries_edge_set(self): }, ) cache = temporal_util.extract_timeseries_schema_cache(schema) - temporal.filter_timeseries_by_timestamp( + temporal.extract_features_timeseries( graph=graph, - schema_cache=cache, + timeseries_schema_cache=cache, target_timestamp=20, + max_timeseries_len=32, ) e1_set = graph.edge_sets["e1"] np.testing.assert_array_equal(e1_set.features["timestamps"][0], [10, 20]) np.testing.assert_array_equal(e1_set.features["edge_sig"][0], [1.0, 2.0]) - def test_filter_timeseries_skips_missing_node_set(self): + def test_filter_timeseries_missing_node_set_raises(self): graph = in_memory_graph.InMemoryGraph(node_sets={}, edge_sets={}) missing_cache = temporal_util.TimeseriesSchemaCache( node_sets={ @@ -264,12 +268,34 @@ def test_filter_timeseries_skips_missing_node_set(self): edge_sets={}, has_timeseries=True, ) - # Should not raise ValueError for missing node set - temporal.filter_timeseries_by_timestamp( - graph=graph, - schema_cache=missing_cache, - target_timestamp=25, + with self.assertRaises(KeyError): + temporal.extract_features_timeseries( + graph=graph, + timeseries_schema_cache=missing_cache, + target_timestamp=25, + max_timeseries_len=32, + ) + + def test_filter_timeseries_missing_edge_set_raises(self): + graph = in_memory_graph.InMemoryGraph(node_sets={}, edge_sets={}) + missing_cache = temporal_util.TimeseriesSchemaCache( + node_sets={}, + edge_sets={ + "missing_edge": [ + temporal_util.TimeseriesGroupSpec( + timestamp_feature_name="t", feature_names=["t"] + ) + ] + }, + has_timeseries=True, ) + with self.assertRaises(KeyError): + temporal.extract_features_timeseries( + graph=graph, + timeseries_schema_cache=missing_cache, + target_timestamp=25, + max_timeseries_len=32, + ) def test_extract_timeseries_schema_cache_and_filter(self): graph, schema = _make_sample_graph_and_schema() @@ -281,8 +307,11 @@ def test_extract_timeseries_schema_cache_and_filter(self): self.assertEqual(group.timestamp_feature_name, "time") self.assertCountEqual(group.feature_names, ["time", "signal"]) - temporal.filter_timeseries_by_timestamp( - graph=graph, schema_cache=cache, target_timestamp=30 + temporal.extract_features_timeseries( + graph=graph, + timeseries_schema_cache=cache, + target_timestamp=30, + max_timeseries_len=32, ) np.testing.assert_array_equal( graph.node_sets["hardware"].features["time"][0], @@ -293,6 +322,326 @@ def test_extract_timeseries_schema_cache_and_filter(self): [1.0, 2.0, 3.0], ) + def test_filter_timeseries_clips_non_timestamp_series(self): + graph = in_memory_graph.InMemoryGraph( + node_sets={ + "n1": in_memory_graph.InMemoryNodeSet( + num_nodes=1, + features={ + "f1": np.array( + [np.array([10, 20, 30, 40])], dtype=np.object_ + ), + }, + ) + }, + edge_sets={}, + ) + schema = schema_lib.GraphSchema( + node_sets={ + "n1": schema_lib.NodeSchema( + features={ + "f1": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + is_timeseries=True, + ), + } + ) + }, + edge_sets={}, + ) + cache = temporal_util.extract_timeseries_schema_cache(schema) + temporal.extract_features_timeseries( + graph=graph, + timeseries_schema_cache=cache, + target_timestamp=25, + max_timeseries_len=2, + ) + np.testing.assert_array_equal( + graph.node_sets["n1"].features["f1"][0], [30, 40] + ) + + def test_clip_timeseries_to_max_len(self): + graph = in_memory_graph.InMemoryGraph( + node_sets={ + "n1": in_memory_graph.InMemoryNodeSet( + num_nodes=2, + features={ + "f1": np.array( + [ + np.array([10, 20, 30, 40, 50]), + np.array([1, 2, 3]), + ], + dtype=np.object_, + ), + "f2": np.array( + [ + [1.0, 2.0, 3.0, 4.0, 5.0], + [6.0, 7.0, 8.0, 9.0, 10.0], + ], + dtype=np.float32, + ), + }, + ) + }, + edge_sets={ + "e1": in_memory_graph.InMemoryEdgeSet( + adjacency=np.array([[0], [1]], dtype=np.int64), + features={ + "f1": np.array( + [np.array([100, 200, 300, 400])], dtype=np.object_ + ), + }, + ) + }, + ) + schema = schema_lib.GraphSchema( + node_sets={ + "n1": schema_lib.NodeSchema( + features={ + "f1": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + is_timeseries=True, + ), + "f2": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + is_timeseries=True, + shape=(5,), + ), + } + ) + }, + edge_sets={ + "e1": schema_lib.EdgeSchema( + source="n1", + target="n1", + features={ + "f1": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + is_timeseries=True, + ) + }, + ) + }, + ) + cache = temporal_util.extract_timeseries_schema_cache(schema) + temporal.extract_features_timeseries( + graph=graph, + timeseries_schema_cache=cache, + max_timeseries_len=3, + target_timestamp=None, + ) + + n1 = graph.node_sets["n1"] + np.testing.assert_array_equal(n1.features["f1"][0], [30, 40, 50]) + np.testing.assert_array_equal(n1.features["f1"][1], [1, 2, 3]) + np.testing.assert_array_equal( + n1.features["f2"], + np.array([[3.0, 4.0, 5.0], [8.0, 9.0, 10.0]], dtype=np.float32), + ) + np.testing.assert_array_equal( + graph.edge_sets["e1"].features["f1"][0], [200, 300, 400] + ) + + def test_clip_timeseries_empty_entities(self): + graph = in_memory_graph.InMemoryGraph( + node_sets={ + "n1": in_memory_graph.InMemoryNodeSet( + num_nodes=0, + features={"f1": np.array([], dtype=np.object_)}, + ) + }, + edge_sets={}, + ) + schema = schema_lib.GraphSchema( + node_sets={ + "n1": schema_lib.NodeSchema( + features={ + "f1": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + is_timeseries=True, + ), + } + ) + }, + edge_sets={}, + ) + cache = temporal_util.extract_timeseries_schema_cache(schema) + temporal.extract_features_timeseries( + graph=graph, + timeseries_schema_cache=cache, + max_timeseries_len=3, + target_timestamp=None, + ) + + def test_invalid_max_timeseries_len_raises(self): + graph, schema = _make_sample_graph_and_schema() + cache = temporal_util.extract_timeseries_schema_cache(schema) + with self.assertRaisesRegex( + ValueError, "max_timeseries_len must be positive" + ): + temporal.extract_features_timeseries( + graph=graph, timeseries_schema_cache=cache, max_timeseries_len=0 + ) + with self.assertRaisesRegex( + ValueError, "max_timeseries_len must be positive" + ): + temporal._compute_group_slices( + timestamp_values=np.array([10, 20]), + node_idxs=np.array([0]), + max_timeseries_len=0, + ) + + def test_compute_group_slices_causal_object_array(self): + ts_values = np.array( + [ + np.array([10, 20, 30, 40, 50]), + np.array([5, 15, 25, 35]), + ], + dtype=np.object_, + ) + node_idxs = np.array([0, 1]) + starts, ends = temporal._compute_group_slices( + timestamp_values=ts_values, + node_idxs=node_idxs, + target_timestamp=28, + max_timeseries_len=2, + ) + np.testing.assert_array_equal(starts, np.array([0, 1])) + np.testing.assert_array_equal(ends, np.array([2, 3])) + + def test_compute_group_slices_causal_dense_array(self): + ts_values = np.array( + [[10, 20, 30, 40], [5, 15, 25, 35]], + dtype=np.int64, + ) + node_idxs = np.array([1, 0]) + starts, ends = temporal._compute_group_slices( + timestamp_values=ts_values, + node_idxs=node_idxs, + max_timeseries_len=4, + target_timestamp=20, + ) + np.testing.assert_array_equal(starts, np.array([0, 0])) + np.testing.assert_array_equal(ends, np.array([2, 2])) + + def test_compute_group_slices_no_target_timestamp(self): + ts_values = np.array( + [np.array([10, 20, 30, 40, 50]), np.array([5, 15, 25])], + dtype=np.object_, + ) + starts, ends = temporal._compute_group_slices( + timestamp_values=ts_values, + node_idxs=np.array([0, 1]), + max_timeseries_len=2, + target_timestamp=None, + ) + np.testing.assert_array_equal(starts, np.array([3, 1])) + np.testing.assert_array_equal(ends, np.array([5, 3])) + + def test_compute_group_slices_empty_node_idxs(self): + ts_values = np.array([np.array([10, 20])], dtype=np.object_) + empty_idxs = np.empty(0, dtype=np.int64) + starts, ends = temporal._compute_group_slices( + timestamp_values=ts_values, + node_idxs=empty_idxs, + max_timeseries_len=5, + target_timestamp=20, + ) + self.assertEmpty(starts) + self.assertEmpty(ends) + + def test_crop_timeseries_per_entity_slice_plan(self): + ts_values = np.array( + [ + np.array([10, 20, 30, 40, 50]), + np.array([5, 15, 25, 35]), + ], + dtype=np.object_, + ) + values = np.array( + [ + np.array([1.0, 2.0, 3.0, 4.0, 5.0]), + np.array([0.5, 1.5, 2.5, 3.5]), + ], + dtype=np.object_, + ) + starts, ends = temporal._compute_group_slices( + timestamp_values=ts_values, + node_idxs=np.array([0, 1]), + target_timestamp=28, + max_timeseries_len=2, + ) + extracted = temporal._crop_timeseries( + feature_values=values, + node_idxs=np.array([0, 1]), + start_indices=starts, + end_indices=ends, + ) + self.assertEqual(extracted.dtype, np.object_) + np.testing.assert_array_equal(extracted[0], [1.0, 2.0]) + np.testing.assert_array_equal(extracted[1], [1.5, 2.5]) + + def test_extract_features_timeseries_from_source_graph(self): + schema = schema_lib.GraphSchema( + node_sets={ + "n1": schema_lib.NodeSchema( + features={ + "node_id": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.PRIMARY_ID, + ), + "f1": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + is_timeseries=True, + ), + } + ) + }, + edge_sets={}, + ) + source_graph = in_memory_graph.InMemoryGraph( + node_sets={ + "n1": in_memory_graph.InMemoryNodeSet( + num_nodes=3, + features={ + "node_id": np.array([100, 101, 102], dtype=np.int64), + "f1": np.array( + [ + np.array([1, 2, 3, 4, 5]), + np.array([10, 20, 30]), + np.array([100, 200, 300, 400]), + ], + dtype=np.object_, + ), + }, + ) + }, + edge_sets={}, + ) + sampled_graph = in_memory_graph.InMemoryGraph( + node_sets={ + "n1": in_memory_graph.InMemoryNodeSet( + num_nodes=2, + features={"#idx": np.array([2, 0], dtype=np.int64)}, + ) + }, + edge_sets={}, + ) + cache = temporal_util.extract_timeseries_schema_cache(schema) + temporal.extract_features_timeseries( + graph=sampled_graph, + source_graph=source_graph, + timeseries_schema_cache=cache, + max_timeseries_len=2, + target_timestamp=None, + ) + + n1 = sampled_graph.node_sets["n1"] + np.testing.assert_array_equal(n1.features["node_id"], [102, 100]) + np.testing.assert_array_equal(n1.features["f1"][0], [300, 400]) + np.testing.assert_array_equal(n1.features["f1"][1], [4, 5]) + if __name__ == "__main__": absltest.main() + diff --git a/dgf/src/util/temporal.py b/dgf/src/util/temporal.py index 2446744..82f6412 100644 --- a/dgf/src/util/temporal.py +++ b/dgf/src/util/temporal.py @@ -132,23 +132,34 @@ def _extract_entity_set_timeseries_specs( features: Dict[str, schema_lib.FeatureSchema], ) -> List[TimeseriesGroupSpec]: """Extracts and groups timeseries feature specs for a node set or edge set.""" - ts_groups: Dict[Optional[str], List[str]] = collections.defaultdict(list) + grouped_features: Dict[str, List[str]] = collections.defaultdict(list) + ungrouped_features: List[str] = [] + for fname, fschema in features.items(): if not fschema.is_timeseries: continue grp = fschema.group or (fname if fschema.is_creation_time else None) - ts_groups[grp].append(fname) + if grp is not None: + grouped_features[grp].append(fname) + else: + ungrouped_features.append(fname) specs = [] - for grp_name, fnames in ts_groups.items(): - ts_feat_name = None - if grp_name is not None: - ts_feat_name = group_creation_time_feature_name(grp_name, features) + for grp_name, fnames in grouped_features.items(): + ts_feat_name = group_creation_time_feature_name(grp_name, features) specs.append( TimeseriesGroupSpec( timestamp_feature_name=ts_feat_name, feature_names=fnames ) ) + + if ungrouped_features: + specs.append( + TimeseriesGroupSpec( + timestamp_feature_name=None, feature_names=ungrouped_features + ) + ) + return specs diff --git a/dgf/src/util/temporal_test.py b/dgf/src/util/temporal_test.py index f0d3486..051d7c2 100644 --- a/dgf/src/util/temporal_test.py +++ b/dgf/src/util/temporal_test.py @@ -62,7 +62,6 @@ def test_extract_timeseries_schema_cache_no_timeseries_features(self): cache = temporal.extract_timeseries_schema_cache(schema) self.assertEqual(cache.node_sets, {"nodes": []}) self.assertEqual(cache.edge_sets, {"edges": []}) - self.assertFalse(cache.has_timeseries) def test_extract_timeseries_schema_cache_grouped(self): schema = schema_lib.GraphSchema( @@ -350,7 +349,9 @@ def test_schema_has_timeseries_features(self): }, edge_sets={}, ) - self.assertTrue(temporal.schema_has_timeseries_features(schema_with_node_ts)) + self.assertTrue( + temporal.schema_has_timeseries_features(schema_with_node_ts) + ) schema_with_edge_ts = schema_lib.GraphSchema( node_sets={"n1": schema_lib.NodeSchema(features={})}, @@ -367,7 +368,9 @@ def test_schema_has_timeseries_features(self): ) }, ) - self.assertTrue(temporal.schema_has_timeseries_features(schema_with_edge_ts)) + self.assertTrue( + temporal.schema_has_timeseries_features(schema_with_edge_ts) + ) def test_schema_has_dynamic_timeseries_features(self): schema_static_ts = schema_lib.GraphSchema(