Skip to content

Add automatic temporal split. #188

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Sep 11, 2026
+160 −16
Merged

Add automatic temporal split. #188

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 42 additions & 13 deletions dgf/src/learning/ten_lines/node_prediction_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -625,6 +625,8 @@ def compute_train_and_valid_node_idxs(
train_seed_nodes: Optional[common.SeedNodeIdxs],
valid_seed_nodes: Optional[common.SeedNodeIdxs],
max_num_valid_examples: Optional[int],
temporal_split: bool = False,
ts_feature: Optional[str] = None,
) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
"""Computes the training and validation seed node indices."""
if not isinstance(graph, in_memory_graph_lib.InMemoryGraph) or (
Expand Down Expand Up @@ -670,12 +672,39 @@ def compute_train_and_valid_node_idxs(
)
return None, None

train_seed_node_idxs, valid_seed_node_idxs = util.split_train_valid(
num_graph_seed_nodes,
validation_ratio,
random_seed,
max_num_valid_examples=max_num_valid_examples,
)
if temporal_split:
if ts_feature is None:
raise ValueError(
"`ts_feature` must be specified when `temporal_split` is True."
)
if ts_feature not in graph.node_sets[target_nodeset].features:
raise ValueError(
f"Timestamp feature '{ts_feature}' not found in"
f" target nodeset '{target_nodeset}'."
)
log.info(
"Splitting the train and validation sets temporally on the creation"
" timestamp feature %r of nodeset %r: the oldest nodes are used for"
" training and the most recent ones for validation (instead of a"
" random split).",
ts_feature,
target_nodeset,
)
timestamps = graph.node_sets[target_nodeset].features[ts_feature]
train_seed_node_idxs, valid_seed_node_idxs = (
util.split_train_valid_temporal(
creation_times=timestamps,
validation_ratio=validation_ratio,
max_num_valid_examples=max_num_valid_examples,
)
)
else:
train_seed_node_idxs, valid_seed_node_idxs = util.split_train_valid(
num_graph_seed_nodes,
validation_ratio,
random_seed,
max_num_valid_examples=max_num_valid_examples,
)
log.info(
"Num. training seed nodes: %d, Num. validation seed nodes: %d",
len(train_seed_node_idxs),
Expand Down Expand Up @@ -715,6 +744,11 @@ def prepare_datasets(
else:
max_num_valid_examples = num_valid_steps * batch_size

target_ts_feature = temporal_util.creation_time_feature_name(
schema.node_sets[target_nodeset].features
)
target_has_creation_time = target_ts_feature is not None

train_seed_node_idxs, valid_seed_node_idxs = (
compute_train_and_valid_node_idxs(
graph,
Expand All @@ -726,6 +760,8 @@ def prepare_datasets(
train_seed_nodes=train_seed_nodes,
valid_seed_nodes=valid_seed_nodes,
max_num_valid_examples=max_num_valid_examples,
temporal_split=temporal_sampling and target_has_creation_time,
ts_feature=target_ts_feature,
)
)

Expand All @@ -740,13 +776,6 @@ def prepare_datasets(
sampling_plan = sampling_config_lib.simple_sampling_config_to_sampling_plan(
sampling_config, schema
)

target_has_creation_time = (
temporal_util.creation_time_feature_name(
schema.node_sets[target_nodeset].features
)
is not None
)
if auto_normalize_config is None:
auto_normalize_config = normalize_lib.AutoNormalizeConfig(
keep_raw_features=keep_raw_features or set(),
Expand Down
111 changes: 111 additions & 0 deletions dgf/src/learning/ten_lines/node_prediction_dataset_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -429,6 +429,117 @@ class root:
schema=schema_without_ts,
)

def test_compute_train_and_valid_node_idxs(self):
graph, _ = gen_test_graph.generate_temporal_in_memory_graph(False)
# 1. Temporal split
train_idx, valid_idx = (
node_prediction_dataset.compute_train_and_valid_node_idxs(
graph=graph,
valid_graph=None,
graph_format="IN_MEMORY_GRAPH",
target_nodeset="n1",
random_seed=42,
validation_ratio=0.5,
train_seed_nodes=None,
valid_seed_nodes=None,
max_num_valid_examples=None,
temporal_split=True,
ts_feature="timestamp",
)
)
self.assertIsNotNone(train_idx)
self.assertIsNotNone(valid_idx)
np.testing.assert_array_equal(train_idx, np.array([0, 1]))
np.testing.assert_array_equal(valid_idx, np.array([2, 3]))

# 2. Random split (temporal_split=False)
train_idx, valid_idx = (
node_prediction_dataset.compute_train_and_valid_node_idxs(
graph=graph,
valid_graph=None,
graph_format="IN_MEMORY_GRAPH",
target_nodeset="n1",
random_seed=42,
validation_ratio=0.5,
train_seed_nodes=None,
valid_seed_nodes=None,
max_num_valid_examples=None,
temporal_split=False,
)
)
self.assertIsNotNone(train_idx)
self.assertIsNotNone(valid_idx)
self.assertLen(train_idx, 2)
self.assertLen(valid_idx, 2)

# 3. Error when temporal_split=True but ts_feature is None
with self.assertRaises(ValueError):
node_prediction_dataset.compute_train_and_valid_node_idxs(
graph=graph,
valid_graph=None,
graph_format="IN_MEMORY_GRAPH",
target_nodeset="n1",
random_seed=42,
validation_ratio=0.5,
train_seed_nodes=None,
valid_seed_nodes=None,
max_num_valid_examples=None,
temporal_split=True,
ts_feature=None,
)

# 4. Error when temporal_split=True but ts_feature is not found in graph
with self.assertRaises(ValueError):
node_prediction_dataset.compute_train_and_valid_node_idxs(
graph=graph,
valid_graph=None,
graph_format="IN_MEMORY_GRAPH",
target_nodeset="n1",
random_seed=42,
validation_ratio=0.5,
train_seed_nodes=None,
valid_seed_nodes=None,
max_num_valid_examples=None,
temporal_split=True,
ts_feature="non_existent_feature",
)

def test_prepare_datasets_automatic_temporal_split(self):
graph, schema = gen_test_graph.generate_temporal_in_memory_graph(False)
train_dataset, valid_dataset = node_prediction_dataset.prepare_datasets(
graph=graph,
valid_graph=None, # pyrefly: ignore[bad-argument-type]
schema=schema,
target_nodeset="n1",
random_seed=42,
batch_size=2,
num_sampling_hops=1,
sampling_width=3,
verbose=0,
graph_format="IN_MEMORY_GRAPH",
validation_ratio=0.5,
train_seed_nodes=None,
valid_seed_nodes=None,
temporal_sampling=True,
nodeset_timestamp_features={"n1": "timestamp"},
edgeset_timestamp_features={"e1": "timestamp"},
num_valid_steps=None,
cache_valid_dataset=False,
cache_normalized_features=False,
cache_normalized_features_device="host",
sampling_plan=None,
)
self.assertIsNotNone(train_dataset.seed_node_idxs)
self.assertIsNotNone(valid_dataset)
assert valid_dataset is not None
self.assertIsNotNone(valid_dataset.seed_node_idxs)
np.testing.assert_array_equal(
np.sort(train_dataset.seed_node_idxs), np.array([0, 1])
)
np.testing.assert_array_equal(
np.sort(valid_dataset.seed_node_idxs), np.array([2, 3])
)


if __name__ == "__main__":
absltest.main()
10 changes: 7 additions & 3 deletions dgf/src/learning/ten_lines/node_prediction_train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,8 +195,10 @@ def train_node_model(
verbose: The verbosity level. Higher values provide more output.
validation_ratio: Ratio of the training dataset used to create the
validation dataset in case no validation dataset is manually provided
e.g., train_seed_nodes and valid_seed_nodes are provided. If set to 0, the
entire dataset is used for training, and the tree is not pruned.
e.g., train_seed_nodes and valid_seed_nodes are provided. If `time_aware`
is True, the split is temporal (past nodes for training, future nodes for
validation) instead of random. If set to 0, the entire dataset is used
for training, and the tree is not pruned.
train_seed_nodes: Optional. A np.ndarray or list of integer indices
specifying the subset of nodes within the `target_nodeset` to be used for
training. If None, the training nodes are determined based on
Expand Down Expand Up @@ -227,7 +229,9 @@ def train_node_model(
time-consuming, but it will increase memory usage.
time_aware: Enables temporal-aware training. If `False` (default), no
temporal masking is applied. If `True`, timestamp features are inferred
from the schema (via features marked as creation timestamps).
from the schema (via features marked as creation timestamps), the target
nodeset is required to have a creation timestamp, and the train and
validation sets are split temporally instead of randomly.
message_pooling: The pooling method to use for aggregating messages.
experimental_preprocess_core_model_config: Advanced option. An optional
callable to modify the `CoreModelConfig` before it is used to build the
Expand Down