From 4705d0159f8d393d89ef08410835ff6a430a7745 Mon Sep 17 00:00:00 2001 From: pc0618 Date: Thu, 22 Jan 2026 00:10:48 +0000 Subject: [PATCH 1/8] Add TGB (Temporal Graph Benchmark) datasets + tasks --- relbench/datasets/__init__.py | 28 +++++ relbench/datasets/tgb.py | 93 +++++++++++++++ relbench/tasks/__init__.py | 196 +++++++++++++++++++++++++++++++ relbench/tasks/tgb.py | 209 ++++++++++++++++++++++++++++++++++ 4 files changed, 526 insertions(+) create mode 100644 relbench/datasets/tgb.py create mode 100644 relbench/tasks/tgb.py diff --git a/relbench/datasets/__init__.py b/relbench/datasets/__init__.py index baf87f36..52256b9f 100644 --- a/relbench/datasets/__init__.py +++ b/relbench/datasets/__init__.py @@ -18,6 +18,7 @@ ratebeer, salt, stack, + tgb, trial, ) from relbench.utils import get_relbench_cache_dir @@ -76,6 +77,12 @@ def download_dataset(name: str) -> None: "locally; skipping download." ) return + if name.startswith("rel-tgb-"): + print( + f"Dataset '{name}' is a community TGB export and must be prepared " + "locally (place db/*.parquet in the cache dir); skipping download." + ) + return if name == "rel-mimic": print("Downloading Mimic dataset...") @@ -157,3 +164,24 @@ def get_dataset(name: str, download=True) -> Dataset: register_dataset("dbinfer-amazon", dbinfer.DBInferAmazonDataset) register_dataset("dbinfer-stackexchange", dbinfer.DBInferStackExchangeDataset) register_dataset("dbinfer-outbrain-small", dbinfer.DBInferOutbrainSmallDataset) + +# Community dataset family: Temporal Graph Benchmark (TGB) +# Names follow the convention `rel--` where +# dataset_name = "tgb" and qualifier is the original TGB dataset id. +register_dataset("rel-tgb-tgbl-wiki", tgb.TGBDataset, tgb_name="tgbl-wiki") +register_dataset("rel-tgb-tgbl-wiki-v2", tgb.TGBDataset, tgb_name="tgbl-wiki-v2") +register_dataset("rel-tgb-tgbl-review", tgb.TGBDataset, tgb_name="tgbl-review") +register_dataset("rel-tgb-tgbl-review-v2", tgb.TGBDataset, tgb_name="tgbl-review-v2") +register_dataset("rel-tgb-tgbl-coin", tgb.TGBDataset, tgb_name="tgbl-coin") +register_dataset("rel-tgb-tgbl-comment", tgb.TGBDataset, tgb_name="tgbl-comment") +register_dataset("rel-tgb-tgbl-flight", tgb.TGBDataset, tgb_name="tgbl-flight") + +register_dataset("rel-tgb-thgl-software", tgb.TGBDataset, tgb_name="thgl-software") +register_dataset("rel-tgb-thgl-forum", tgb.TGBDataset, tgb_name="thgl-forum") +register_dataset("rel-tgb-thgl-github", tgb.TGBDataset, tgb_name="thgl-github") +register_dataset("rel-tgb-thgl-myket", tgb.TGBDataset, tgb_name="thgl-myket") + +register_dataset("rel-tgb-tgbn-trade", tgb.TGBDataset, tgb_name="tgbn-trade") +register_dataset("rel-tgb-tgbn-genre", tgb.TGBDataset, tgb_name="tgbn-genre") +register_dataset("rel-tgb-tgbn-reddit", tgb.TGBDataset, tgb_name="tgbn-reddit") +register_dataset("rel-tgb-tgbn-token", tgb.TGBDataset, tgb_name="tgbn-token") diff --git a/relbench/datasets/tgb.py b/relbench/datasets/tgb.py new file mode 100644 index 00000000..b62c95aa --- /dev/null +++ b/relbench/datasets/tgb.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import pandas as pd + +from relbench.base import Database, Dataset + + +@dataclass(frozen=True) +class TGBCutoffs: + val_timestamp_s: int + test_timestamp_s: int + + +_TGB_CUTOFFS: dict[str, TGBCutoffs] = { + # Dynamic link property prediction (tgbl-*) + # NOTE: These cutoffs match the pre-built RelBench exports used in our + # conversion pipeline. They are expected to be consistent with the + # official TGB temporal split (70/15/15). + "tgbl-wiki": TGBCutoffs(val_timestamp_s=1862653, test_timestamp_s=2218300), + "tgbl-wiki-v2": TGBCutoffs(val_timestamp_s=1862653, test_timestamp_s=2218300), + "tgbl-review": TGBCutoffs(val_timestamp_s=1464912000, test_timestamp_s=1488844800), + "tgbl-review-v2": TGBCutoffs(val_timestamp_s=1464912000, test_timestamp_s=1488844800), + "tgbl-coin": TGBCutoffs(val_timestamp_s=1662096249, test_timestamp_s=1664482319), + "tgbl-comment": TGBCutoffs(val_timestamp_s=1282869285, test_timestamp_s=1288838725), + "tgbl-flight": TGBCutoffs(val_timestamp_s=1638162000, test_timestamp_s=1653796800), + # Temporal heterogeneous link prediction (thgl-*) + "thgl-software": TGBCutoffs(val_timestamp_s=1706003880, test_timestamp_s=1706315669), + "thgl-forum": TGBCutoffs(val_timestamp_s=1390426563, test_timestamp_s=1390838358), + "thgl-github": TGBCutoffs(val_timestamp_s=1711075987, test_timestamp_s=1711482874), + "thgl-myket": TGBCutoffs(val_timestamp_s=1603724860, test_timestamp_s=1606341312), + # Dynamic node property prediction (tgbn-*) + # Nodeprop timestamps may be stored as years in the raw sources; our export + # normalizes timestamps to UNIX seconds. + "tgbn-trade": TGBCutoffs(val_timestamp_s=1262304000, test_timestamp_s=1388534400), + "tgbn-genre": TGBCutoffs(val_timestamp_s=1216427762, test_timestamp_s=1230448684), + "tgbn-reddit": TGBCutoffs(val_timestamp_s=1279485233, test_timestamp_s=1286653871), + "tgbn-token": TGBCutoffs(val_timestamp_s=1522889022, test_timestamp_s=1525386888), +} + + +class TGBDataset(Dataset): + r"""Community dataset family: Temporal Graph Benchmark (TGB) exports. + + This dataset class expects a pre-built RelBench database at: + `cache_dir/db/*.parquet`. + + The recommended workflow for contributions (see CONTRIBUTING.md): + 1) Materialize `db/` (Parquet) for each dataset. + 2) Zip it as `db.zip` and publish it. + 3) Add the sha256 hash to `relbench/datasets/hashes.json`. + + Notes: + - Primary/foreign keys are stored as int64 (DBML: `bigint`) to avoid 32-bit + overflow and to match PyArrow's default integer type. + - We store only the relational schema + timestamps; "splits" are derived + from cutoffs (val/test timestamps). + """ + + url = "https://tgb.complexdatalab.com/" + + def __init__(self, *, tgb_name: str, cache_dir: Optional[str] = None) -> None: + if tgb_name not in _TGB_CUTOFFS: + raise ValueError( + f"Unknown tgb_name='{tgb_name}'. Known keys: {sorted(_TGB_CUTOFFS.keys())}" + ) + self.tgb_name = str(tgb_name) + + cutoffs = _TGB_CUTOFFS[self.tgb_name] + # TGB exports store timestamps as timezone-aware UTC (timestamp[ns, UTC]). + # Keep dataset cutoffs in UTC as well to avoid tz-mismatch in task builders. + self.val_timestamp = pd.to_datetime(int(cutoffs.val_timestamp_s), unit="s", utc=True) + self.test_timestamp = pd.to_datetime(int(cutoffs.test_timestamp_s), unit="s", utc=True) + + super().__init__(cache_dir=cache_dir) + + def make_db(self) -> Database: + if self.cache_dir is None: + raise RuntimeError("TGBDataset requires cache_dir to locate the cached db.") + + db_dir = Path(self.cache_dir) / "db" + if db_dir.exists() and any(db_dir.glob("*.parquet")): + return Database.load(db_dir) + + raise RuntimeError( + f"TGB dataset '{self.tgb_name}' not found at {db_dir}. " + "This dataset is distributed as a pre-built RelBench database (db.zip). " + "Please run `download_dataset(...)` or place `db/*.parquet` in the cache directory." + ) + diff --git a/relbench/tasks/__init__.py b/relbench/tasks/__init__.py index 84694dfc..546bf4fe 100644 --- a/relbench/tasks/__init__.py +++ b/relbench/tasks/__init__.py @@ -5,6 +5,7 @@ from functools import lru_cache from typing import List +import pandas as pd import pooch from relbench.base import AutoCompleteTask, BaseTask, TaskType @@ -20,6 +21,7 @@ mimic, ratebeer, stack, + tgb, trial, ) from relbench.utils import get_relbench_cache_dir @@ -81,6 +83,12 @@ def download_task(dataset_name: str, task_name: str) -> None: "generated locally; skipping download." ) return + if dataset_name.startswith("rel-tgb-"): + print( + f"Task '{dataset_name}/{task_name}' is a community TGB task and must be " + "generated locally; skipping download." + ) + return DOWNLOAD_REGISTRY.fetch( f"{dataset_name}/tasks/{task_name}.zip", @@ -547,3 +555,191 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: register_task("dbinfer-stackexchange", "churn", dbinfer.StackExchangeChurnTask) register_task("dbinfer-stackexchange", "upvote", dbinfer.StackExchangeUpvoteTask) register_task("dbinfer-outbrain-small", "ctr", dbinfer.OutbrainCTRTask) + +# Community task family: Temporal Graph Benchmark (TGB) +# +# Task naming follows: +# - recommendation tasks: -- +# We use "next" as the single-word suffix for all TGB link-style tasks. + +# tgbl-* (link prediction) +register_task( + "rel-tgb-tgbl-wiki", + "src-dst-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="src_nodes", dst_entity_table="dst_nodes"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) +register_task( + "rel-tgb-tgbl-wiki-v2", + "src-dst-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="src_nodes", dst_entity_table="dst_nodes"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) +register_task( + "rel-tgb-tgbl-review", + "src-dst-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), + timedelta=pd.Timedelta(days=30), + eval_k=10, +) +register_task( + "rel-tgb-tgbl-review-v2", + "src-dst-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), + timedelta=pd.Timedelta(days=30), + eval_k=10, +) +register_task( + "rel-tgb-tgbl-coin", + "src-dst-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), + timedelta=pd.Timedelta(days=7), + eval_k=10, +) +register_task( + "rel-tgb-tgbl-comment", + "src-dst-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), + timedelta=pd.Timedelta(days=7), + eval_k=10, +) +register_task( + "rel-tgb-tgbl-flight", + "src-dst-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), + timedelta=pd.Timedelta(days=30), + eval_k=10, +) + +# thgl-* (heterogeneous link prediction) – expose type-pair tasks using +# node-type table names from the exports. +register_task( + "rel-tgb-thgl-software", + "type0-type1-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_1"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) +register_task( + "rel-tgb-thgl-software", + "type0-type3-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_3"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) +register_task( + "rel-tgb-thgl-software", + "type1-type2-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_1", dst_entity_table="nodes_type_2"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) +register_task( + "rel-tgb-thgl-software", + "type3-type2-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_3", dst_entity_table="nodes_type_2"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) + +register_task( + "rel-tgb-thgl-github", + "type0-type1-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_1"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) +register_task( + "rel-tgb-thgl-github", + "type0-type3-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_3"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) +register_task( + "rel-tgb-thgl-github", + "type1-type2-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_1", dst_entity_table="nodes_type_2"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) +register_task( + "rel-tgb-thgl-github", + "type3-type2-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_3", dst_entity_table="nodes_type_2"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) + +register_task( + "rel-tgb-thgl-forum", + "type0-type0-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_0"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) +register_task( + "rel-tgb-thgl-forum", + "type0-type1-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_1"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) + +register_task( + "rel-tgb-thgl-myket", + "type0-type1-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_1"), + timedelta=pd.Timedelta(days=7), + eval_k=10, +) + +# tgbn-* (node property prediction as node→label recommendation) +register_task( + "rel-tgb-tgbn-trade", + "node-label-next", + tgb.TGBNodeLabelNextTask, + timedelta=pd.Timedelta(days=365), + eval_k=10, +) +register_task( + "rel-tgb-tgbn-genre", + "node-label-next", + tgb.TGBNodeLabelNextTask, + timedelta=pd.Timedelta(days=7), + eval_k=10, +) +register_task( + "rel-tgb-tgbn-reddit", + "node-label-next", + tgb.TGBNodeLabelNextTask, + timedelta=pd.Timedelta(days=7), + eval_k=10, +) +register_task( + "rel-tgb-tgbn-token", + "node-label-next", + tgb.TGBNodeLabelNextTask, + timedelta=pd.Timedelta(days=7), + eval_k=10, +) diff --git a/relbench/tasks/tgb.py b/relbench/tasks/tgb.py new file mode 100644 index 00000000..0f82943b --- /dev/null +++ b/relbench/tasks/tgb.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import pandas as pd + +from relbench.base import Database, RecommendationTask, Table, TaskType +from relbench.metrics import ( + link_prediction_map, + link_prediction_ndcg, + link_prediction_precision, + link_prediction_recall, +) + + +def _unique_list(values: pd.Series) -> list[int]: + # Preserve deterministic ordering per group as first-appearance order. + # Pandas' `unique()` preserves order. + return [int(x) for x in values.unique().tolist()] + + +@dataclass(frozen=True) +class TGBAggregatedEventsSpec: + src_entity_table: str + dst_entity_table: str + event_tables: tuple[str, ...] = () + event_time_col: str = "event_ts" + src_col: str = "src_id" + dst_col: str = "dst_id" + + +class TGBSrcDstNextTask(RecommendationTask): + r"""Recommendation task for TGB-style edge streams. + + For each anchor timestamp `t`, predict the set of destination entities that + connect to a source entity within the horizon `(t, t + timedelta]`. + + This task is intentionally generic to support: + - homogeneous `tgbl-*` exports: `nodes` + `events` + - bipartite `tgbl-wiki*` exports: `src_nodes`/`dst_nodes` + `events` + - heterogeneous `thgl-*` exports: multiple `events_edge_type_*` tables + aggregated by source/destination node types + """ + + task_type = TaskType.LINK_PREDICTION + metrics = [ + link_prediction_recall, + link_prediction_precision, + link_prediction_map, + link_prediction_ndcg, + ] + + # Convention: keep these stable across all TGB tasks. + src_entity_col = "src_id" + dst_entity_col = "dst_id" + time_col = "timestamp" + + def __init__( + self, + dataset, + *, + spec: TGBAggregatedEventsSpec, + timedelta: pd.Timedelta, + eval_k: int = 10, + cache_dir: Optional[str] = None, + ) -> None: + self.spec = spec + self.timedelta = pd.Timedelta(timedelta) + self.eval_k = int(eval_k) + self.src_entity_table = spec.src_entity_table + self.dst_entity_table = spec.dst_entity_table + super().__init__(dataset, cache_dir=cache_dir) + + def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Table: + event_tables = ( + self.spec.event_tables + if self.spec.event_tables + else tuple(name for name in db.table_dict.keys() if name.startswith("events")) + ) + rows: list[pd.DataFrame] = [] + for t in timestamps: + start = pd.Timestamp(t) + end = start + self.timedelta + + window_frames: list[pd.DataFrame] = [] + for table_name in event_tables: + table = db.table_dict[table_name] + # For heterogeneous exports, keep only event tables whose FK schema + # matches the intended src/dst entity tables. + if table.fkey_col_to_pkey_table.get(self.spec.src_col) != self.src_entity_table: + continue + if table.fkey_col_to_pkey_table.get(self.spec.dst_col) != self.dst_entity_table: + continue + + ev = table.df + mask = (ev[self.spec.event_time_col] > start) & (ev[self.spec.event_time_col] <= end) + if mask.any(): + window_frames.append(ev.loc[mask, [self.spec.src_col, self.spec.dst_col]]) + + if not window_frames: + continue + + win = pd.concat(window_frames, axis=0, ignore_index=True) + grouped = ( + win.groupby(self.spec.src_col, sort=False)[self.spec.dst_col] + .agg(_unique_list) + .reset_index() + .rename(columns={self.spec.src_col: self.src_entity_col, self.spec.dst_col: self.dst_entity_col}) + ) + grouped[self.time_col] = start + rows.append(grouped[[self.time_col, self.src_entity_col, self.dst_entity_col]]) + + if rows: + out = pd.concat(rows, axis=0, ignore_index=True) + else: + out = pd.DataFrame({self.time_col: [], self.src_entity_col: [], self.dst_entity_col: []}) + + return Table( + df=out, + fkey_col_to_pkey_table={ + self.src_entity_col: self.src_entity_table, + self.dst_entity_col: self.dst_entity_table, + }, + pkey_col=None, + time_col=self.time_col, + ) + + +class TGBNodeLabelNextTask(RecommendationTask): + r"""Node property prediction as a recommendation task. + + For each anchor timestamp `t`, predict the set of labels associated with a + node within the horizon `(t, t + timedelta]`. + + This task assumes the standard `tgbn-*` export schema: + - `nodes(node_id)` + - `labels(label_id)` + - `label_events(label_event_id, src_id, label_ts)` + - `label_event_items(item_id, label_event_id, label_id, label_weight)` + + Note: This task treats labels as unweighted positives for evaluation + (binary relevance at top-k). If graded relevance is desired, a custom task + evaluator would be required. + """ + + task_type = TaskType.LINK_PREDICTION + metrics = [ + link_prediction_recall, + link_prediction_precision, + link_prediction_map, + link_prediction_ndcg, + ] + + src_entity_col = "src_id" + dst_entity_col = "label_id" + time_col = "timestamp" + + def __init__( + self, + dataset, + *, + timedelta: pd.Timedelta, + eval_k: int = 10, + cache_dir: Optional[str] = None, + ) -> None: + self.timedelta = pd.Timedelta(timedelta) + self.eval_k = int(eval_k) + self.src_entity_table = "nodes" + self.dst_entity_table = "labels" + super().__init__(dataset, cache_dir=cache_dir) + + def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Table: + le = db.table_dict["label_events"].df + items = db.table_dict["label_event_items"].df + + rows: list[pd.DataFrame] = [] + for t in timestamps: + start = pd.Timestamp(t) + end = start + self.timedelta + mask = (le["label_ts"] > start) & (le["label_ts"] <= end) + if not mask.any(): + continue + + le_win = le.loc[mask, ["label_event_id", "src_id"]] + joined = le_win.merge(items[["label_event_id", "label_id"]], on="label_event_id", how="inner") + grouped = ( + joined.groupby("src_id", sort=False)["label_id"] + .agg(_unique_list) + .reset_index() + .rename(columns={"label_id": self.dst_entity_col, "src_id": self.src_entity_col}) + ) + grouped[self.time_col] = start + rows.append(grouped[[self.time_col, self.src_entity_col, self.dst_entity_col]]) + + if rows: + out = pd.concat(rows, axis=0, ignore_index=True) + else: + out = pd.DataFrame({self.time_col: [], self.src_entity_col: [], self.dst_entity_col: []}) + + return Table( + df=out, + fkey_col_to_pkey_table={ + self.src_entity_col: self.src_entity_table, + self.dst_entity_col: self.dst_entity_table, + }, + pkey_col=None, + time_col=self.time_col, + ) From c51b389252ed4a587a2c8a86b0c5ea79b053ec2b Mon Sep 17 00:00:00 2001 From: pc0618 Date: Thu, 22 Jan 2026 00:41:05 +0000 Subject: [PATCH 2/8] Fix TGB task defs and add hashes --- relbench/datasets/__init__.py | 6 ----- relbench/datasets/hashes.json | 20 +++++++++++++++- relbench/tasks/__init__.py | 44 +++++++++++++++++++++-------------- relbench/tasks/hashes.json | 29 ++++++++++++++++++++++- relbench/tasks/tgb.py | 23 ++++++++++-------- 5 files changed, 87 insertions(+), 35 deletions(-) diff --git a/relbench/datasets/__init__.py b/relbench/datasets/__init__.py index 52256b9f..f8c142cb 100644 --- a/relbench/datasets/__init__.py +++ b/relbench/datasets/__init__.py @@ -77,12 +77,6 @@ def download_dataset(name: str) -> None: "locally; skipping download." ) return - if name.startswith("rel-tgb-"): - print( - f"Dataset '{name}' is a community TGB export and must be prepared " - "locally (place db/*.parquet in the cache dir); skipping download." - ) - return if name == "rel-mimic": print("Downloading Mimic dataset...") diff --git a/relbench/datasets/hashes.json b/relbench/datasets/hashes.json index e2dee7c0..6f6d3c0c 100644 --- a/relbench/datasets/hashes.json +++ b/relbench/datasets/hashes.json @@ -16,5 +16,23 @@ "dbinfer-outbrain-small/db.zip": "d186a71fc534bcac4299c616110805a997fa56fd16c15110d02e1c8eb3975210", "dbinfer-retailrocket/db.zip": "6dbe83488e11c0d159b4592aa4cff57f6fa22b9f4bbf4ae38d0388d536897d75", "dbinfer-seznam/db.zip": "77314bf874dc495e8a4a61f2dc5f12982bbec3c5b6b7af5555e9a2bb587154d9", - "dbinfer-stackexchange/db.zip": "4c8b8ac38b56d57bc2ead3c12be1237b69cae76062ed3f176fc02ad327a84d19" + "dbinfer-stackexchange/db.zip": "4c8b8ac38b56d57bc2ead3c12be1237b69cae76062ed3f176fc02ad327a84d19", + + "rel-tgb-tgbl-coin/db.zip": "058917b0b5ac29cadf0c12c07516ff6cc14c4a5eea0a027902425c9e05bf2690", + "rel-tgb-tgbl-comment/db.zip": "3d2ec40a8fb32969615db31fb8fa81a45d4b9a91a8d40e5f8ee7358ea1c8402a", + "rel-tgb-tgbl-flight/db.zip": "eabc26da95efcad0dd2661eb683470a962355bf7229e8490df653cd990dd3627", + "rel-tgb-tgbl-review/db.zip": "274a21fe7aac039b032f75aeba3efbcaf8af5531ea2c106a61b15651440d694e", + "rel-tgb-tgbl-review-v2/db.zip": "a5719eb25505d3500c910c022a68ff886ad9e6b1687035d0f95263f1f2874deb", + "rel-tgb-tgbl-wiki/db.zip": "f32ad5382de5deccb55ecd465e969fae3d8e49168de0731c7ff7ae8f04754dbd", + "rel-tgb-tgbl-wiki-v2/db.zip": "6d620c982a7209e109d052ca964d5c8f8cbae8575f883c65d9ce377343a6a597", + + "rel-tgb-tgbn-genre/db.zip": "e46aecc28315ca9872117817ce65bc1f4d00ed261ab6d35038a84bcf0ebab7bf", + "rel-tgb-tgbn-reddit/db.zip": "bc444f5cbaf7004ef7c6c8f98ae12ae72e8427ccc9710e96e46b0fc156cc952e", + "rel-tgb-tgbn-token/db.zip": "44a0f90a62642b8054ef1fbee35c862a67e5d0cec92f08b8fd44c503e0986be8", + "rel-tgb-tgbn-trade/db.zip": "8a983ae6281ea058dc8d84f6ce339d2dacf0ae9bb98bb82cca3fed54ec3fe370", + + "rel-tgb-thgl-forum/db.zip": "ef3185fd1558711c51306524d07dddb9356d3108eb0af28727ac75c1f73d3c3f", + "rel-tgb-thgl-github/db.zip": "88105bcfc3d1409f5408c3108a6a300c0782bf85264a1e7153bcc01828368078", + "rel-tgb-thgl-myket/db.zip": "4dce093d60c3009b0c85ac0e93da7a6df9b9a1904c6dafd8cf730d0878e0a5aa", + "rel-tgb-thgl-software/db.zip": "9fe3316bf11908065f1eda28f14ef3cbab64e5166bf0ebf7e359315c0528b99d" } diff --git a/relbench/tasks/__init__.py b/relbench/tasks/__init__.py index 546bf4fe..98c3fa1e 100644 --- a/relbench/tasks/__init__.py +++ b/relbench/tasks/__init__.py @@ -83,12 +83,6 @@ def download_task(dataset_name: str, task_name: str) -> None: "generated locally; skipping download." ) return - if dataset_name.startswith("rel-tgb-"): - print( - f"Task '{dataset_name}/{task_name}' is a community TGB task and must be " - "generated locally; skipping download." - ) - return DOWNLOAD_REGISTRY.fetch( f"{dataset_name}/tasks/{task_name}.zip", @@ -584,7 +578,7 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: "src-dst-next", tgb.TGBSrcDstNextTask, spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), - timedelta=pd.Timedelta(days=30), + timedelta=pd.Timedelta(days=180), eval_k=10, ) register_task( @@ -592,7 +586,7 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: "src-dst-next", tgb.TGBSrcDstNextTask, spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), - timedelta=pd.Timedelta(days=30), + timedelta=pd.Timedelta(days=180), eval_k=10, ) register_task( @@ -600,7 +594,7 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: "src-dst-next", tgb.TGBSrcDstNextTask, spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), - timedelta=pd.Timedelta(days=7), + timedelta=pd.Timedelta(days=14), eval_k=10, ) register_task( @@ -608,7 +602,7 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: "src-dst-next", tgb.TGBSrcDstNextTask, spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), - timedelta=pd.Timedelta(days=7), + timedelta=pd.Timedelta(days=56), eval_k=10, ) register_task( @@ -616,7 +610,7 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: "src-dst-next", tgb.TGBSrcDstNextTask, spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), - timedelta=pd.Timedelta(days=30), + timedelta=pd.Timedelta(days=90), eval_k=10, ) @@ -665,25 +659,41 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: ) register_task( "rel-tgb-thgl-github", - "type0-type3-next", + "type1-type1-next", tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_3"), + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_1", dst_entity_table="nodes_type_1"), timedelta=pd.Timedelta(days=1), eval_k=10, ) register_task( "rel-tgb-thgl-github", - "type1-type2-next", + "type2-type0-next", tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_1", dst_entity_table="nodes_type_2"), + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_2", dst_entity_table="nodes_type_0"), timedelta=pd.Timedelta(days=1), eval_k=10, ) register_task( "rel-tgb-thgl-github", - "type3-type2-next", + "type2-type1-next", tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_3", dst_entity_table="nodes_type_2"), + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_2", dst_entity_table="nodes_type_1"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) +register_task( + "rel-tgb-thgl-github", + "type2-type3-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_2", dst_entity_table="nodes_type_3"), + timedelta=pd.Timedelta(days=1), + eval_k=10, +) +register_task( + "rel-tgb-thgl-github", + "type3-type1-next", + tgb.TGBSrcDstNextTask, + spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_3", dst_entity_table="nodes_type_1"), timedelta=pd.Timedelta(days=1), eval_k=10, ) diff --git a/relbench/tasks/hashes.json b/relbench/tasks/hashes.json index f62d6e8a..820506f5 100644 --- a/relbench/tasks/hashes.json +++ b/relbench/tasks/hashes.json @@ -80,5 +80,32 @@ "dbinfer-seznam/tasks/charge.zip": "6e62750bcd39fe5fdd4d550325953f117517ce153cd5d84a2c5e7ded2caa1b60", "dbinfer-seznam/tasks/prepay.zip": "377564195dd8f8603859771e952fb707a489a5247abf274a5f02ca1bfa3f441d", "dbinfer-stackexchange/tasks/churn.zip": "d649b0a70dc8be9fca797c82e14d19ed4311637711144f4c5b74249081b107ca", - "dbinfer-stackexchange/tasks/upvote.zip": "af8f516422e3e19c0d76f13c00c4c0577cdc218b84468006eda369de8f9c7953" + "dbinfer-stackexchange/tasks/upvote.zip": "af8f516422e3e19c0d76f13c00c4c0577cdc218b84468006eda369de8f9c7953", + + "rel-tgb-tgbl-coin/tasks/src-dst-next.zip": "3e9f11e204bc0cc61f9ba07c5f32abc78e1de4f482804652b307aa231171301a", + "rel-tgb-tgbl-comment/tasks/src-dst-next.zip": "5cb7f475382937666237471112f225cc3aa542c5842e32b948dcb8f5002aba20", + "rel-tgb-tgbl-flight/tasks/src-dst-next.zip": "7e37fd976eb8a468b06fdcd0be6ba987dfca117cb06f7dcc18f6d080ab39ee21", + "rel-tgb-tgbl-review/tasks/src-dst-next.zip": "3c87d7c3fa7cc5f60c59882f56cf49a37a342bfc9766bd067d57ee8e637578e5", + "rel-tgb-tgbl-review-v2/tasks/src-dst-next.zip": "da7b52b1ce1f6505bf786adc9dc124d40c3d8876490e437465f0fb393e18e0b7", + "rel-tgb-tgbl-wiki/tasks/src-dst-next.zip": "29e929befbee3141b1b9c9cb53393e5091a4bfd6cbd8b9cdca47c5084b50122b", + "rel-tgb-tgbl-wiki-v2/tasks/src-dst-next.zip": "33fda785b98feefa76863c6fed006823ca762eb6fedeae4f3a37ba606c62f4a1", + + "rel-tgb-thgl-forum/tasks/type0-type0-next.zip": "5c2c02a2692f80c5f9e0b61f640e14bd48b3d5dc191d12516d02462796636a71", + "rel-tgb-thgl-forum/tasks/type0-type1-next.zip": "3a95191573489f21236fdd008d12615f678266693f5b6fde3a2aba66fcf2d8eb", + "rel-tgb-thgl-github/tasks/type0-type1-next.zip": "de64b8a851f3a129d6142d89e67429dcd44d8e5e43aeada77dc21e560465e859", + "rel-tgb-thgl-github/tasks/type1-type1-next.zip": "efa5b06c2db56673e9e6d6ba3440cea52a0a3258235654ea16edb661d761a5e9", + "rel-tgb-thgl-github/tasks/type2-type0-next.zip": "0f53dc7f8c0307209759aae7a73f83b5f3a77125d203cb894c7456b12310078b", + "rel-tgb-thgl-github/tasks/type2-type1-next.zip": "b5fd82bc7f8f4cf4d423df6eb8d3ecc1ab0e87bb892872dfe6951738163b6da2", + "rel-tgb-thgl-github/tasks/type2-type3-next.zip": "48ea7dcd7fda0cc8057e7e507add042129fd98c6ce01b9252299028094df1bb1", + "rel-tgb-thgl-github/tasks/type3-type1-next.zip": "90ea060aa79e70f11a812c46cc860b618c5786924fa2a1fabd2cfe5f6d8dea7b", + "rel-tgb-thgl-myket/tasks/type0-type1-next.zip": "61f7fdfa03da8dbaee3143c29c16a804afc5299e58a95c3b74622bc81b244348", + "rel-tgb-thgl-software/tasks/type0-type1-next.zip": "370b2427c7f9b25a306e91477aaaca8fa6dcad9cb0861882b02015b1debe1e27", + "rel-tgb-thgl-software/tasks/type0-type3-next.zip": "c0648c2726a4c0b632b74bc39773d15423b1dd1a97b2a55f4123626f5ae3225b", + "rel-tgb-thgl-software/tasks/type1-type2-next.zip": "22e4e8b0578884bd4d5bf9dfb7a07407522a8e4bbd428ff9c47f6a09a2c8e0f1", + "rel-tgb-thgl-software/tasks/type3-type2-next.zip": "d74a1b6ba94e523865d784b7a46453ae72607ab73f79c836e64ae1017ad78d0e", + + "rel-tgb-tgbn-genre/tasks/node-label-next.zip": "044fdd588ebca1830f35862aebef4e997cd9e877d51050ddf36e1d9789baff04", + "rel-tgb-tgbn-reddit/tasks/node-label-next.zip": "087dfc47cd122b8ba3d696998bd3725a94d9cc78217c8451b9c667d93a72ddfe", + "rel-tgb-tgbn-token/tasks/node-label-next.zip": "0e3c3fb43d17b4a46b8ed82135118f8b572952ce0b9efd86d18e6ad580f7587d", + "rel-tgb-tgbn-trade/tasks/node-label-next.zip": "105cc706f346848df2a290e9ddbdca1797d0bd3b3fe0fa1e6af2d3a6234f470f" } diff --git a/relbench/tasks/tgb.py b/relbench/tasks/tgb.py index 0f82943b..2af30b41 100644 --- a/relbench/tasks/tgb.py +++ b/relbench/tasks/tgb.py @@ -96,20 +96,21 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab ev = table.df mask = (ev[self.spec.event_time_col] > start) & (ev[self.spec.event_time_col] <= end) if mask.any(): - window_frames.append(ev.loc[mask, [self.spec.src_col, self.spec.dst_col]]) + window_frames.append(ev.loc[mask, [self.spec.src_col, self.spec.dst_col, self.spec.event_time_col]]) if not window_frames: continue win = pd.concat(window_frames, axis=0, ignore_index=True) - grouped = ( - win.groupby(self.spec.src_col, sort=False)[self.spec.dst_col] - .agg(_unique_list) - .reset_index() - .rename(columns={self.spec.src_col: self.src_entity_col, self.spec.dst_col: self.dst_entity_col}) + # "next" semantics: pick the first destination per source in the window. + win = win.sort_values(self.spec.event_time_col, kind="mergesort") + first = win.drop_duplicates(subset=[self.spec.src_col], keep="first") + out = first[[self.spec.src_col, self.spec.dst_col]].rename( + columns={self.spec.src_col: self.src_entity_col} ) - grouped[self.time_col] = start - rows.append(grouped[[self.time_col, self.src_entity_col, self.dst_entity_col]]) + out[self.dst_entity_col] = out[self.spec.dst_col].apply(lambda x: [int(x)]) + out[self.time_col] = start + rows.append(out[[self.time_col, self.src_entity_col, self.dst_entity_col]]) if rows: out = pd.concat(rows, axis=0, ignore_index=True) @@ -182,8 +183,10 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab if not mask.any(): continue - le_win = le.loc[mask, ["label_event_id", "src_id"]] - joined = le_win.merge(items[["label_event_id", "label_id"]], on="label_event_id", how="inner") + le_win = le.loc[mask, ["label_event_id", "src_id", "label_ts"]].sort_values("label_ts", kind="mergesort") + # "next" semantics: pick the first label_event per node in the window. + le_first = le_win.drop_duplicates(subset=["src_id"], keep="first")[["label_event_id", "src_id"]] + joined = le_first.merge(items[["label_event_id", "label_id"]], on="label_event_id", how="inner") grouped = ( joined.groupby("src_id", sort=False)["label_id"] .agg(_unique_list) From b9b4fd7e497213727d493cb52f2dddf178150027 Mon Sep 17 00:00:00 2001 From: pc0618 Date: Thu, 29 Jan 2026 06:19:52 +0000 Subject: [PATCH 3/8] Update TGB tasks and refresh hashes --- relbench/datasets/hashes.json | 22 +- relbench/modeling/utils.py | 22 +- relbench/tasks/__init__.py | 276 +++++-------- relbench/tasks/hashes.json | 69 ++-- relbench/tasks/tgb.py | 720 +++++++++++++++++++++++++++------- 5 files changed, 735 insertions(+), 374 deletions(-) diff --git a/relbench/datasets/hashes.json b/relbench/datasets/hashes.json index 6f6d3c0c..fd62d217 100644 --- a/relbench/datasets/hashes.json +++ b/relbench/datasets/hashes.json @@ -18,21 +18,21 @@ "dbinfer-seznam/db.zip": "77314bf874dc495e8a4a61f2dc5f12982bbec3c5b6b7af5555e9a2bb587154d9", "dbinfer-stackexchange/db.zip": "4c8b8ac38b56d57bc2ead3c12be1237b69cae76062ed3f176fc02ad327a84d19", - "rel-tgb-tgbl-coin/db.zip": "058917b0b5ac29cadf0c12c07516ff6cc14c4a5eea0a027902425c9e05bf2690", - "rel-tgb-tgbl-comment/db.zip": "3d2ec40a8fb32969615db31fb8fa81a45d4b9a91a8d40e5f8ee7358ea1c8402a", - "rel-tgb-tgbl-flight/db.zip": "eabc26da95efcad0dd2661eb683470a962355bf7229e8490df653cd990dd3627", - "rel-tgb-tgbl-review/db.zip": "274a21fe7aac039b032f75aeba3efbcaf8af5531ea2c106a61b15651440d694e", - "rel-tgb-tgbl-review-v2/db.zip": "a5719eb25505d3500c910c022a68ff886ad9e6b1687035d0f95263f1f2874deb", - "rel-tgb-tgbl-wiki/db.zip": "f32ad5382de5deccb55ecd465e969fae3d8e49168de0731c7ff7ae8f04754dbd", - "rel-tgb-tgbl-wiki-v2/db.zip": "6d620c982a7209e109d052ca964d5c8f8cbae8575f883c65d9ce377343a6a597", + "rel-tgb-tgbl-coin/db.zip": "6de3b62bcfc59bb18ff8de0614ce5cbcf21179ad56cb58920d25ade32ec43e00", + "rel-tgb-tgbl-comment/db.zip": "4eb41776954efa2cb06cdb64e093b182335b8683e60236a6645cf4bcd83597be", + "rel-tgb-tgbl-flight/db.zip": "e82646893a45be6c5312e5dc1a774c08d929ecb9eb309c0ad9fdec6dae3b156a", + "rel-tgb-tgbl-review/db.zip": "63b405aafd8092cbda297694e649329ca806ccc3c49c8e6f53eb4a3c19075091", + "rel-tgb-tgbl-review-v2/db.zip": "a5f1a7a9661a700ebb1a3b43df74037853ad3bd6ed54c74342045d2dc8448bc2", + "rel-tgb-tgbl-wiki/db.zip": "ad7b55d1d7b7125c06588db0b8cbebe87c629a95a6c9a911369f89aca9dffdc9", + "rel-tgb-tgbl-wiki-v2/db.zip": "5eb56f8e459405e3b554ce56585903c80038d84a99ca4333b00ce32c8c7a38f1", "rel-tgb-tgbn-genre/db.zip": "e46aecc28315ca9872117817ce65bc1f4d00ed261ab6d35038a84bcf0ebab7bf", "rel-tgb-tgbn-reddit/db.zip": "bc444f5cbaf7004ef7c6c8f98ae12ae72e8427ccc9710e96e46b0fc156cc952e", "rel-tgb-tgbn-token/db.zip": "44a0f90a62642b8054ef1fbee35c862a67e5d0cec92f08b8fd44c503e0986be8", "rel-tgb-tgbn-trade/db.zip": "8a983ae6281ea058dc8d84f6ce339d2dacf0ae9bb98bb82cca3fed54ec3fe370", - "rel-tgb-thgl-forum/db.zip": "ef3185fd1558711c51306524d07dddb9356d3108eb0af28727ac75c1f73d3c3f", - "rel-tgb-thgl-github/db.zip": "88105bcfc3d1409f5408c3108a6a300c0782bf85264a1e7153bcc01828368078", - "rel-tgb-thgl-myket/db.zip": "4dce093d60c3009b0c85ac0e93da7a6df9b9a1904c6dafd8cf730d0878e0a5aa", - "rel-tgb-thgl-software/db.zip": "9fe3316bf11908065f1eda28f14ef3cbab64e5166bf0ebf7e359315c0528b99d" + "rel-tgb-thgl-forum/db.zip": "fdb20c1afc542e8026b85df850b5e8a539694c182eb4530924fad65964461256", + "rel-tgb-thgl-github/db.zip": "f69c1d49779a4dbead101e7325447854c8890982d254efb2b722117819ba8304", + "rel-tgb-thgl-myket/db.zip": "67a3af25553ceabe8a7eff8906c2213bd82e932cdaeb48ffa326b39d30e0f0cf", + "rel-tgb-thgl-software/db.zip": "35816ef6a07be2291349f081814e18b498a2e46e006ed21be36fdb1a8a0eb90d" } diff --git a/relbench/modeling/utils.py b/relbench/modeling/utils.py index 28011ec7..070d1b81 100644 --- a/relbench/modeling/utils.py +++ b/relbench/modeling/utils.py @@ -10,11 +10,23 @@ def to_unix_time(ser: pd.Series) -> np.ndarray: r"""Converts a :class:`pandas.Timestamp` series to UNIX timestamp (in seconds).""" - assert ser.dtype in [np.dtype("datetime64[s]"), np.dtype("datetime64[ns]")] - unix_time = ser.astype("int64").values - if ser.dtype == np.dtype("datetime64[ns]"): - unix_time //= 10**9 - return unix_time + # Accept tz-aware timestamps and normalize to UTC. + # Many external parquet exports use `datetime64[ns, UTC]`. + if pd.api.types.is_datetime64_any_dtype(ser.dtype) or pd.api.types.is_datetime64tz_dtype(ser.dtype): + ts = pd.to_datetime(ser, utc=True) + unix_ns = ts.astype("int64").to_numpy(copy=False) + return (unix_ns // 1_000_000_000).astype(np.int64, copy=False) + + # Allow integer/float timestamps that are already in seconds. + if pd.api.types.is_integer_dtype(ser.dtype): + return ser.astype("int64").to_numpy(copy=False) + if pd.api.types.is_float_dtype(ser.dtype): + return ser.astype("int64").to_numpy(copy=False) + + # Fallback: parse strings/objects as datetimes. + ts = pd.to_datetime(ser, utc=True) + unix_ns = ts.astype("int64").to_numpy(copy=False) + return (unix_ns // 1_000_000_000).astype(np.int64, copy=False) def remove_pkey_fkey(col_to_stype: Dict[str, Any], table: Table) -> dict: diff --git a/relbench/tasks/__init__.py b/relbench/tasks/__init__.py index 98c3fa1e..efd69910 100644 --- a/relbench/tasks/__init__.py +++ b/relbench/tasks/__init__.py @@ -5,7 +5,6 @@ from functools import lru_cache from typing import List -import pandas as pd import pooch from relbench.base import AutoCompleteTask, BaseTask, TaskType @@ -552,204 +551,121 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: # Community task family: Temporal Graph Benchmark (TGB) # -# Task naming follows: -# - recommendation tasks: -- -# We use "next" as the single-word suffix for all TGB link-style tasks. +# TGB official evaluation for link prediction is one-vs-many with pre-generated +# negatives and reports MRR/Hits@k. For heterogeneous datasets, this is done +# per edge type; we therefore expose one task per `events_edge_type_*` table. +# +# In addition, we expose exporter-defined "next interaction" recommendation +# tasks (`*-next`) to validate RelBench recommendation baselines on the same +# exported schema. These tasks use standard RelBench top-k metrics +# (precision/recall/MAP). -# tgbl-* (link prediction) -register_task( +# tgbl-* (link prediction; single event table) +for dataset_name in [ "rel-tgb-tgbl-wiki", - "src-dst-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="src_nodes", dst_entity_table="dst_nodes"), - timedelta=pd.Timedelta(days=1), - eval_k=10, -) -register_task( "rel-tgb-tgbl-wiki-v2", - "src-dst-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="src_nodes", dst_entity_table="dst_nodes"), - timedelta=pd.Timedelta(days=1), - eval_k=10, -) -register_task( "rel-tgb-tgbl-review", - "src-dst-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), - timedelta=pd.Timedelta(days=180), - eval_k=10, -) -register_task( "rel-tgb-tgbl-review-v2", - "src-dst-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), - timedelta=pd.Timedelta(days=180), - eval_k=10, -) -register_task( "rel-tgb-tgbl-coin", - "src-dst-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), - timedelta=pd.Timedelta(days=14), - eval_k=10, -) -register_task( "rel-tgb-tgbl-comment", - "src-dst-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), - timedelta=pd.Timedelta(days=56), - eval_k=10, -) -register_task( "rel-tgb-tgbl-flight", - "src-dst-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes", dst_entity_table="nodes"), - timedelta=pd.Timedelta(days=90), - eval_k=10, -) +]: + register_task( + dataset_name, + "src-dst-next", + tgb.TGBNextLinkPredTask, + tgb_task_name="src-dst-next", + eval_k=10, + ) + register_task( + dataset_name, + "src-dst-mrr", + tgb.TGBOneVsManyLinkPredTask, + spec=tgb.TGBLinkPredSpec(event_table="events"), + k_value=10, + ) -# thgl-* (heterogeneous link prediction) – expose type-pair tasks using -# node-type table names from the exports. -register_task( - "rel-tgb-thgl-software", - "type0-type1-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_1"), - timedelta=pd.Timedelta(days=1), - eval_k=10, -) -register_task( - "rel-tgb-thgl-software", - "type0-type3-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_3"), - timedelta=pd.Timedelta(days=1), - eval_k=10, -) -register_task( - "rel-tgb-thgl-software", - "type1-type2-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_1", dst_entity_table="nodes_type_2"), - timedelta=pd.Timedelta(days=1), - eval_k=10, -) -register_task( - "rel-tgb-thgl-software", - "type3-type2-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_3", dst_entity_table="nodes_type_2"), - timedelta=pd.Timedelta(days=1), - eval_k=10, -) -register_task( - "rel-tgb-thgl-github", - "type0-type1-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_1"), - timedelta=pd.Timedelta(days=1), - eval_k=10, -) -register_task( - "rel-tgb-thgl-github", - "type1-type1-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_1", dst_entity_table="nodes_type_1"), - timedelta=pd.Timedelta(days=1), - eval_k=10, -) -register_task( - "rel-tgb-thgl-github", - "type2-type0-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_2", dst_entity_table="nodes_type_0"), - timedelta=pd.Timedelta(days=1), - eval_k=10, -) -register_task( - "rel-tgb-thgl-github", - "type2-type1-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_2", dst_entity_table="nodes_type_1"), - timedelta=pd.Timedelta(days=1), - eval_k=10, -) -register_task( - "rel-tgb-thgl-github", - "type2-type3-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_2", dst_entity_table="nodes_type_3"), - timedelta=pd.Timedelta(days=1), - eval_k=10, +def _register_thgl_edge_type_tasks(dataset_name: str, edge_types: list[int]) -> None: + for et in edge_types: + register_task( + dataset_name, + f"edge-type-{int(et)}-mrr", + tgb.TGBOneVsManyLinkPredTask, + spec=tgb.TGBLinkPredSpec(event_table=f"events_edge_type_{int(et)}"), + k_value=10, + ) + + +def _register_thgl_next_tasks(dataset_name: str, next_tasks: list[str]) -> None: + for name in next_tasks: + register_task( + dataset_name, + name, + tgb.TGBNextLinkPredTask, + tgb_task_name=name, + eval_k=10, + ) + + +_register_thgl_edge_type_tasks("rel-tgb-thgl-software", list(range(14))) +_register_thgl_edge_type_tasks("rel-tgb-thgl-github", list(range(14))) +_register_thgl_edge_type_tasks("rel-tgb-thgl-forum", [0, 1]) +_register_thgl_edge_type_tasks("rel-tgb-thgl-myket", [0, 1]) + +_register_thgl_next_tasks( + "rel-tgb-thgl-software", + [ + "type0-type1-next", + "type0-type3-next", + "type1-type2-next", + "type3-type2-next", + ], ) -register_task( +_register_thgl_next_tasks( "rel-tgb-thgl-github", - "type3-type1-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_3", dst_entity_table="nodes_type_1"), - timedelta=pd.Timedelta(days=1), - eval_k=10, -) - -register_task( - "rel-tgb-thgl-forum", - "type0-type0-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_0"), - timedelta=pd.Timedelta(days=1), - eval_k=10, + [ + "type0-type1-next", + "type1-type1-next", + "type2-type0-next", + "type2-type1-next", + "type2-type3-next", + "type3-type1-next", + ], ) -register_task( +_register_thgl_next_tasks( "rel-tgb-thgl-forum", - "type0-type1-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_1"), - timedelta=pd.Timedelta(days=1), - eval_k=10, + [ + "type0-type0-next", + "type0-type1-next", + ], ) - -register_task( +_register_thgl_next_tasks( "rel-tgb-thgl-myket", - "type0-type1-next", - tgb.TGBSrcDstNextTask, - spec=tgb.TGBAggregatedEventsSpec(src_entity_table="nodes_type_0", dst_entity_table="nodes_type_1"), - timedelta=pd.Timedelta(days=7), - eval_k=10, + [ + "type0-type1-next", + ], ) -# tgbn-* (node property prediction as node→label recommendation) -register_task( +# tgbn-* (node property prediction; NDCG@10 over label distributions) +for dataset_name in [ "rel-tgb-tgbn-trade", - "node-label-next", - tgb.TGBNodeLabelNextTask, - timedelta=pd.Timedelta(days=365), - eval_k=10, -) -register_task( "rel-tgb-tgbn-genre", - "node-label-next", - tgb.TGBNodeLabelNextTask, - timedelta=pd.Timedelta(days=7), - eval_k=10, -) -register_task( "rel-tgb-tgbn-reddit", - "node-label-next", - tgb.TGBNodeLabelNextTask, - timedelta=pd.Timedelta(days=7), - eval_k=10, -) -register_task( "rel-tgb-tgbn-token", - "node-label-next", - tgb.TGBNodeLabelNextTask, - timedelta=pd.Timedelta(days=7), - eval_k=10, -) +]: + register_task( + dataset_name, + "node-label-next", + tgb.TGBNextLinkPredTask, + tgb_task_name="node-label-next", + dst_entity_col="label_id", + eval_k=10, + ) + register_task( + dataset_name, + "node-label-ndcg", + tgb.TGBNodePropNDCGTask, + spec=tgb.TGBNodePropSpec(), + k=10, + ) diff --git a/relbench/tasks/hashes.json b/relbench/tasks/hashes.json index 820506f5..86289918 100644 --- a/relbench/tasks/hashes.json +++ b/relbench/tasks/hashes.json @@ -82,30 +82,47 @@ "dbinfer-stackexchange/tasks/churn.zip": "d649b0a70dc8be9fca797c82e14d19ed4311637711144f4c5b74249081b107ca", "dbinfer-stackexchange/tasks/upvote.zip": "af8f516422e3e19c0d76f13c00c4c0577cdc218b84468006eda369de8f9c7953", - "rel-tgb-tgbl-coin/tasks/src-dst-next.zip": "3e9f11e204bc0cc61f9ba07c5f32abc78e1de4f482804652b307aa231171301a", - "rel-tgb-tgbl-comment/tasks/src-dst-next.zip": "5cb7f475382937666237471112f225cc3aa542c5842e32b948dcb8f5002aba20", - "rel-tgb-tgbl-flight/tasks/src-dst-next.zip": "7e37fd976eb8a468b06fdcd0be6ba987dfca117cb06f7dcc18f6d080ab39ee21", - "rel-tgb-tgbl-review/tasks/src-dst-next.zip": "3c87d7c3fa7cc5f60c59882f56cf49a37a342bfc9766bd067d57ee8e637578e5", - "rel-tgb-tgbl-review-v2/tasks/src-dst-next.zip": "da7b52b1ce1f6505bf786adc9dc124d40c3d8876490e437465f0fb393e18e0b7", - "rel-tgb-tgbl-wiki/tasks/src-dst-next.zip": "29e929befbee3141b1b9c9cb53393e5091a4bfd6cbd8b9cdca47c5084b50122b", - "rel-tgb-tgbl-wiki-v2/tasks/src-dst-next.zip": "33fda785b98feefa76863c6fed006823ca762eb6fedeae4f3a37ba606c62f4a1", - - "rel-tgb-thgl-forum/tasks/type0-type0-next.zip": "5c2c02a2692f80c5f9e0b61f640e14bd48b3d5dc191d12516d02462796636a71", - "rel-tgb-thgl-forum/tasks/type0-type1-next.zip": "3a95191573489f21236fdd008d12615f678266693f5b6fde3a2aba66fcf2d8eb", - "rel-tgb-thgl-github/tasks/type0-type1-next.zip": "de64b8a851f3a129d6142d89e67429dcd44d8e5e43aeada77dc21e560465e859", - "rel-tgb-thgl-github/tasks/type1-type1-next.zip": "efa5b06c2db56673e9e6d6ba3440cea52a0a3258235654ea16edb661d761a5e9", - "rel-tgb-thgl-github/tasks/type2-type0-next.zip": "0f53dc7f8c0307209759aae7a73f83b5f3a77125d203cb894c7456b12310078b", - "rel-tgb-thgl-github/tasks/type2-type1-next.zip": "b5fd82bc7f8f4cf4d423df6eb8d3ecc1ab0e87bb892872dfe6951738163b6da2", - "rel-tgb-thgl-github/tasks/type2-type3-next.zip": "48ea7dcd7fda0cc8057e7e507add042129fd98c6ce01b9252299028094df1bb1", - "rel-tgb-thgl-github/tasks/type3-type1-next.zip": "90ea060aa79e70f11a812c46cc860b618c5786924fa2a1fabd2cfe5f6d8dea7b", - "rel-tgb-thgl-myket/tasks/type0-type1-next.zip": "61f7fdfa03da8dbaee3143c29c16a804afc5299e58a95c3b74622bc81b244348", - "rel-tgb-thgl-software/tasks/type0-type1-next.zip": "370b2427c7f9b25a306e91477aaaca8fa6dcad9cb0861882b02015b1debe1e27", - "rel-tgb-thgl-software/tasks/type0-type3-next.zip": "c0648c2726a4c0b632b74bc39773d15423b1dd1a97b2a55f4123626f5ae3225b", - "rel-tgb-thgl-software/tasks/type1-type2-next.zip": "22e4e8b0578884bd4d5bf9dfb7a07407522a8e4bbd428ff9c47f6a09a2c8e0f1", - "rel-tgb-thgl-software/tasks/type3-type2-next.zip": "d74a1b6ba94e523865d784b7a46453ae72607ab73f79c836e64ae1017ad78d0e", - - "rel-tgb-tgbn-genre/tasks/node-label-next.zip": "044fdd588ebca1830f35862aebef4e997cd9e877d51050ddf36e1d9789baff04", - "rel-tgb-tgbn-reddit/tasks/node-label-next.zip": "087dfc47cd122b8ba3d696998bd3725a94d9cc78217c8451b9c667d93a72ddfe", - "rel-tgb-tgbn-token/tasks/node-label-next.zip": "0e3c3fb43d17b4a46b8ed82135118f8b572952ce0b9efd86d18e6ad580f7587d", - "rel-tgb-tgbn-trade/tasks/node-label-next.zip": "105cc706f346848df2a290e9ddbdca1797d0bd3b3fe0fa1e6af2d3a6234f470f" + "rel-tgb-tgbl-coin/tasks/src-dst-mrr.zip": "702ff505d95955b29c9998eb4f9c2e38e6c112c2b9e21ed7c2aa75ae4a923280", + "rel-tgb-tgbl-comment/tasks/src-dst-mrr.zip": "50846c6b82433bcac2bb025c3e7edfed46b651a9d06b462b058a071f08741fdc", + "rel-tgb-tgbl-flight/tasks/src-dst-mrr.zip": "10527b2122a73bcb35a6b3e6230ded99242825bf955c363b710cce20545c8104", + "rel-tgb-tgbl-review-v2/tasks/src-dst-mrr.zip": "aeded6586c00995a33881305717c3a625f96a1cc7253d2806af6125b53173fb3", + "rel-tgb-tgbl-review/tasks/src-dst-mrr.zip": "aeded6586c00995a33881305717c3a625f96a1cc7253d2806af6125b53173fb3", + "rel-tgb-tgbl-wiki-v2/tasks/src-dst-mrr.zip": "9517a458de6dc921e54af91154f1823d7bdbfa16207650df16a32bcd9b1aab92", + "rel-tgb-tgbl-wiki/tasks/src-dst-mrr.zip": "a105ef6cf8dcdd5f94d1fbcfa856a3ecf10821a9d217b53c8266eca2d8b04b1b", + "rel-tgb-tgbn-genre/tasks/node-label-ndcg.zip": "4fe66c0e9e0442fa9c4979fa386f8249f812b19b817faaefae46f3b82f8c3dce", + "rel-tgb-tgbn-reddit/tasks/node-label-ndcg.zip": "bdcdb15597719541e2873408d36ccc7589ddc25e8fd9be60d3bc2928a4fb02d5", + "rel-tgb-tgbn-token/tasks/node-label-ndcg.zip": "64daa8ca8d3ff62bff0929f8af8dcbb0836c004646102df94df8236b13fffba1", + "rel-tgb-tgbn-trade/tasks/node-label-ndcg.zip": "f69920871dac5e365fa521d180d9a234802318bdbb313a59a362fe6c68f14e1d", + "rel-tgb-thgl-forum/tasks/edge-type-0-mrr.zip": "14975fa000b5a4d8458fe5f9fd23d6662af8205d4ef9d3463da2051b16b3a5a8", + "rel-tgb-thgl-forum/tasks/edge-type-1-mrr.zip": "fffe1b4dc3ed985e1d442d3db56da2ebfb18b05ccf3fefe0d13ae64431993027", + "rel-tgb-thgl-github/tasks/edge-type-0-mrr.zip": "9c346bf63c2dcc36171b62884c18c4062781ca689ec861177d70f80531856e43", + "rel-tgb-thgl-github/tasks/edge-type-1-mrr.zip": "e0db81eec7a8b12c488fe7c773f7226e7f7914ec5613b9ba11192a3bb2c19e21", + "rel-tgb-thgl-github/tasks/edge-type-10-mrr.zip": "dcfaa9094d7d0f7944dda32075d165c5f97c4d0d2530dced6e27121163599a65", + "rel-tgb-thgl-github/tasks/edge-type-11-mrr.zip": "7f707920f58c4fea8dabcfe0e511be1a60603ba73301f28a698a4fce87e804da", + "rel-tgb-thgl-github/tasks/edge-type-12-mrr.zip": "58cf770c6917081865c4c8f5660670e3f821618bc9a3a03f57d5a12891254ea4", + "rel-tgb-thgl-github/tasks/edge-type-13-mrr.zip": "cfaa8350612cc1aa99d7dcaccd32849bae294e3b2ef8e3cf2232c8b82e45650c", + "rel-tgb-thgl-github/tasks/edge-type-2-mrr.zip": "d42be110b69c9dae1d952459118864af5701a18393a3ca11ddb96766889ea5be", + "rel-tgb-thgl-github/tasks/edge-type-3-mrr.zip": "61e31b44e361a23c498e7aacdc271ba8b4b276b0134d1c07e0ecd694ab732665", + "rel-tgb-thgl-github/tasks/edge-type-4-mrr.zip": "61832694d14aa45ddf393122e1530d63a0c93d902913767d29928bbadf208f0e", + "rel-tgb-thgl-github/tasks/edge-type-5-mrr.zip": "a5378c8797a79b655b6eb4678ce2b6f5d1d39e992e6ceda08a1a6f3f7b26cd66", + "rel-tgb-thgl-github/tasks/edge-type-6-mrr.zip": "76d606d9f7d6c69be22b39e7b1086a7d75fd47b29f682538153fdc9ccf08e1bf", + "rel-tgb-thgl-github/tasks/edge-type-7-mrr.zip": "7c9405cc05cad4f299967aee0510eaf718a2dee1884bac790e76c81263b72f22", + "rel-tgb-thgl-github/tasks/edge-type-8-mrr.zip": "44cad0c9358293530781d20b78cb87bd21ebf8fea36628c0875ed4558a6a9117", + "rel-tgb-thgl-github/tasks/edge-type-9-mrr.zip": "19be77af66f1cb0cb640158f97d42a0d16431eb79bd23ed1aaace4b68519c8a9", + "rel-tgb-thgl-myket/tasks/edge-type-0-mrr.zip": "979133caaa02ddf444e3b23ada4b5cbe11147b18adab90537740cea914e36daa", + "rel-tgb-thgl-myket/tasks/edge-type-1-mrr.zip": "f4f84ca245d499c0773ec0c95619e5f6041461c3233cce96ad3a6302da41c12c", + "rel-tgb-thgl-software/tasks/edge-type-0-mrr.zip": "6262267aae8309bda302c64e365d227d036f95f934ee8ba6fb038f20d28b3fed", + "rel-tgb-thgl-software/tasks/edge-type-1-mrr.zip": "7de5c8664e9b78460cad9345a4a537e9d75641f0875605a5151e189ab568fd0e", + "rel-tgb-thgl-software/tasks/edge-type-10-mrr.zip": "b3a1ed7410ed737fb2ca6ecd4f1de59cc0d7e0cfd29fdcb68ccd13e76e1997e3", + "rel-tgb-thgl-software/tasks/edge-type-11-mrr.zip": "9e24b4c1365dfb989fc6c4db4cc19146e998f852205fca0fe7f1fe4a047ef44d", + "rel-tgb-thgl-software/tasks/edge-type-12-mrr.zip": "1580ff8adde8202edae05a555b12a7fe38ce0addd360010ec925b4a7503350d2", + "rel-tgb-thgl-software/tasks/edge-type-13-mrr.zip": "e4a51a47777c34748fec03813bb2368bf841f90479bd37ca41c653b9203491cb", + "rel-tgb-thgl-software/tasks/edge-type-2-mrr.zip": "1b8177e304b62939b3081c054b56b99d5874300e101057593668af455558b385", + "rel-tgb-thgl-software/tasks/edge-type-3-mrr.zip": "e0837e81655cfe953b805af162e7904dfe457a93783701d50a26d819dfec848f", + "rel-tgb-thgl-software/tasks/edge-type-4-mrr.zip": "0d15eeaf8d9e456912518aedd2e6a239b6fddc6941095b95240ef41fb34f14cc", + "rel-tgb-thgl-software/tasks/edge-type-5-mrr.zip": "1df8a435f4e978c976288c3d0a53f7430e759d27cd91ad966a1af6c3814da480", + "rel-tgb-thgl-software/tasks/edge-type-6-mrr.zip": "51fa5917363476973c048ef19306188d0d8cf6c218e316f4949d08d10c9092f0", + "rel-tgb-thgl-software/tasks/edge-type-7-mrr.zip": "ee8a0e231255e15fe58283b9f9243ca438695e1760db8cb24168b6bfbb6c078a", + "rel-tgb-thgl-software/tasks/edge-type-8-mrr.zip": "26e37301204380e27934ed7046201ebd4edf72ac7b430c8ef792ad3defc9e638", + "rel-tgb-thgl-software/tasks/edge-type-9-mrr.zip": "d68f2bcafee7a7d7761e4b4eadec36c500cb3e8711d1b23196ecb160962b5b9e" } diff --git a/relbench/tasks/tgb.py b/relbench/tasks/tgb.py index 2af30b41..ba370db1 100644 --- a/relbench/tasks/tgb.py +++ b/relbench/tasks/tgb.py @@ -1,212 +1,628 @@ from __future__ import annotations +import pickle +import re from dataclasses import dataclass -from typing import Optional +from functools import lru_cache +from pathlib import Path +from typing import Any, Optional +import numpy as np import pandas as pd -from relbench.base import Database, RecommendationTask, Table, TaskType +from relbench.base import BaseTask, Dataset, RecommendationTask, Table, TaskType from relbench.metrics import ( link_prediction_map, - link_prediction_ndcg, link_prediction_precision, link_prediction_recall, ) -def _unique_list(values: pd.Series) -> list[int]: - # Preserve deterministic ordering per group as first-appearance order. - # Pandas' `unique()` preserves order. - return [int(x) for x in values.unique().tolist()] +def _to_unix_seconds(ts: pd.Series) -> np.ndarray: + ts = pd.to_datetime(ts, utc=True) + # Timestamp[ns] -> seconds + return (ts.astype("int64").to_numpy(copy=False) // 1_000_000_000).astype(np.int64, copy=False) + + +def _tgb_eval_hits_and_mrr( + y_pred_pos: np.ndarray, + y_pred_neg: np.ndarray, + *, + k_value: int, +) -> dict[str, float]: + r"""Match TGB's link prediction evaluator for one-vs-many ranking. + + This mirrors `tgb.linkproppred.evaluate.Evaluator._eval_hits_and_mrr`: + - optimistic rank counts negatives strictly greater than positive + - pessimistic rank counts negatives greater-or-equal to positive + - final rank is the average of the two (plus 1) + """ + + y_pred_pos = np.asarray(y_pred_pos).reshape(-1, 1) + y_pred_neg = np.asarray(y_pred_neg) + if y_pred_neg.ndim != 2 or y_pred_neg.shape[0] != y_pred_pos.shape[0]: + raise ValueError( + "Expected y_pred_neg with shape (N, num_neg) matching y_pred_pos (N,). " + f"Got y_pred_pos={y_pred_pos.shape}, y_pred_neg={y_pred_neg.shape}." + ) + + optimistic_rank = (y_pred_neg > y_pred_pos).sum(axis=1) + pessimistic_rank = (y_pred_neg >= y_pred_pos).sum(axis=1) + ranking_list = 0.5 * (optimistic_rank + pessimistic_rank) + 1 + + hits_k = (ranking_list <= int(k_value)).astype(np.float32).mean().item() + mrr = (1.0 / ranking_list.astype(np.float32)).mean().item() + return {f"hits@{int(k_value)}": hits_k, "mrr": mrr} @dataclass(frozen=True) -class TGBAggregatedEventsSpec: - src_entity_table: str - dst_entity_table: str - event_tables: tuple[str, ...] = () - event_time_col: str = "event_ts" +class TGBLinkPredSpec: + r"""Specification for a TGB link prediction task over a single event table.""" + + event_table: str src_col: str = "src_id" dst_col: str = "dst_id" + # Matches the exporter convention (`*_exports/**/db/events*.parquet`). + time_col: str = "event_ts" -class TGBSrcDstNextTask(RecommendationTask): - r"""Recommendation task for TGB-style edge streams. +class TGBOneVsManyLinkPredTask(BaseTask): + r"""TGB-style link prediction task with official one-vs-many MRR/Hits@k evaluation. - For each anchor timestamp `t`, predict the set of destination entities that - connect to a source entity within the horizon `(t, t + timedelta]`. + Important: Exact parity with TGB requires: + - task tables (`train/val/test.parquet`) listing the *exact* positive edges + for each split (matching TGB masks, not just timestamp cutoffs), and + - the TGB-provided pre-generated negative samples placed at: + `/negatives/val_ns.pkl` and `.../test_ns.pkl`. - This task is intentionally generic to support: - - homogeneous `tgbl-*` exports: `nodes` + `events` - - bipartite `tgbl-wiki*` exports: `src_nodes`/`dst_nodes` + `events` - - heterogeneous `thgl-*` exports: multiple `events_edge_type_*` tables - aggregated by source/destination node types + For heterogeneous `thgl-*` datasets, TGB negative samples are keyed by + `(timestamp_s, src_global_id, edge_type_id)`, which requires the dataset to + ship mapping files that relate RelBench's per-type local node ids back to + the original TGB global node ids. """ task_type = TaskType.LINK_PREDICTION - metrics = [ - link_prediction_recall, - link_prediction_precision, - link_prediction_map, - link_prediction_ndcg, - ] - - # Convention: keep these stable across all TGB tasks. - src_entity_col = "src_id" - dst_entity_col = "dst_id" - time_col = "timestamp" + # BaseTask enforces `dataset.test_timestamp - dataset.val_timestamp >= timedelta`. + # We override table construction, so we just need a non-negative timedelta. + timedelta = pd.Timedelta(seconds=1) + metrics = [] + num_eval_timestamps = 1 def __init__( self, - dataset, + dataset: Dataset, *, - spec: TGBAggregatedEventsSpec, - timedelta: pd.Timedelta, - eval_k: int = 10, + spec: TGBLinkPredSpec, + k_value: int = 10, cache_dir: Optional[str] = None, ) -> None: self.spec = spec - self.timedelta = pd.Timedelta(timedelta) - self.eval_k = int(eval_k) - self.src_entity_table = spec.src_entity_table - self.dst_entity_table = spec.dst_entity_table + self.k_value = int(k_value) + + self.time_col = spec.time_col + self.src_entity_col = spec.src_col + self.dst_entity_col = spec.dst_col + + # Infer FK target tables + (optional) edge_type id from the exported database. + db = dataset.get_db() + if spec.event_table not in db.table_dict: + raise ValueError( + f"Event table '{spec.event_table}' not found in dataset db. " + f"Available tables: {sorted(db.table_dict.keys())}" + ) + event_tbl = db.table_dict[spec.event_table] + self.src_entity_table = event_tbl.fkey_col_to_pkey_table.get(spec.src_col) + self.dst_entity_table = event_tbl.fkey_col_to_pkey_table.get(spec.dst_col) + if self.src_entity_table is None or self.dst_entity_table is None: + raise ValueError( + f"Event table '{spec.event_table}' must have fkeys for " + f"'{spec.src_col}' and '{spec.dst_col}'. Got {event_tbl.fkey_col_to_pkey_table}." + ) + + m = re.fullmatch(r"events_edge_type_(\d+)", spec.event_table) + self.edge_type_id: Optional[int] = int(m.group(1)) if m else None + super().__init__(dataset, cache_dir=cache_dir) - def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Table: - event_tables = ( - self.spec.event_tables - if self.spec.event_tables - else tuple(name for name in db.table_dict.keys() if name.startswith("events")) + def make_table(self, db, timestamps): # pragma: no cover + raise RuntimeError( + "TGBOneVsManyLinkPredTask expects precomputed task tables " + "(train/val/test.parquet) and overrides _get_table()." ) - rows: list[pd.DataFrame] = [] - for t in timestamps: - start = pd.Timestamp(t) - end = start + self.timedelta - - window_frames: list[pd.DataFrame] = [] - for table_name in event_tables: - table = db.table_dict[table_name] - # For heterogeneous exports, keep only event tables whose FK schema - # matches the intended src/dst entity tables. - if table.fkey_col_to_pkey_table.get(self.spec.src_col) != self.src_entity_table: - continue - if table.fkey_col_to_pkey_table.get(self.spec.dst_col) != self.dst_entity_table: - continue - - ev = table.df - mask = (ev[self.spec.event_time_col] > start) & (ev[self.spec.event_time_col] <= end) - if mask.any(): - window_frames.append(ev.loc[mask, [self.spec.src_col, self.spec.dst_col, self.spec.event_time_col]]) - - if not window_frames: - continue - - win = pd.concat(window_frames, axis=0, ignore_index=True) - # "next" semantics: pick the first destination per source in the window. - win = win.sort_values(self.spec.event_time_col, kind="mergesort") - first = win.drop_duplicates(subset=[self.spec.src_col], keep="first") - out = first[[self.spec.src_col, self.spec.dst_col]].rename( - columns={self.spec.src_col: self.src_entity_col} + + def filter_dangling_entities(self, table: Table) -> Table: + # Keep parity with TGB negatives by dropping rows with invalid ids. + if self.src_entity_table: + table.df = table.df[table.df[self.src_entity_col] < len(self.dataset.get_db().table_dict[self.src_entity_table])] + if self.dst_entity_table: + table.df = table.df[table.df[self.dst_entity_col] < len(self.dataset.get_db().table_dict[self.dst_entity_table])] + table.df = table.df.reset_index(drop=True) + return table + + def _get_table(self, split: str) -> Table: + # Load cached task tables rather than generating by timestamps. + if split not in ["train", "val", "test"]: + raise ValueError(f"Unknown split '{split}'.") + table_path = f"{self.cache_dir}/{split}.parquet" + if not self.cache_dir or not Path(table_path).exists(): + raise RuntimeError( + "Exact TGB parity requires precomputed task tables. " + f"Missing {table_path}. Use download=True or place the parquet files in cache." ) - out[self.dst_entity_col] = out[self.spec.dst_col].apply(lambda x: [int(x)]) - out[self.time_col] = start - rows.append(out[[self.time_col, self.src_entity_col, self.dst_entity_col]]) - if rows: - out = pd.concat(rows, axis=0, ignore_index=True) + table = Table.load(table_path) + return self.filter_dangling_entities(table) + + def _negatives_path(self, split: str) -> Path: + if split not in ["val", "test"]: + raise ValueError("Negative samples are only defined for val/test splits.") + if self.dataset.cache_dir is None: + raise RuntimeError("Dataset has no cache_dir; cannot locate TGB negatives.") + return Path(self.dataset.cache_dir) / "negatives" / f"{split}_ns.pkl" + + # Large datasets ship multi-GB negative dictionaries. Keep at most one split + # in memory at a time to reduce peak RAM usage. + @lru_cache(maxsize=1) + def _load_negatives(self, split: str) -> dict[Any, Any]: + path = self._negatives_path(split) + if not path.exists(): + raise RuntimeError( + f"Missing TGB negative samples at {path}. " + "To match TGB exactly, include the official negative sample pickle " + "in the dataset download under `negatives/`." + ) + with path.open("rb") as f: + return pickle.load(f) + + def _mapping_paths(self) -> dict[str, Path]: + if self.dataset.cache_dir is None: + raise RuntimeError("Dataset has no cache_dir; cannot locate mapping files.") + base = Path(self.dataset.cache_dir) / "mappings" + return { + "node_type": base / "node_type.npy", + "local_id": base / "local_id.npy", + } + + @lru_cache(maxsize=None) + def _load_global_to_local(self) -> tuple[np.ndarray, np.ndarray]: + paths = self._mapping_paths() + node_type_path = paths["node_type"] + local_id_path = paths["local_id"] + if not node_type_path.exists() or not local_id_path.exists(): + raise RuntimeError( + "Missing heterogeneous mapping files. Expected:\n" + f"- {node_type_path}\n- {local_id_path}\n" + "These are required to map TGB global ids to RelBench per-type local ids." + ) + node_type = np.load(node_type_path) + local_id = np.load(local_id_path) + return node_type.astype(np.int64, copy=False), local_id.astype(np.int64, copy=False) + + @lru_cache(maxsize=None) + def _load_local_to_global(self, node_type_id: int) -> np.ndarray: + if self.dataset.cache_dir is None: + raise RuntimeError("Dataset has no cache_dir; cannot locate mapping files.") + p = Path(self.dataset.cache_dir) / "mappings" / f"globals_type_{int(node_type_id)}.npy" + if not p.exists(): + raise RuntimeError( + f"Missing mapping file {p}. This is required to map local ids " + "back to TGB global ids for negative-sample lookup." + ) + return np.load(p).astype(np.int64, copy=False) + + def _node_type_id_from_table(self, table_name: str) -> Optional[int]: + m = re.fullmatch(r"nodes_type_(\d+)", str(table_name)) + return int(m.group(1)) if m else None + + def _bipartite_offset(self) -> Optional[int]: + # Export convention for bipartite tgbl-wiki*: src ids are 0..num_src-1, + # dst ids are 0..num_dst-1 in RelBench, but TGB negatives use a single + # global id space with dst shifted by `num_src`. + if self.src_entity_table == "src_nodes" and self.dst_entity_table == "dst_nodes": + return len(self.dataset.get_db().table_dict["src_nodes"]) + return None + + def _src_local_to_global(self, src_local: np.ndarray) -> np.ndarray: + src_type = self._node_type_id_from_table(self.src_entity_table) + if src_type is None: + # homogeneous/bipartite case: local ids are global ids + return src_local.astype(np.int64, copy=False) + globals_ = self._load_local_to_global(src_type) + return globals_[src_local.astype(np.int64, copy=False)] + + def _dst_global_to_local(self, dst_global: np.ndarray) -> np.ndarray: + dst_type = self._node_type_id_from_table(self.dst_entity_table) + if dst_type is None: + offset = self._bipartite_offset() + if offset is None: + return dst_global.astype(np.int64, copy=False) + dst_global = dst_global.astype(np.int64, copy=False) + out = dst_global - int(offset) + if (out < 0).any(): + raise RuntimeError("Bipartite negatives contain ids outside destination range.") + return out.astype(np.int64, copy=False) + node_type, local_id = self._load_global_to_local() + dst_global = dst_global.astype(np.int64, copy=False) + bad = node_type[dst_global] != dst_type + if bad.any(): + raise RuntimeError("Negative samples contain destination nodes of unexpected type.") + return local_id[dst_global].astype(np.int64, copy=False) + + def get_negative_dsts_local(self, *, split: str, table: Optional[Table] = None) -> list[np.ndarray]: + r"""Return negative destination ids (local to dst entity table) for each row. + + This is intended to help users reproduce the TGB evaluation protocol, + i.e., score the true destination vs the provided negatives. + """ + if split not in ["val", "test"]: + raise ValueError("Negatives are only defined for val/test splits.") + if table is None: + table = self.get_table(split, mask_input_cols=False) + + df = table.df + ts_s = _to_unix_seconds(df[self.time_col]) + src_local = df[self.src_entity_col].astype("int64").to_numpy() + src_global = self._src_local_to_global(src_local) + neg_dict = self._load_negatives(split) + + negs_local: list[np.ndarray] = [] + if self.edge_type_id is None: + # tgbl-* negatives are keyed by (src, dst, t) + dst_local = df[self.dst_entity_col].astype("int64").to_numpy() + offset = self._bipartite_offset() + if offset is None: + dst_key = dst_local + else: + dst_key = dst_local + int(offset) + for s, d, t in zip(src_global.tolist(), dst_key.tolist(), ts_s.tolist()): + negs_g = np.asarray(neg_dict[(s, d, t)], dtype=np.int64) + negs_local.append(self._dst_global_to_local(negs_g)) else: - out = pd.DataFrame({self.time_col: [], self.src_entity_col: [], self.dst_entity_col: []}) - - return Table( - df=out, - fkey_col_to_pkey_table={ - self.src_entity_col: self.src_entity_table, - self.dst_entity_col: self.dst_entity_table, - }, - pkey_col=None, - time_col=self.time_col, - ) + # thgl-* negatives are keyed by (t, src, edge_type) + et = int(self.edge_type_id) + for t, s in zip(ts_s.tolist(), src_global.tolist()): + negs_g = np.asarray(neg_dict[(t, s, et)], dtype=np.int64) + negs_local.append(self._dst_global_to_local(negs_g)) + return negs_local + + def evaluate( + self, + pred: Any, + target_table: Optional[Table] = None, + metrics: Optional[list] = None, + ) -> dict[str, float]: + r"""Evaluate predictions using TGB's one-vs-many MRR/Hits@k. + + Expected `pred` formats: + - dict with keys `y_pred_pos` (shape [N]) and `y_pred_neg` (shape [N, K]) + - numpy array of shape [N, 1+K] where pred[:, 0] is positive score + """ + if target_table is None: + target_table = self.get_table("test", mask_input_cols=False) + + if isinstance(pred, dict): + y_pred_pos = np.asarray(pred.get("y_pred_pos")) + y_pred_neg = np.asarray(pred.get("y_pred_neg")) + else: + arr = np.asarray(pred) + if arr.ndim != 2 or arr.shape[1] < 2: + raise ValueError( + "For TGBOneVsManyLinkPredTask, pred must be a dict with " + "'y_pred_pos'/'y_pred_neg' or an array shaped (N, 1+K)." + ) + y_pred_pos = arr[:, 0] + y_pred_neg = arr[:, 1:] + + if y_pred_pos.shape[0] != len(target_table): + raise ValueError( + f"Prediction length {y_pred_pos.shape[0]} does not match target table rows {len(target_table)}." + ) + return _tgb_eval_hits_and_mrr(y_pred_pos, y_pred_neg, k_value=self.k_value) -class TGBNodeLabelNextTask(RecommendationTask): - r"""Node property prediction as a recommendation task. - For each anchor timestamp `t`, predict the set of labels associated with a - node within the horizon `(t, t + timedelta]`. +class TGBNextLinkPredTask(RecommendationTask): + r"""TGB-style "next" link prediction task in RelBench recommendation format. - This task assumes the standard `tgbn-*` export schema: - - `nodes(node_id)` - - `labels(label_id)` - - `label_events(label_event_id, src_id, label_ts)` - - `label_event_items(item_id, label_event_id, label_id, label_weight)` + This task is backed by precomputed task tables (train/val/test.parquet) in + RelBench RecommendationTask format: + - columns: `timestamp`, `src_id`, and a list-valued destination column + (e.g. `dst_id` or `label_id`) - Note: This task treats labels as unweighted positives for evaluation - (binary relevance at top-k). If graded relevance is desired, a custom task - evaluator would be required. + Notes: + - This task is intended to validate RelBench baselines (e.g. + `examples/gnn_recommendation.py`, `examples/tgn_attention_recommendation.py`) + on the translated TGB exports. + - We do not currently support building these tables from scratch, since + exact parity with TGB-style "next interaction" semantics is exporter-defined. """ task_type = TaskType.LINK_PREDICTION - metrics = [ - link_prediction_recall, - link_prediction_precision, - link_prediction_map, - link_prediction_ndcg, - ] - - src_entity_col = "src_id" - dst_entity_col = "label_id" - time_col = "timestamp" + # BaseTask enforces `dataset.test_timestamp - dataset.val_timestamp >= timedelta`. + # We only need a non-negative timedelta since we rely on cached tables. + timedelta = pd.Timedelta(seconds=1) + num_eval_timestamps = 1 + metrics = [link_prediction_precision, link_prediction_recall, link_prediction_map] def __init__( self, - dataset, + dataset: Dataset, *, - timedelta: pd.Timedelta, + tgb_task_name: str, eval_k: int = 10, + src_entity_col: str = "src_id", + dst_entity_col: str = "dst_id", + time_col: str = "timestamp", cache_dir: Optional[str] = None, ) -> None: - self.timedelta = pd.Timedelta(timedelta) self.eval_k = int(eval_k) - self.src_entity_table = "nodes" - self.dst_entity_table = "labels" + + # Column conventions of exported task tables. + self.time_col = str(time_col) + self.src_entity_col = str(src_entity_col) + self.dst_entity_col = str(dst_entity_col) + + # Infer src/dst entity tables from the cached task table metadata if present. + src_table = None + dst_table = None + if cache_dir is not None: + cache_path = Path(cache_dir) + for split in ["train", "val", "test"]: + p = cache_path / f"{split}.parquet" + if p.exists(): + tbl = Table.load(p) + src_table = tbl.fkey_col_to_pkey_table.get(self.src_entity_col) + dst_table = tbl.fkey_col_to_pkey_table.get(self.dst_entity_col) + break + + if src_table is None or dst_table is None: + # Fallback to inferring from task naming conventions: + # - `src-dst-next` for tgbl-* exports (bipartite or homogeneous) + # - `typeX-typeY-next` for thgl-* exports (per node type) + task_name = str(tgb_task_name) + m = re.fullmatch(r"type(\d+)-type(\d+)-next", task_name) + if m is not None: + src_table = f"nodes_type_{int(m.group(1))}" + dst_table = f"nodes_type_{int(m.group(2))}" + elif task_name == "src-dst-next": + db = dataset.get_db() + if "src_nodes" in db.table_dict and "dst_nodes" in db.table_dict: + src_table = "src_nodes" + dst_table = "dst_nodes" + else: + src_table = "nodes" + dst_table = "nodes" + elif task_name == "node-label-next": + # tgbn-* exports: predict labels for a node. + src_table = "nodes" + dst_table = "labels" + else: + raise ValueError( + f"Unable to infer src/dst entity tables for task_name={task_name!r}. " + "Provide cached task tables with fkey metadata." + ) + + self.src_entity_table = str(src_table) + self.dst_entity_table = str(dst_table) + + # Validate that entity tables exist without materializing the full DB. + if dataset.cache_dir is None: + raise RuntimeError("TGBNextLinkPredTask requires dataset.cache_dir to validate entity tables.") + db_dir = Path(dataset.cache_dir) / "db" + src_path = db_dir / f"{self.src_entity_table}.parquet" + dst_path = db_dir / f"{self.dst_entity_table}.parquet" + if not src_path.exists(): + raise ValueError(f"src_entity_table='{self.src_entity_table}' not found in dataset db at {src_path}.") + if not dst_path.exists(): + raise ValueError(f"dst_entity_table='{self.dst_entity_table}' not found in dataset db at {dst_path}.") + super().__init__(dataset, cache_dir=cache_dir) - def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Table: - le = db.table_dict["label_events"].df - items = db.table_dict["label_event_items"].df - - rows: list[pd.DataFrame] = [] - for t in timestamps: - start = pd.Timestamp(t) - end = start + self.timedelta - mask = (le["label_ts"] > start) & (le["label_ts"] <= end) - if not mask.any(): - continue - - le_win = le.loc[mask, ["label_event_id", "src_id", "label_ts"]].sort_values("label_ts", kind="mergesort") - # "next" semantics: pick the first label_event per node in the window. - le_first = le_win.drop_duplicates(subset=["src_id"], keep="first")[["label_event_id", "src_id"]] - joined = le_first.merge(items[["label_event_id", "label_id"]], on="label_event_id", how="inner") - grouped = ( - joined.groupby("src_id", sort=False)["label_id"] - .agg(_unique_list) - .reset_index() - .rename(columns={"label_id": self.dst_entity_col, "src_id": self.src_entity_col}) + def make_table(self, db, timestamps): # pragma: no cover + raise RuntimeError( + "TGBNextLinkPredTask expects precomputed task tables " + "(train/val/test.parquet)." + ) + + +@dataclass(frozen=True) +class TGBNodePropSpec: + r"""Specification for a TGB node property prediction task.""" + + label_events_table: str = "label_events" + label_items_table: str = "label_event_items" + labels_table: str = "labels" + node_col: str = "src_id" + label_event_id_col: str = "label_event_id" + label_id_col: str = "label_id" + label_weight_col: str = "label_weight" + time_col: str = "label_ts" + + +def _tgb_nodeprop_ndcg_at_k( + *, + topk_label_ids: np.ndarray, # [N, K] + topk_scores: np.ndarray, # [N, K] (unused except for ordering) + true_label_ids: list[np.ndarray], + true_label_w: list[np.ndarray], + k: int, +) -> float: + """Compute NDCG@k with the same (exponential) gain convention as sklearn.ndcg_score. + + TGB's nodeproppred evaluator uses `sklearn.metrics.ndcg_score(y_true, y_pred, k=10)`. + For sparse ground truth, we compute the equivalent quantity using only the top-k + predicted labels and the top-k true labels. + """ + k = int(k) + if topk_label_ids.ndim != 2 or topk_scores.ndim != 2: + raise ValueError("topk_label_ids/topk_scores must be 2D arrays (N,K).") + if topk_label_ids.shape != topk_scores.shape: + raise ValueError("topk_label_ids and topk_scores must have the same shape.") + + n = int(topk_label_ids.shape[0]) + if len(true_label_ids) != n or len(true_label_w) != n: + raise ValueError("true_label_ids/true_label_w must be lists of length N.") + + discounts = 1.0 / np.log2(np.arange(k, dtype=np.float64) + 2.0) + + ndcgs: list[float] = [] + for i in range(n): + ids = np.asarray(true_label_ids[i], dtype=np.int64) + rel = np.asarray(true_label_w[i], dtype=np.float64) + if ids.size == 0: + ndcgs.append(0.0) + continue + + # Ideal DCG: sort true relevances and take top-k. + rel_sorted = np.sort(rel)[::-1] + rel_top = rel_sorted[:k] + idcg = ((np.exp2(rel_top) - 1.0) * discounts[: rel_top.shape[0]]).sum() + if idcg <= 0: + ndcgs.append(0.0) + continue + + # DCG: lookup true relevance for predicted top-k labels. + rel_map = {int(l): float(w) for l, w in zip(ids.tolist(), rel.tolist())} + pred_ids = topk_label_ids[i, :k] + gains = np.fromiter((np.exp2(rel_map.get(int(l), 0.0)) - 1.0 for l in pred_ids.tolist()), dtype=np.float64) + dcg = (gains * discounts[: gains.shape[0]]).sum() + ndcgs.append(float(dcg / idcg)) + + return float(np.mean(ndcgs)) if ndcgs else 0.0 + + +class TGBNodePropNDCGTask(BaseTask): + r"""TGB-style node property prediction with official NDCG@10 evaluation. + + This task assumes the `tgbn-*` export schema includes: + - `label_events(label_event_id, src_id, label_ts)` + - `label_event_items(label_event_id, label_id, label_weight)` + - `labels(label_id, ...)` + + Important: Exact parity with TGB requires scoring *all* labels (or at least + retrieving the top-k labels under that full scoring) for each label event. + """ + + task_type = TaskType.MULTILABEL_CLASSIFICATION + timedelta = pd.Timedelta(seconds=1) + metrics = [] + num_eval_timestamps = 1 + + def __init__( + self, + dataset: Dataset, + *, + spec: TGBNodePropSpec = TGBNodePropSpec(), + k: int = 10, + cache_dir: Optional[str] = None, + ) -> None: + self.spec = spec + self.k = int(k) + + db = dataset.get_db() + if spec.label_events_table not in db.table_dict: + raise ValueError(f"Missing table {spec.label_events_table} in dataset db.") + if spec.label_items_table not in db.table_dict: + raise ValueError(f"Missing table {spec.label_items_table} in dataset db.") + if spec.labels_table not in db.table_dict: + raise ValueError(f"Missing table {spec.labels_table} in dataset db.") + + label_events = db.table_dict[spec.label_events_table] + self.entity_table = label_events.fkey_col_to_pkey_table.get(spec.node_col) + if self.entity_table is None: + raise ValueError( + f"Expected {spec.label_events_table}.{spec.node_col} to be a foreign key. " + f"Got {label_events.fkey_col_to_pkey_table}." ) - grouped[self.time_col] = start - rows.append(grouped[[self.time_col, self.src_entity_col, self.dst_entity_col]]) - if rows: - out = pd.concat(rows, axis=0, ignore_index=True) + self.time_col = spec.time_col + self.entity_col = spec.node_col + self.label_event_id_col = spec.label_event_id_col + + super().__init__(dataset, cache_dir=cache_dir) + + def make_table(self, db, timestamps): # pragma: no cover + raise RuntimeError( + "TGBNodePropNDCGTask expects precomputed task tables (train/val/test.parquet) " + "and overrides _get_table()." + ) + + def filter_dangling_entities(self, table: Table) -> Table: + # Drop rows with invalid node ids. + db = self.dataset.get_db() + num_nodes = len(db.table_dict[self.entity_table]) + bad = table.df[self.entity_col] >= num_nodes + if bad.any(): + table.df = table.df[~bad] + return table + + @lru_cache(maxsize=1) + def _label_csr(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Return CSR arrays over label_event_id -> (label_id, label_weight).""" + db = self.dataset.get_db(upto_test_timestamp=False) + items = db.table_dict[self.spec.label_items_table].df[ + [self.spec.label_event_id_col, self.spec.label_id_col, self.spec.label_weight_col] + ].copy() + + event_ids = items[self.spec.label_event_id_col].astype("int64").to_numpy() + order = np.argsort(event_ids, kind="mergesort") + event_ids = event_ids[order] + label_ids = items[self.spec.label_id_col].astype("int64").to_numpy()[order] + label_w = items[self.spec.label_weight_col].astype("float64").to_numpy()[order] + + num_events = len(db.table_dict[self.spec.label_events_table]) + counts = np.bincount(event_ids, minlength=num_events).astype(np.int64, copy=False) + indptr = np.empty(num_events + 1, dtype=np.int64) + indptr[0] = 0 + np.cumsum(counts, out=indptr[1:]) + return indptr, label_ids, label_w + + def _truth_for_events(self, label_event_ids: np.ndarray) -> tuple[list[np.ndarray], list[np.ndarray]]: + indptr, label_ids, label_w = self._label_csr() + true_ids: list[np.ndarray] = [] + true_w: list[np.ndarray] = [] + for e in label_event_ids.astype(np.int64, copy=False).tolist(): + e = int(e) + start = int(indptr[e]) + end = int(indptr[e + 1]) + true_ids.append(np.asarray(label_ids[start:end], dtype=np.int64)) + true_w.append(np.asarray(label_w[start:end], dtype=np.float64)) + return true_ids, true_w + + def evaluate( + self, + pred: Any, + target_table: Optional[Table] = None, + metrics: Optional[list] = None, + ) -> dict[str, float]: + if target_table is None: + target_table = self.get_table("test", mask_input_cols=False) + + if isinstance(pred, dict): + y_pred = np.asarray(pred.get("y_pred")) else: - out = pd.DataFrame({self.time_col: [], self.src_entity_col: [], self.dst_entity_col: []}) - - return Table( - df=out, - fkey_col_to_pkey_table={ - self.src_entity_col: self.src_entity_table, - self.dst_entity_col: self.dst_entity_table, - }, - pkey_col=None, - time_col=self.time_col, + y_pred = np.asarray(pred) + + if y_pred.ndim != 2: + raise ValueError("Expected predictions with shape (N, num_labels).") + if y_pred.shape[0] != len(target_table): + raise ValueError(f"Prediction rows {y_pred.shape[0]} != target rows {len(target_table)}.") + + # Extract top-k label ids per row. + k = int(self.k) + topk = np.argpartition(-y_pred, kth=min(k - 1, y_pred.shape[1] - 1), axis=1)[:, :k] + topk_scores = np.take_along_axis(y_pred, topk, axis=1) + order = np.argsort(-topk_scores, axis=1, kind="mergesort") + topk = np.take_along_axis(topk, order, axis=1) + topk_scores = np.take_along_axis(topk_scores, order, axis=1) + + label_event_ids = target_table.df[self.label_event_id_col].astype("int64").to_numpy() + true_ids, true_w = self._truth_for_events(label_event_ids) + ndcg = _tgb_nodeprop_ndcg_at_k( + topk_label_ids=topk, + topk_scores=topk_scores, + true_label_ids=true_ids, + true_label_w=true_w, + k=k, ) + return {f"ndcg@{k}": float(ndcg)} From dcd63c473f85f295dea92567f21c39e94dc6fe14 Mon Sep 17 00:00:00 2001 From: pc0618 Date: Thu, 29 Jan 2026 07:32:21 +0000 Subject: [PATCH 4/8] Drop TGB *-next task registrations --- relbench/tasks/__init__.py | 65 -------------------------------------- 1 file changed, 65 deletions(-) diff --git a/relbench/tasks/__init__.py b/relbench/tasks/__init__.py index efd69910..1b35b742 100644 --- a/relbench/tasks/__init__.py +++ b/relbench/tasks/__init__.py @@ -554,11 +554,6 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: # TGB official evaluation for link prediction is one-vs-many with pre-generated # negatives and reports MRR/Hits@k. For heterogeneous datasets, this is done # per edge type; we therefore expose one task per `events_edge_type_*` table. -# -# In addition, we expose exporter-defined "next interaction" recommendation -# tasks (`*-next`) to validate RelBench recommendation baselines on the same -# exported schema. These tasks use standard RelBench top-k metrics -# (precision/recall/MAP). # tgbl-* (link prediction; single event table) for dataset_name in [ @@ -570,13 +565,6 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: "rel-tgb-tgbl-comment", "rel-tgb-tgbl-flight", ]: - register_task( - dataset_name, - "src-dst-next", - tgb.TGBNextLinkPredTask, - tgb_task_name="src-dst-next", - eval_k=10, - ) register_task( dataset_name, "src-dst-mrr", @@ -597,56 +585,11 @@ def _register_thgl_edge_type_tasks(dataset_name: str, edge_types: list[int]) -> ) -def _register_thgl_next_tasks(dataset_name: str, next_tasks: list[str]) -> None: - for name in next_tasks: - register_task( - dataset_name, - name, - tgb.TGBNextLinkPredTask, - tgb_task_name=name, - eval_k=10, - ) - - _register_thgl_edge_type_tasks("rel-tgb-thgl-software", list(range(14))) _register_thgl_edge_type_tasks("rel-tgb-thgl-github", list(range(14))) _register_thgl_edge_type_tasks("rel-tgb-thgl-forum", [0, 1]) _register_thgl_edge_type_tasks("rel-tgb-thgl-myket", [0, 1]) -_register_thgl_next_tasks( - "rel-tgb-thgl-software", - [ - "type0-type1-next", - "type0-type3-next", - "type1-type2-next", - "type3-type2-next", - ], -) -_register_thgl_next_tasks( - "rel-tgb-thgl-github", - [ - "type0-type1-next", - "type1-type1-next", - "type2-type0-next", - "type2-type1-next", - "type2-type3-next", - "type3-type1-next", - ], -) -_register_thgl_next_tasks( - "rel-tgb-thgl-forum", - [ - "type0-type0-next", - "type0-type1-next", - ], -) -_register_thgl_next_tasks( - "rel-tgb-thgl-myket", - [ - "type0-type1-next", - ], -) - # tgbn-* (node property prediction; NDCG@10 over label distributions) for dataset_name in [ "rel-tgb-tgbn-trade", @@ -654,14 +597,6 @@ def _register_thgl_next_tasks(dataset_name: str, next_tasks: list[str]) -> None: "rel-tgb-tgbn-reddit", "rel-tgb-tgbn-token", ]: - register_task( - dataset_name, - "node-label-next", - tgb.TGBNextLinkPredTask, - tgb_task_name="node-label-next", - dst_entity_col="label_id", - eval_k=10, - ) register_task( dataset_name, "node-label-ndcg", From 1dbb2bae7d49cd0a6c5c701507dd98f19f28649f Mon Sep 17 00:00:00 2001 From: pc0618 Date: Thu, 29 Jan 2026 07:48:03 +0000 Subject: [PATCH 5/8] Trim comments in TGB modules --- relbench/datasets/__init__.py | 4 +-- relbench/datasets/tgb.py | 28 +---------------- relbench/modeling/utils.py | 6 +--- relbench/tasks/__init__.py | 9 +----- relbench/tasks/tgb.py | 57 ++--------------------------------- 5 files changed, 7 insertions(+), 97 deletions(-) diff --git a/relbench/datasets/__init__.py b/relbench/datasets/__init__.py index f8c142cb..2fde19e1 100644 --- a/relbench/datasets/__init__.py +++ b/relbench/datasets/__init__.py @@ -159,9 +159,7 @@ def get_dataset(name: str, download=True) -> Dataset: register_dataset("dbinfer-stackexchange", dbinfer.DBInferStackExchangeDataset) register_dataset("dbinfer-outbrain-small", dbinfer.DBInferOutbrainSmallDataset) -# Community dataset family: Temporal Graph Benchmark (TGB) -# Names follow the convention `rel--` where -# dataset_name = "tgb" and qualifier is the original TGB dataset id. +# Temporal Graph Benchmark (TGB) register_dataset("rel-tgb-tgbl-wiki", tgb.TGBDataset, tgb_name="tgbl-wiki") register_dataset("rel-tgb-tgbl-wiki-v2", tgb.TGBDataset, tgb_name="tgbl-wiki-v2") register_dataset("rel-tgb-tgbl-review", tgb.TGBDataset, tgb_name="tgbl-review") diff --git a/relbench/datasets/tgb.py b/relbench/datasets/tgb.py index b62c95aa..e8ba5b85 100644 --- a/relbench/datasets/tgb.py +++ b/relbench/datasets/tgb.py @@ -16,10 +16,6 @@ class TGBCutoffs: _TGB_CUTOFFS: dict[str, TGBCutoffs] = { - # Dynamic link property prediction (tgbl-*) - # NOTE: These cutoffs match the pre-built RelBench exports used in our - # conversion pipeline. They are expected to be consistent with the - # official TGB temporal split (70/15/15). "tgbl-wiki": TGBCutoffs(val_timestamp_s=1862653, test_timestamp_s=2218300), "tgbl-wiki-v2": TGBCutoffs(val_timestamp_s=1862653, test_timestamp_s=2218300), "tgbl-review": TGBCutoffs(val_timestamp_s=1464912000, test_timestamp_s=1488844800), @@ -27,14 +23,10 @@ class TGBCutoffs: "tgbl-coin": TGBCutoffs(val_timestamp_s=1662096249, test_timestamp_s=1664482319), "tgbl-comment": TGBCutoffs(val_timestamp_s=1282869285, test_timestamp_s=1288838725), "tgbl-flight": TGBCutoffs(val_timestamp_s=1638162000, test_timestamp_s=1653796800), - # Temporal heterogeneous link prediction (thgl-*) "thgl-software": TGBCutoffs(val_timestamp_s=1706003880, test_timestamp_s=1706315669), "thgl-forum": TGBCutoffs(val_timestamp_s=1390426563, test_timestamp_s=1390838358), "thgl-github": TGBCutoffs(val_timestamp_s=1711075987, test_timestamp_s=1711482874), "thgl-myket": TGBCutoffs(val_timestamp_s=1603724860, test_timestamp_s=1606341312), - # Dynamic node property prediction (tgbn-*) - # Nodeprop timestamps may be stored as years in the raw sources; our export - # normalizes timestamps to UNIX seconds. "tgbn-trade": TGBCutoffs(val_timestamp_s=1262304000, test_timestamp_s=1388534400), "tgbn-genre": TGBCutoffs(val_timestamp_s=1216427762, test_timestamp_s=1230448684), "tgbn-reddit": TGBCutoffs(val_timestamp_s=1279485233, test_timestamp_s=1286653871), @@ -43,22 +35,7 @@ class TGBCutoffs: class TGBDataset(Dataset): - r"""Community dataset family: Temporal Graph Benchmark (TGB) exports. - - This dataset class expects a pre-built RelBench database at: - `cache_dir/db/*.parquet`. - - The recommended workflow for contributions (see CONTRIBUTING.md): - 1) Materialize `db/` (Parquet) for each dataset. - 2) Zip it as `db.zip` and publish it. - 3) Add the sha256 hash to `relbench/datasets/hashes.json`. - - Notes: - - Primary/foreign keys are stored as int64 (DBML: `bigint`) to avoid 32-bit - overflow and to match PyArrow's default integer type. - - We store only the relational schema + timestamps; "splits" are derived - from cutoffs (val/test timestamps). - """ + r"""Temporal Graph Benchmark (TGB) datasets exported to RelBench format.""" url = "https://tgb.complexdatalab.com/" @@ -70,8 +47,6 @@ def __init__(self, *, tgb_name: str, cache_dir: Optional[str] = None) -> None: self.tgb_name = str(tgb_name) cutoffs = _TGB_CUTOFFS[self.tgb_name] - # TGB exports store timestamps as timezone-aware UTC (timestamp[ns, UTC]). - # Keep dataset cutoffs in UTC as well to avoid tz-mismatch in task builders. self.val_timestamp = pd.to_datetime(int(cutoffs.val_timestamp_s), unit="s", utc=True) self.test_timestamp = pd.to_datetime(int(cutoffs.test_timestamp_s), unit="s", utc=True) @@ -90,4 +65,3 @@ def make_db(self) -> Database: "This dataset is distributed as a pre-built RelBench database (db.zip). " "Please run `download_dataset(...)` or place `db/*.parquet` in the cache directory." ) - diff --git a/relbench/modeling/utils.py b/relbench/modeling/utils.py index 070d1b81..95b0a507 100644 --- a/relbench/modeling/utils.py +++ b/relbench/modeling/utils.py @@ -9,21 +9,17 @@ def to_unix_time(ser: pd.Series) -> np.ndarray: - r"""Converts a :class:`pandas.Timestamp` series to UNIX timestamp (in seconds).""" - # Accept tz-aware timestamps and normalize to UTC. - # Many external parquet exports use `datetime64[ns, UTC]`. + r"""Convert a timestamp-like series to UNIX seconds.""" if pd.api.types.is_datetime64_any_dtype(ser.dtype) or pd.api.types.is_datetime64tz_dtype(ser.dtype): ts = pd.to_datetime(ser, utc=True) unix_ns = ts.astype("int64").to_numpy(copy=False) return (unix_ns // 1_000_000_000).astype(np.int64, copy=False) - # Allow integer/float timestamps that are already in seconds. if pd.api.types.is_integer_dtype(ser.dtype): return ser.astype("int64").to_numpy(copy=False) if pd.api.types.is_float_dtype(ser.dtype): return ser.astype("int64").to_numpy(copy=False) - # Fallback: parse strings/objects as datetimes. ts = pd.to_datetime(ser, utc=True) unix_ns = ts.astype("int64").to_numpy(copy=False) return (unix_ns // 1_000_000_000).astype(np.int64, copy=False) diff --git a/relbench/tasks/__init__.py b/relbench/tasks/__init__.py index 1b35b742..80451efc 100644 --- a/relbench/tasks/__init__.py +++ b/relbench/tasks/__init__.py @@ -549,13 +549,7 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: register_task("dbinfer-stackexchange", "upvote", dbinfer.StackExchangeUpvoteTask) register_task("dbinfer-outbrain-small", "ctr", dbinfer.OutbrainCTRTask) -# Community task family: Temporal Graph Benchmark (TGB) -# -# TGB official evaluation for link prediction is one-vs-many with pre-generated -# negatives and reports MRR/Hits@k. For heterogeneous datasets, this is done -# per edge type; we therefore expose one task per `events_edge_type_*` table. - -# tgbl-* (link prediction; single event table) +# Temporal Graph Benchmark (TGB) for dataset_name in [ "rel-tgb-tgbl-wiki", "rel-tgb-tgbl-wiki-v2", @@ -590,7 +584,6 @@ def _register_thgl_edge_type_tasks(dataset_name: str, edge_types: list[int]) -> _register_thgl_edge_type_tasks("rel-tgb-thgl-forum", [0, 1]) _register_thgl_edge_type_tasks("rel-tgb-thgl-myket", [0, 1]) -# tgbn-* (node property prediction; NDCG@10 over label distributions) for dataset_name in [ "rel-tgb-tgbn-trade", "rel-tgb-tgbn-genre", diff --git a/relbench/tasks/tgb.py b/relbench/tasks/tgb.py index ba370db1..158b09ce 100644 --- a/relbench/tasks/tgb.py +++ b/relbench/tasks/tgb.py @@ -20,7 +20,6 @@ def _to_unix_seconds(ts: pd.Series) -> np.ndarray: ts = pd.to_datetime(ts, utc=True) - # Timestamp[ns] -> seconds return (ts.astype("int64").to_numpy(copy=False) // 1_000_000_000).astype(np.int64, copy=False) @@ -30,13 +29,7 @@ def _tgb_eval_hits_and_mrr( *, k_value: int, ) -> dict[str, float]: - r"""Match TGB's link prediction evaluator for one-vs-many ranking. - - This mirrors `tgb.linkproppred.evaluate.Evaluator._eval_hits_and_mrr`: - - optimistic rank counts negatives strictly greater than positive - - pessimistic rank counts negatives greater-or-equal to positive - - final rank is the average of the two (plus 1) - """ + r"""Compute one-vs-many Hits@k and MRR with tie-aware ranking.""" y_pred_pos = np.asarray(y_pred_pos).reshape(-1, 1) y_pred_neg = np.asarray(y_pred_neg) @@ -62,28 +55,13 @@ class TGBLinkPredSpec: event_table: str src_col: str = "src_id" dst_col: str = "dst_id" - # Matches the exporter convention (`*_exports/**/db/events*.parquet`). time_col: str = "event_ts" class TGBOneVsManyLinkPredTask(BaseTask): - r"""TGB-style link prediction task with official one-vs-many MRR/Hits@k evaluation. - - Important: Exact parity with TGB requires: - - task tables (`train/val/test.parquet`) listing the *exact* positive edges - for each split (matching TGB masks, not just timestamp cutoffs), and - - the TGB-provided pre-generated negative samples placed at: - `/negatives/val_ns.pkl` and `.../test_ns.pkl`. - - For heterogeneous `thgl-*` datasets, TGB negative samples are keyed by - `(timestamp_s, src_global_id, edge_type_id)`, which requires the dataset to - ship mapping files that relate RelBench's per-type local node ids back to - the original TGB global node ids. - """ + r"""TGB link prediction with official one-vs-many MRR/Hits@k evaluation.""" task_type = TaskType.LINK_PREDICTION - # BaseTask enforces `dataset.test_timestamp - dataset.val_timestamp >= timedelta`. - # We override table construction, so we just need a non-negative timedelta. timedelta = pd.Timedelta(seconds=1) metrics = [] num_eval_timestamps = 1 @@ -103,7 +81,6 @@ def __init__( self.src_entity_col = spec.src_col self.dst_entity_col = spec.dst_col - # Infer FK target tables + (optional) edge_type id from the exported database. db = dataset.get_db() if spec.event_table not in db.table_dict: raise ValueError( @@ -131,7 +108,6 @@ def make_table(self, db, timestamps): # pragma: no cover ) def filter_dangling_entities(self, table: Table) -> Table: - # Keep parity with TGB negatives by dropping rows with invalid ids. if self.src_entity_table: table.df = table.df[table.df[self.src_entity_col] < len(self.dataset.get_db().table_dict[self.src_entity_table])] if self.dst_entity_table: @@ -140,7 +116,6 @@ def filter_dangling_entities(self, table: Table) -> Table: return table def _get_table(self, split: str) -> Table: - # Load cached task tables rather than generating by timestamps. if split not in ["train", "val", "test"]: raise ValueError(f"Unknown split '{split}'.") table_path = f"{self.cache_dir}/{split}.parquet" @@ -160,8 +135,6 @@ def _negatives_path(self, split: str) -> Path: raise RuntimeError("Dataset has no cache_dir; cannot locate TGB negatives.") return Path(self.dataset.cache_dir) / "negatives" / f"{split}_ns.pkl" - # Large datasets ship multi-GB negative dictionaries. Keep at most one split - # in memory at a time to reduce peak RAM usage. @lru_cache(maxsize=1) def _load_negatives(self, split: str) -> dict[Any, Any]: path = self._negatives_path(split) @@ -215,9 +188,6 @@ def _node_type_id_from_table(self, table_name: str) -> Optional[int]: return int(m.group(1)) if m else None def _bipartite_offset(self) -> Optional[int]: - # Export convention for bipartite tgbl-wiki*: src ids are 0..num_src-1, - # dst ids are 0..num_dst-1 in RelBench, but TGB negatives use a single - # global id space with dst shifted by `num_src`. if self.src_entity_table == "src_nodes" and self.dst_entity_table == "dst_nodes": return len(self.dataset.get_db().table_dict["src_nodes"]) return None @@ -225,7 +195,6 @@ def _bipartite_offset(self) -> Optional[int]: def _src_local_to_global(self, src_local: np.ndarray) -> np.ndarray: src_type = self._node_type_id_from_table(self.src_entity_table) if src_type is None: - # homogeneous/bipartite case: local ids are global ids return src_local.astype(np.int64, copy=False) globals_ = self._load_local_to_global(src_type) return globals_[src_local.astype(np.int64, copy=False)] @@ -267,7 +236,6 @@ def get_negative_dsts_local(self, *, split: str, table: Optional[Table] = None) negs_local: list[np.ndarray] = [] if self.edge_type_id is None: - # tgbl-* negatives are keyed by (src, dst, t) dst_local = df[self.dst_entity_col].astype("int64").to_numpy() offset = self._bipartite_offset() if offset is None: @@ -278,7 +246,6 @@ def get_negative_dsts_local(self, *, split: str, table: Optional[Table] = None) negs_g = np.asarray(neg_dict[(s, d, t)], dtype=np.int64) negs_local.append(self._dst_global_to_local(negs_g)) else: - # thgl-* negatives are keyed by (t, src, edge_type) et = int(self.edge_type_id) for t, s in zip(ts_s.tolist(), src_global.tolist()): negs_g = np.asarray(neg_dict[(t, s, et)], dtype=np.int64) @@ -338,8 +305,6 @@ class TGBNextLinkPredTask(RecommendationTask): """ task_type = TaskType.LINK_PREDICTION - # BaseTask enforces `dataset.test_timestamp - dataset.val_timestamp >= timedelta`. - # We only need a non-negative timedelta since we rely on cached tables. timedelta = pd.Timedelta(seconds=1) num_eval_timestamps = 1 metrics = [link_prediction_precision, link_prediction_recall, link_prediction_map] @@ -357,12 +322,10 @@ def __init__( ) -> None: self.eval_k = int(eval_k) - # Column conventions of exported task tables. self.time_col = str(time_col) self.src_entity_col = str(src_entity_col) self.dst_entity_col = str(dst_entity_col) - # Infer src/dst entity tables from the cached task table metadata if present. src_table = None dst_table = None if cache_dir is not None: @@ -376,9 +339,6 @@ def __init__( break if src_table is None or dst_table is None: - # Fallback to inferring from task naming conventions: - # - `src-dst-next` for tgbl-* exports (bipartite or homogeneous) - # - `typeX-typeY-next` for thgl-* exports (per node type) task_name = str(tgb_task_name) m = re.fullmatch(r"type(\d+)-type(\d+)-next", task_name) if m is not None: @@ -393,7 +353,6 @@ def __init__( src_table = "nodes" dst_table = "nodes" elif task_name == "node-label-next": - # tgbn-* exports: predict labels for a node. src_table = "nodes" dst_table = "labels" else: @@ -405,7 +364,6 @@ def __init__( self.src_entity_table = str(src_table) self.dst_entity_table = str(dst_table) - # Validate that entity tables exist without materializing the full DB. if dataset.cache_dir is None: raise RuntimeError("TGBNextLinkPredTask requires dataset.cache_dir to validate entity tables.") db_dir = Path(dataset.cache_dir) / "db" @@ -447,12 +405,7 @@ def _tgb_nodeprop_ndcg_at_k( true_label_w: list[np.ndarray], k: int, ) -> float: - """Compute NDCG@k with the same (exponential) gain convention as sklearn.ndcg_score. - - TGB's nodeproppred evaluator uses `sklearn.metrics.ndcg_score(y_true, y_pred, k=10)`. - For sparse ground truth, we compute the equivalent quantity using only the top-k - predicted labels and the top-k true labels. - """ + """Compute NDCG@k with the same gain convention as sklearn.ndcg_score.""" k = int(k) if topk_label_ids.ndim != 2 or topk_scores.ndim != 2: raise ValueError("topk_label_ids/topk_scores must be 2D arrays (N,K).") @@ -473,7 +426,6 @@ def _tgb_nodeprop_ndcg_at_k( ndcgs.append(0.0) continue - # Ideal DCG: sort true relevances and take top-k. rel_sorted = np.sort(rel)[::-1] rel_top = rel_sorted[:k] idcg = ((np.exp2(rel_top) - 1.0) * discounts[: rel_top.shape[0]]).sum() @@ -481,7 +433,6 @@ def _tgb_nodeprop_ndcg_at_k( ndcgs.append(0.0) continue - # DCG: lookup true relevance for predicted top-k labels. rel_map = {int(l): float(w) for l, w in zip(ids.tolist(), rel.tolist())} pred_ids = topk_label_ids[i, :k] gains = np.fromiter((np.exp2(rel_map.get(int(l), 0.0)) - 1.0 for l in pred_ids.tolist()), dtype=np.float64) @@ -548,7 +499,6 @@ def make_table(self, db, timestamps): # pragma: no cover ) def filter_dangling_entities(self, table: Table) -> Table: - # Drop rows with invalid node ids. db = self.dataset.get_db() num_nodes = len(db.table_dict[self.entity_table]) bad = table.df[self.entity_col] >= num_nodes @@ -608,7 +558,6 @@ def evaluate( if y_pred.shape[0] != len(target_table): raise ValueError(f"Prediction rows {y_pred.shape[0]} != target rows {len(target_table)}.") - # Extract top-k label ids per row. k = int(self.k) topk = np.argpartition(-y_pred, kth=min(k - 1, y_pred.shape[1] - 1), axis=1)[:, :k] topk_scores = np.take_along_axis(y_pred, topk, axis=1) From 80f57ed86a268a92ce731a05b1a3981c7885eaef Mon Sep 17 00:00:00 2001 From: pc0618 Date: Tue, 3 Feb 2026 21:26:01 +0000 Subject: [PATCH 6/8] Document TGB community datasets --- README.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/README.md b/README.md index 26bb17fa..971c5bc3 100644 --- a/README.md +++ b/README.md @@ -157,6 +157,22 @@ from relbench.datasets import get_dataset from relbench.tasks import get_task ``` +## Community Datasets: Temporal Graph Benchmark (TGB) + +RelBench supports artifact-backed community datasets exported from the Temporal Graph Benchmark (TGB): +- Dynamic link prediction: `rel-tgb-tgbl-*` (task: `src-dst-mrr`) +- Temporal heterogeneous link prediction: `rel-tgb-thgl-*` (tasks: `edge-type--mrr`) +- Dynamic node property prediction: `rel-tgb-tgbn-*` (task: `node-label-ndcg`) + +These datasets/tasks are distributed as pre-built `db.zip` and task-table zip artifacts (not raw-source processing in RelBench). + +Example: +```python +dataset = get_dataset("rel-tgb-tgbl-wiki-v2", download=True) +task = get_task("rel-tgb-tgbl-wiki-v2", "src-dst-mrr", download=True) +val_table = task.get_table("val", mask_input_cols=False) +``` + Get a dataset, e.g., `rel-amazon`: ```python dataset: Dataset = get_dataset("rel-amazon", download=True) From 1b611f62445d90219488b6b54f939603cec6d5dd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 3 Feb 2026 23:50:48 +0000 Subject: [PATCH 7/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- relbench/datasets/tgb.py | 16 +++++-- relbench/modeling/utils.py | 4 +- relbench/tasks/tgb.py | 96 +++++++++++++++++++++++++++++--------- 3 files changed, 88 insertions(+), 28 deletions(-) diff --git a/relbench/datasets/tgb.py b/relbench/datasets/tgb.py index e8ba5b85..ae9f914a 100644 --- a/relbench/datasets/tgb.py +++ b/relbench/datasets/tgb.py @@ -19,11 +19,15 @@ class TGBCutoffs: "tgbl-wiki": TGBCutoffs(val_timestamp_s=1862653, test_timestamp_s=2218300), "tgbl-wiki-v2": TGBCutoffs(val_timestamp_s=1862653, test_timestamp_s=2218300), "tgbl-review": TGBCutoffs(val_timestamp_s=1464912000, test_timestamp_s=1488844800), - "tgbl-review-v2": TGBCutoffs(val_timestamp_s=1464912000, test_timestamp_s=1488844800), + "tgbl-review-v2": TGBCutoffs( + val_timestamp_s=1464912000, test_timestamp_s=1488844800 + ), "tgbl-coin": TGBCutoffs(val_timestamp_s=1662096249, test_timestamp_s=1664482319), "tgbl-comment": TGBCutoffs(val_timestamp_s=1282869285, test_timestamp_s=1288838725), "tgbl-flight": TGBCutoffs(val_timestamp_s=1638162000, test_timestamp_s=1653796800), - "thgl-software": TGBCutoffs(val_timestamp_s=1706003880, test_timestamp_s=1706315669), + "thgl-software": TGBCutoffs( + val_timestamp_s=1706003880, test_timestamp_s=1706315669 + ), "thgl-forum": TGBCutoffs(val_timestamp_s=1390426563, test_timestamp_s=1390838358), "thgl-github": TGBCutoffs(val_timestamp_s=1711075987, test_timestamp_s=1711482874), "thgl-myket": TGBCutoffs(val_timestamp_s=1603724860, test_timestamp_s=1606341312), @@ -47,8 +51,12 @@ def __init__(self, *, tgb_name: str, cache_dir: Optional[str] = None) -> None: self.tgb_name = str(tgb_name) cutoffs = _TGB_CUTOFFS[self.tgb_name] - self.val_timestamp = pd.to_datetime(int(cutoffs.val_timestamp_s), unit="s", utc=True) - self.test_timestamp = pd.to_datetime(int(cutoffs.test_timestamp_s), unit="s", utc=True) + self.val_timestamp = pd.to_datetime( + int(cutoffs.val_timestamp_s), unit="s", utc=True + ) + self.test_timestamp = pd.to_datetime( + int(cutoffs.test_timestamp_s), unit="s", utc=True + ) super().__init__(cache_dir=cache_dir) diff --git a/relbench/modeling/utils.py b/relbench/modeling/utils.py index 95b0a507..f21f9297 100644 --- a/relbench/modeling/utils.py +++ b/relbench/modeling/utils.py @@ -10,7 +10,9 @@ def to_unix_time(ser: pd.Series) -> np.ndarray: r"""Convert a timestamp-like series to UNIX seconds.""" - if pd.api.types.is_datetime64_any_dtype(ser.dtype) or pd.api.types.is_datetime64tz_dtype(ser.dtype): + if pd.api.types.is_datetime64_any_dtype( + ser.dtype + ) or pd.api.types.is_datetime64tz_dtype(ser.dtype): ts = pd.to_datetime(ser, utc=True) unix_ns = ts.astype("int64").to_numpy(copy=False) return (unix_ns // 1_000_000_000).astype(np.int64, copy=False) diff --git a/relbench/tasks/tgb.py b/relbench/tasks/tgb.py index 158b09ce..240757d7 100644 --- a/relbench/tasks/tgb.py +++ b/relbench/tasks/tgb.py @@ -20,7 +20,9 @@ def _to_unix_seconds(ts: pd.Series) -> np.ndarray: ts = pd.to_datetime(ts, utc=True) - return (ts.astype("int64").to_numpy(copy=False) // 1_000_000_000).astype(np.int64, copy=False) + return (ts.astype("int64").to_numpy(copy=False) // 1_000_000_000).astype( + np.int64, copy=False + ) def _tgb_eval_hits_and_mrr( @@ -109,9 +111,15 @@ def make_table(self, db, timestamps): # pragma: no cover def filter_dangling_entities(self, table: Table) -> Table: if self.src_entity_table: - table.df = table.df[table.df[self.src_entity_col] < len(self.dataset.get_db().table_dict[self.src_entity_table])] + table.df = table.df[ + table.df[self.src_entity_col] + < len(self.dataset.get_db().table_dict[self.src_entity_table]) + ] if self.dst_entity_table: - table.df = table.df[table.df[self.dst_entity_col] < len(self.dataset.get_db().table_dict[self.dst_entity_table])] + table.df = table.df[ + table.df[self.dst_entity_col] + < len(self.dataset.get_db().table_dict[self.dst_entity_table]) + ] table.df = table.df.reset_index(drop=True) return table @@ -169,13 +177,19 @@ def _load_global_to_local(self) -> tuple[np.ndarray, np.ndarray]: ) node_type = np.load(node_type_path) local_id = np.load(local_id_path) - return node_type.astype(np.int64, copy=False), local_id.astype(np.int64, copy=False) + return node_type.astype(np.int64, copy=False), local_id.astype( + np.int64, copy=False + ) @lru_cache(maxsize=None) def _load_local_to_global(self, node_type_id: int) -> np.ndarray: if self.dataset.cache_dir is None: raise RuntimeError("Dataset has no cache_dir; cannot locate mapping files.") - p = Path(self.dataset.cache_dir) / "mappings" / f"globals_type_{int(node_type_id)}.npy" + p = ( + Path(self.dataset.cache_dir) + / "mappings" + / f"globals_type_{int(node_type_id)}.npy" + ) if not p.exists(): raise RuntimeError( f"Missing mapping file {p}. This is required to map local ids " @@ -188,7 +202,10 @@ def _node_type_id_from_table(self, table_name: str) -> Optional[int]: return int(m.group(1)) if m else None def _bipartite_offset(self) -> Optional[int]: - if self.src_entity_table == "src_nodes" and self.dst_entity_table == "dst_nodes": + if ( + self.src_entity_table == "src_nodes" + and self.dst_entity_table == "dst_nodes" + ): return len(self.dataset.get_db().table_dict["src_nodes"]) return None @@ -208,20 +225,26 @@ def _dst_global_to_local(self, dst_global: np.ndarray) -> np.ndarray: dst_global = dst_global.astype(np.int64, copy=False) out = dst_global - int(offset) if (out < 0).any(): - raise RuntimeError("Bipartite negatives contain ids outside destination range.") + raise RuntimeError( + "Bipartite negatives contain ids outside destination range." + ) return out.astype(np.int64, copy=False) node_type, local_id = self._load_global_to_local() dst_global = dst_global.astype(np.int64, copy=False) bad = node_type[dst_global] != dst_type if bad.any(): - raise RuntimeError("Negative samples contain destination nodes of unexpected type.") + raise RuntimeError( + "Negative samples contain destination nodes of unexpected type." + ) return local_id[dst_global].astype(np.int64, copy=False) - def get_negative_dsts_local(self, *, split: str, table: Optional[Table] = None) -> list[np.ndarray]: + def get_negative_dsts_local( + self, *, split: str, table: Optional[Table] = None + ) -> list[np.ndarray]: r"""Return negative destination ids (local to dst entity table) for each row. - This is intended to help users reproduce the TGB evaluation protocol, - i.e., score the true destination vs the provided negatives. + This is intended to help users reproduce the TGB evaluation protocol, i.e., + score the true destination vs the provided negatives. """ if split not in ["val", "test"]: raise ValueError("Negatives are only defined for val/test splits.") @@ -365,14 +388,20 @@ def __init__( self.dst_entity_table = str(dst_table) if dataset.cache_dir is None: - raise RuntimeError("TGBNextLinkPredTask requires dataset.cache_dir to validate entity tables.") + raise RuntimeError( + "TGBNextLinkPredTask requires dataset.cache_dir to validate entity tables." + ) db_dir = Path(dataset.cache_dir) / "db" src_path = db_dir / f"{self.src_entity_table}.parquet" dst_path = db_dir / f"{self.dst_entity_table}.parquet" if not src_path.exists(): - raise ValueError(f"src_entity_table='{self.src_entity_table}' not found in dataset db at {src_path}.") + raise ValueError( + f"src_entity_table='{self.src_entity_table}' not found in dataset db at {src_path}." + ) if not dst_path.exists(): - raise ValueError(f"dst_entity_table='{self.dst_entity_table}' not found in dataset db at {dst_path}.") + raise ValueError( + f"dst_entity_table='{self.dst_entity_table}' not found in dataset db at {dst_path}." + ) super().__init__(dataset, cache_dir=cache_dir) @@ -435,7 +464,10 @@ def _tgb_nodeprop_ndcg_at_k( rel_map = {int(l): float(w) for l, w in zip(ids.tolist(), rel.tolist())} pred_ids = topk_label_ids[i, :k] - gains = np.fromiter((np.exp2(rel_map.get(int(l), 0.0)) - 1.0 for l in pred_ids.tolist()), dtype=np.float64) + gains = np.fromiter( + (np.exp2(rel_map.get(int(l), 0.0)) - 1.0 for l in pred_ids.tolist()), + dtype=np.float64, + ) dcg = (gains * discounts[: gains.shape[0]]).sum() ndcgs.append(float(dcg / idcg)) @@ -510,9 +542,17 @@ def filter_dangling_entities(self, table: Table) -> Table: def _label_csr(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Return CSR arrays over label_event_id -> (label_id, label_weight).""" db = self.dataset.get_db(upto_test_timestamp=False) - items = db.table_dict[self.spec.label_items_table].df[ - [self.spec.label_event_id_col, self.spec.label_id_col, self.spec.label_weight_col] - ].copy() + items = ( + db.table_dict[self.spec.label_items_table] + .df[ + [ + self.spec.label_event_id_col, + self.spec.label_id_col, + self.spec.label_weight_col, + ] + ] + .copy() + ) event_ids = items[self.spec.label_event_id_col].astype("int64").to_numpy() order = np.argsort(event_ids, kind="mergesort") @@ -521,13 +561,17 @@ def _label_csr(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: label_w = items[self.spec.label_weight_col].astype("float64").to_numpy()[order] num_events = len(db.table_dict[self.spec.label_events_table]) - counts = np.bincount(event_ids, minlength=num_events).astype(np.int64, copy=False) + counts = np.bincount(event_ids, minlength=num_events).astype( + np.int64, copy=False + ) indptr = np.empty(num_events + 1, dtype=np.int64) indptr[0] = 0 np.cumsum(counts, out=indptr[1:]) return indptr, label_ids, label_w - def _truth_for_events(self, label_event_ids: np.ndarray) -> tuple[list[np.ndarray], list[np.ndarray]]: + def _truth_for_events( + self, label_event_ids: np.ndarray + ) -> tuple[list[np.ndarray], list[np.ndarray]]: indptr, label_ids, label_w = self._label_csr() true_ids: list[np.ndarray] = [] true_w: list[np.ndarray] = [] @@ -556,16 +600,22 @@ def evaluate( if y_pred.ndim != 2: raise ValueError("Expected predictions with shape (N, num_labels).") if y_pred.shape[0] != len(target_table): - raise ValueError(f"Prediction rows {y_pred.shape[0]} != target rows {len(target_table)}.") + raise ValueError( + f"Prediction rows {y_pred.shape[0]} != target rows {len(target_table)}." + ) k = int(self.k) - topk = np.argpartition(-y_pred, kth=min(k - 1, y_pred.shape[1] - 1), axis=1)[:, :k] + topk = np.argpartition(-y_pred, kth=min(k - 1, y_pred.shape[1] - 1), axis=1)[ + :, :k + ] topk_scores = np.take_along_axis(y_pred, topk, axis=1) order = np.argsort(-topk_scores, axis=1, kind="mergesort") topk = np.take_along_axis(topk, order, axis=1) topk_scores = np.take_along_axis(topk_scores, order, axis=1) - label_event_ids = target_table.df[self.label_event_id_col].astype("int64").to_numpy() + label_event_ids = ( + target_table.df[self.label_event_id_col].astype("int64").to_numpy() + ) true_ids, true_w = self._truth_for_events(label_event_ids) ndcg = _tgb_nodeprop_ndcg_at_k( topk_label_ids=topk, From 23b161c4ce782e81d6a3d299959619a17d7a3ac9 Mon Sep 17 00:00:00 2001 From: pc0618 Date: Tue, 24 Feb 2026 04:37:16 +0000 Subject: [PATCH 8/8] Rename TGB slugs; drop README section --- README.md | 16 ------- relbench/datasets/__init__.py | 34 +++++++------- relbench/datasets/hashes.json | 30 ++++++------ relbench/tasks/__init__.py | 30 ++++++------ relbench/tasks/hashes.json | 86 +++++++++++++++++------------------ 5 files changed, 90 insertions(+), 106 deletions(-) diff --git a/README.md b/README.md index 971c5bc3..26bb17fa 100644 --- a/README.md +++ b/README.md @@ -157,22 +157,6 @@ from relbench.datasets import get_dataset from relbench.tasks import get_task ``` -## Community Datasets: Temporal Graph Benchmark (TGB) - -RelBench supports artifact-backed community datasets exported from the Temporal Graph Benchmark (TGB): -- Dynamic link prediction: `rel-tgb-tgbl-*` (task: `src-dst-mrr`) -- Temporal heterogeneous link prediction: `rel-tgb-thgl-*` (tasks: `edge-type--mrr`) -- Dynamic node property prediction: `rel-tgb-tgbn-*` (task: `node-label-ndcg`) - -These datasets/tasks are distributed as pre-built `db.zip` and task-table zip artifacts (not raw-source processing in RelBench). - -Example: -```python -dataset = get_dataset("rel-tgb-tgbl-wiki-v2", download=True) -task = get_task("rel-tgb-tgbl-wiki-v2", "src-dst-mrr", download=True) -val_table = task.get_table("val", mask_input_cols=False) -``` - Get a dataset, e.g., `rel-amazon`: ```python dataset: Dataset = get_dataset("rel-amazon", download=True) diff --git a/relbench/datasets/__init__.py b/relbench/datasets/__init__.py index 2fde19e1..4a9ae440 100644 --- a/relbench/datasets/__init__.py +++ b/relbench/datasets/__init__.py @@ -160,20 +160,20 @@ def get_dataset(name: str, download=True) -> Dataset: register_dataset("dbinfer-outbrain-small", dbinfer.DBInferOutbrainSmallDataset) # Temporal Graph Benchmark (TGB) -register_dataset("rel-tgb-tgbl-wiki", tgb.TGBDataset, tgb_name="tgbl-wiki") -register_dataset("rel-tgb-tgbl-wiki-v2", tgb.TGBDataset, tgb_name="tgbl-wiki-v2") -register_dataset("rel-tgb-tgbl-review", tgb.TGBDataset, tgb_name="tgbl-review") -register_dataset("rel-tgb-tgbl-review-v2", tgb.TGBDataset, tgb_name="tgbl-review-v2") -register_dataset("rel-tgb-tgbl-coin", tgb.TGBDataset, tgb_name="tgbl-coin") -register_dataset("rel-tgb-tgbl-comment", tgb.TGBDataset, tgb_name="tgbl-comment") -register_dataset("rel-tgb-tgbl-flight", tgb.TGBDataset, tgb_name="tgbl-flight") - -register_dataset("rel-tgb-thgl-software", tgb.TGBDataset, tgb_name="thgl-software") -register_dataset("rel-tgb-thgl-forum", tgb.TGBDataset, tgb_name="thgl-forum") -register_dataset("rel-tgb-thgl-github", tgb.TGBDataset, tgb_name="thgl-github") -register_dataset("rel-tgb-thgl-myket", tgb.TGBDataset, tgb_name="thgl-myket") - -register_dataset("rel-tgb-tgbn-trade", tgb.TGBDataset, tgb_name="tgbn-trade") -register_dataset("rel-tgb-tgbn-genre", tgb.TGBDataset, tgb_name="tgbn-genre") -register_dataset("rel-tgb-tgbn-reddit", tgb.TGBDataset, tgb_name="tgbn-reddit") -register_dataset("rel-tgb-tgbn-token", tgb.TGBDataset, tgb_name="tgbn-token") +register_dataset("tgbl-wiki", tgb.TGBDataset, tgb_name="tgbl-wiki") +register_dataset("tgbl-wiki-v2", tgb.TGBDataset, tgb_name="tgbl-wiki-v2") +register_dataset("tgbl-review", tgb.TGBDataset, tgb_name="tgbl-review") +register_dataset("tgbl-review-v2", tgb.TGBDataset, tgb_name="tgbl-review-v2") +register_dataset("tgbl-coin", tgb.TGBDataset, tgb_name="tgbl-coin") +register_dataset("tgbl-comment", tgb.TGBDataset, tgb_name="tgbl-comment") +register_dataset("tgbl-flight", tgb.TGBDataset, tgb_name="tgbl-flight") + +register_dataset("thgl-software", tgb.TGBDataset, tgb_name="thgl-software") +register_dataset("thgl-forum", tgb.TGBDataset, tgb_name="thgl-forum") +register_dataset("thgl-github", tgb.TGBDataset, tgb_name="thgl-github") +register_dataset("thgl-myket", tgb.TGBDataset, tgb_name="thgl-myket") + +register_dataset("tgbn-trade", tgb.TGBDataset, tgb_name="tgbn-trade") +register_dataset("tgbn-genre", tgb.TGBDataset, tgb_name="tgbn-genre") +register_dataset("tgbn-reddit", tgb.TGBDataset, tgb_name="tgbn-reddit") +register_dataset("tgbn-token", tgb.TGBDataset, tgb_name="tgbn-token") diff --git a/relbench/datasets/hashes.json b/relbench/datasets/hashes.json index fd62d217..80d9b7f6 100644 --- a/relbench/datasets/hashes.json +++ b/relbench/datasets/hashes.json @@ -18,21 +18,21 @@ "dbinfer-seznam/db.zip": "77314bf874dc495e8a4a61f2dc5f12982bbec3c5b6b7af5555e9a2bb587154d9", "dbinfer-stackexchange/db.zip": "4c8b8ac38b56d57bc2ead3c12be1237b69cae76062ed3f176fc02ad327a84d19", - "rel-tgb-tgbl-coin/db.zip": "6de3b62bcfc59bb18ff8de0614ce5cbcf21179ad56cb58920d25ade32ec43e00", - "rel-tgb-tgbl-comment/db.zip": "4eb41776954efa2cb06cdb64e093b182335b8683e60236a6645cf4bcd83597be", - "rel-tgb-tgbl-flight/db.zip": "e82646893a45be6c5312e5dc1a774c08d929ecb9eb309c0ad9fdec6dae3b156a", - "rel-tgb-tgbl-review/db.zip": "63b405aafd8092cbda297694e649329ca806ccc3c49c8e6f53eb4a3c19075091", - "rel-tgb-tgbl-review-v2/db.zip": "a5f1a7a9661a700ebb1a3b43df74037853ad3bd6ed54c74342045d2dc8448bc2", - "rel-tgb-tgbl-wiki/db.zip": "ad7b55d1d7b7125c06588db0b8cbebe87c629a95a6c9a911369f89aca9dffdc9", - "rel-tgb-tgbl-wiki-v2/db.zip": "5eb56f8e459405e3b554ce56585903c80038d84a99ca4333b00ce32c8c7a38f1", + "tgbl-coin/db.zip": "6de3b62bcfc59bb18ff8de0614ce5cbcf21179ad56cb58920d25ade32ec43e00", + "tgbl-comment/db.zip": "4eb41776954efa2cb06cdb64e093b182335b8683e60236a6645cf4bcd83597be", + "tgbl-flight/db.zip": "e82646893a45be6c5312e5dc1a774c08d929ecb9eb309c0ad9fdec6dae3b156a", + "tgbl-review/db.zip": "63b405aafd8092cbda297694e649329ca806ccc3c49c8e6f53eb4a3c19075091", + "tgbl-review-v2/db.zip": "a5f1a7a9661a700ebb1a3b43df74037853ad3bd6ed54c74342045d2dc8448bc2", + "tgbl-wiki/db.zip": "ad7b55d1d7b7125c06588db0b8cbebe87c629a95a6c9a911369f89aca9dffdc9", + "tgbl-wiki-v2/db.zip": "5eb56f8e459405e3b554ce56585903c80038d84a99ca4333b00ce32c8c7a38f1", - "rel-tgb-tgbn-genre/db.zip": "e46aecc28315ca9872117817ce65bc1f4d00ed261ab6d35038a84bcf0ebab7bf", - "rel-tgb-tgbn-reddit/db.zip": "bc444f5cbaf7004ef7c6c8f98ae12ae72e8427ccc9710e96e46b0fc156cc952e", - "rel-tgb-tgbn-token/db.zip": "44a0f90a62642b8054ef1fbee35c862a67e5d0cec92f08b8fd44c503e0986be8", - "rel-tgb-tgbn-trade/db.zip": "8a983ae6281ea058dc8d84f6ce339d2dacf0ae9bb98bb82cca3fed54ec3fe370", + "tgbn-genre/db.zip": "e46aecc28315ca9872117817ce65bc1f4d00ed261ab6d35038a84bcf0ebab7bf", + "tgbn-reddit/db.zip": "bc444f5cbaf7004ef7c6c8f98ae12ae72e8427ccc9710e96e46b0fc156cc952e", + "tgbn-token/db.zip": "44a0f90a62642b8054ef1fbee35c862a67e5d0cec92f08b8fd44c503e0986be8", + "tgbn-trade/db.zip": "8a983ae6281ea058dc8d84f6ce339d2dacf0ae9bb98bb82cca3fed54ec3fe370", - "rel-tgb-thgl-forum/db.zip": "fdb20c1afc542e8026b85df850b5e8a539694c182eb4530924fad65964461256", - "rel-tgb-thgl-github/db.zip": "f69c1d49779a4dbead101e7325447854c8890982d254efb2b722117819ba8304", - "rel-tgb-thgl-myket/db.zip": "67a3af25553ceabe8a7eff8906c2213bd82e932cdaeb48ffa326b39d30e0f0cf", - "rel-tgb-thgl-software/db.zip": "35816ef6a07be2291349f081814e18b498a2e46e006ed21be36fdb1a8a0eb90d" + "thgl-forum/db.zip": "fdb20c1afc542e8026b85df850b5e8a539694c182eb4530924fad65964461256", + "thgl-github/db.zip": "f69c1d49779a4dbead101e7325447854c8890982d254efb2b722117819ba8304", + "thgl-myket/db.zip": "67a3af25553ceabe8a7eff8906c2213bd82e932cdaeb48ffa326b39d30e0f0cf", + "thgl-software/db.zip": "35816ef6a07be2291349f081814e18b498a2e46e006ed21be36fdb1a8a0eb90d" } diff --git a/relbench/tasks/__init__.py b/relbench/tasks/__init__.py index 80451efc..c2eb9085 100644 --- a/relbench/tasks/__init__.py +++ b/relbench/tasks/__init__.py @@ -551,13 +551,13 @@ def get_task(dataset_name: str, task_name: str, download=False) -> BaseTask: # Temporal Graph Benchmark (TGB) for dataset_name in [ - "rel-tgb-tgbl-wiki", - "rel-tgb-tgbl-wiki-v2", - "rel-tgb-tgbl-review", - "rel-tgb-tgbl-review-v2", - "rel-tgb-tgbl-coin", - "rel-tgb-tgbl-comment", - "rel-tgb-tgbl-flight", + "tgbl-wiki", + "tgbl-wiki-v2", + "tgbl-review", + "tgbl-review-v2", + "tgbl-coin", + "tgbl-comment", + "tgbl-flight", ]: register_task( dataset_name, @@ -579,16 +579,16 @@ def _register_thgl_edge_type_tasks(dataset_name: str, edge_types: list[int]) -> ) -_register_thgl_edge_type_tasks("rel-tgb-thgl-software", list(range(14))) -_register_thgl_edge_type_tasks("rel-tgb-thgl-github", list(range(14))) -_register_thgl_edge_type_tasks("rel-tgb-thgl-forum", [0, 1]) -_register_thgl_edge_type_tasks("rel-tgb-thgl-myket", [0, 1]) +_register_thgl_edge_type_tasks("thgl-software", list(range(14))) +_register_thgl_edge_type_tasks("thgl-github", list(range(14))) +_register_thgl_edge_type_tasks("thgl-forum", [0, 1]) +_register_thgl_edge_type_tasks("thgl-myket", [0, 1]) for dataset_name in [ - "rel-tgb-tgbn-trade", - "rel-tgb-tgbn-genre", - "rel-tgb-tgbn-reddit", - "rel-tgb-tgbn-token", + "tgbn-trade", + "tgbn-genre", + "tgbn-reddit", + "tgbn-token", ]: register_task( dataset_name, diff --git a/relbench/tasks/hashes.json b/relbench/tasks/hashes.json index 86289918..bc61201b 100644 --- a/relbench/tasks/hashes.json +++ b/relbench/tasks/hashes.json @@ -82,47 +82,47 @@ "dbinfer-stackexchange/tasks/churn.zip": "d649b0a70dc8be9fca797c82e14d19ed4311637711144f4c5b74249081b107ca", "dbinfer-stackexchange/tasks/upvote.zip": "af8f516422e3e19c0d76f13c00c4c0577cdc218b84468006eda369de8f9c7953", - "rel-tgb-tgbl-coin/tasks/src-dst-mrr.zip": "702ff505d95955b29c9998eb4f9c2e38e6c112c2b9e21ed7c2aa75ae4a923280", - "rel-tgb-tgbl-comment/tasks/src-dst-mrr.zip": "50846c6b82433bcac2bb025c3e7edfed46b651a9d06b462b058a071f08741fdc", - "rel-tgb-tgbl-flight/tasks/src-dst-mrr.zip": "10527b2122a73bcb35a6b3e6230ded99242825bf955c363b710cce20545c8104", - "rel-tgb-tgbl-review-v2/tasks/src-dst-mrr.zip": "aeded6586c00995a33881305717c3a625f96a1cc7253d2806af6125b53173fb3", - "rel-tgb-tgbl-review/tasks/src-dst-mrr.zip": "aeded6586c00995a33881305717c3a625f96a1cc7253d2806af6125b53173fb3", - "rel-tgb-tgbl-wiki-v2/tasks/src-dst-mrr.zip": "9517a458de6dc921e54af91154f1823d7bdbfa16207650df16a32bcd9b1aab92", - "rel-tgb-tgbl-wiki/tasks/src-dst-mrr.zip": "a105ef6cf8dcdd5f94d1fbcfa856a3ecf10821a9d217b53c8266eca2d8b04b1b", - "rel-tgb-tgbn-genre/tasks/node-label-ndcg.zip": "4fe66c0e9e0442fa9c4979fa386f8249f812b19b817faaefae46f3b82f8c3dce", - "rel-tgb-tgbn-reddit/tasks/node-label-ndcg.zip": "bdcdb15597719541e2873408d36ccc7589ddc25e8fd9be60d3bc2928a4fb02d5", - "rel-tgb-tgbn-token/tasks/node-label-ndcg.zip": "64daa8ca8d3ff62bff0929f8af8dcbb0836c004646102df94df8236b13fffba1", - "rel-tgb-tgbn-trade/tasks/node-label-ndcg.zip": "f69920871dac5e365fa521d180d9a234802318bdbb313a59a362fe6c68f14e1d", - "rel-tgb-thgl-forum/tasks/edge-type-0-mrr.zip": "14975fa000b5a4d8458fe5f9fd23d6662af8205d4ef9d3463da2051b16b3a5a8", - "rel-tgb-thgl-forum/tasks/edge-type-1-mrr.zip": "fffe1b4dc3ed985e1d442d3db56da2ebfb18b05ccf3fefe0d13ae64431993027", - "rel-tgb-thgl-github/tasks/edge-type-0-mrr.zip": "9c346bf63c2dcc36171b62884c18c4062781ca689ec861177d70f80531856e43", - "rel-tgb-thgl-github/tasks/edge-type-1-mrr.zip": "e0db81eec7a8b12c488fe7c773f7226e7f7914ec5613b9ba11192a3bb2c19e21", - "rel-tgb-thgl-github/tasks/edge-type-10-mrr.zip": "dcfaa9094d7d0f7944dda32075d165c5f97c4d0d2530dced6e27121163599a65", - "rel-tgb-thgl-github/tasks/edge-type-11-mrr.zip": "7f707920f58c4fea8dabcfe0e511be1a60603ba73301f28a698a4fce87e804da", - "rel-tgb-thgl-github/tasks/edge-type-12-mrr.zip": "58cf770c6917081865c4c8f5660670e3f821618bc9a3a03f57d5a12891254ea4", - "rel-tgb-thgl-github/tasks/edge-type-13-mrr.zip": "cfaa8350612cc1aa99d7dcaccd32849bae294e3b2ef8e3cf2232c8b82e45650c", - "rel-tgb-thgl-github/tasks/edge-type-2-mrr.zip": "d42be110b69c9dae1d952459118864af5701a18393a3ca11ddb96766889ea5be", - "rel-tgb-thgl-github/tasks/edge-type-3-mrr.zip": "61e31b44e361a23c498e7aacdc271ba8b4b276b0134d1c07e0ecd694ab732665", - "rel-tgb-thgl-github/tasks/edge-type-4-mrr.zip": "61832694d14aa45ddf393122e1530d63a0c93d902913767d29928bbadf208f0e", - "rel-tgb-thgl-github/tasks/edge-type-5-mrr.zip": "a5378c8797a79b655b6eb4678ce2b6f5d1d39e992e6ceda08a1a6f3f7b26cd66", - "rel-tgb-thgl-github/tasks/edge-type-6-mrr.zip": "76d606d9f7d6c69be22b39e7b1086a7d75fd47b29f682538153fdc9ccf08e1bf", - "rel-tgb-thgl-github/tasks/edge-type-7-mrr.zip": "7c9405cc05cad4f299967aee0510eaf718a2dee1884bac790e76c81263b72f22", - "rel-tgb-thgl-github/tasks/edge-type-8-mrr.zip": "44cad0c9358293530781d20b78cb87bd21ebf8fea36628c0875ed4558a6a9117", - "rel-tgb-thgl-github/tasks/edge-type-9-mrr.zip": "19be77af66f1cb0cb640158f97d42a0d16431eb79bd23ed1aaace4b68519c8a9", - "rel-tgb-thgl-myket/tasks/edge-type-0-mrr.zip": "979133caaa02ddf444e3b23ada4b5cbe11147b18adab90537740cea914e36daa", - "rel-tgb-thgl-myket/tasks/edge-type-1-mrr.zip": "f4f84ca245d499c0773ec0c95619e5f6041461c3233cce96ad3a6302da41c12c", - "rel-tgb-thgl-software/tasks/edge-type-0-mrr.zip": "6262267aae8309bda302c64e365d227d036f95f934ee8ba6fb038f20d28b3fed", - "rel-tgb-thgl-software/tasks/edge-type-1-mrr.zip": "7de5c8664e9b78460cad9345a4a537e9d75641f0875605a5151e189ab568fd0e", - "rel-tgb-thgl-software/tasks/edge-type-10-mrr.zip": "b3a1ed7410ed737fb2ca6ecd4f1de59cc0d7e0cfd29fdcb68ccd13e76e1997e3", - "rel-tgb-thgl-software/tasks/edge-type-11-mrr.zip": "9e24b4c1365dfb989fc6c4db4cc19146e998f852205fca0fe7f1fe4a047ef44d", - "rel-tgb-thgl-software/tasks/edge-type-12-mrr.zip": "1580ff8adde8202edae05a555b12a7fe38ce0addd360010ec925b4a7503350d2", - "rel-tgb-thgl-software/tasks/edge-type-13-mrr.zip": "e4a51a47777c34748fec03813bb2368bf841f90479bd37ca41c653b9203491cb", - "rel-tgb-thgl-software/tasks/edge-type-2-mrr.zip": "1b8177e304b62939b3081c054b56b99d5874300e101057593668af455558b385", - "rel-tgb-thgl-software/tasks/edge-type-3-mrr.zip": "e0837e81655cfe953b805af162e7904dfe457a93783701d50a26d819dfec848f", - "rel-tgb-thgl-software/tasks/edge-type-4-mrr.zip": "0d15eeaf8d9e456912518aedd2e6a239b6fddc6941095b95240ef41fb34f14cc", - "rel-tgb-thgl-software/tasks/edge-type-5-mrr.zip": "1df8a435f4e978c976288c3d0a53f7430e759d27cd91ad966a1af6c3814da480", - "rel-tgb-thgl-software/tasks/edge-type-6-mrr.zip": "51fa5917363476973c048ef19306188d0d8cf6c218e316f4949d08d10c9092f0", - "rel-tgb-thgl-software/tasks/edge-type-7-mrr.zip": "ee8a0e231255e15fe58283b9f9243ca438695e1760db8cb24168b6bfbb6c078a", - "rel-tgb-thgl-software/tasks/edge-type-8-mrr.zip": "26e37301204380e27934ed7046201ebd4edf72ac7b430c8ef792ad3defc9e638", - "rel-tgb-thgl-software/tasks/edge-type-9-mrr.zip": "d68f2bcafee7a7d7761e4b4eadec36c500cb3e8711d1b23196ecb160962b5b9e" + "tgbl-coin/tasks/src-dst-mrr.zip": "702ff505d95955b29c9998eb4f9c2e38e6c112c2b9e21ed7c2aa75ae4a923280", + "tgbl-comment/tasks/src-dst-mrr.zip": "50846c6b82433bcac2bb025c3e7edfed46b651a9d06b462b058a071f08741fdc", + "tgbl-flight/tasks/src-dst-mrr.zip": "10527b2122a73bcb35a6b3e6230ded99242825bf955c363b710cce20545c8104", + "tgbl-review-v2/tasks/src-dst-mrr.zip": "aeded6586c00995a33881305717c3a625f96a1cc7253d2806af6125b53173fb3", + "tgbl-review/tasks/src-dst-mrr.zip": "aeded6586c00995a33881305717c3a625f96a1cc7253d2806af6125b53173fb3", + "tgbl-wiki-v2/tasks/src-dst-mrr.zip": "9517a458de6dc921e54af91154f1823d7bdbfa16207650df16a32bcd9b1aab92", + "tgbl-wiki/tasks/src-dst-mrr.zip": "a105ef6cf8dcdd5f94d1fbcfa856a3ecf10821a9d217b53c8266eca2d8b04b1b", + "tgbn-genre/tasks/node-label-ndcg.zip": "4fe66c0e9e0442fa9c4979fa386f8249f812b19b817faaefae46f3b82f8c3dce", + "tgbn-reddit/tasks/node-label-ndcg.zip": "bdcdb15597719541e2873408d36ccc7589ddc25e8fd9be60d3bc2928a4fb02d5", + "tgbn-token/tasks/node-label-ndcg.zip": "64daa8ca8d3ff62bff0929f8af8dcbb0836c004646102df94df8236b13fffba1", + "tgbn-trade/tasks/node-label-ndcg.zip": "f69920871dac5e365fa521d180d9a234802318bdbb313a59a362fe6c68f14e1d", + "thgl-forum/tasks/edge-type-0-mrr.zip": "14975fa000b5a4d8458fe5f9fd23d6662af8205d4ef9d3463da2051b16b3a5a8", + "thgl-forum/tasks/edge-type-1-mrr.zip": "fffe1b4dc3ed985e1d442d3db56da2ebfb18b05ccf3fefe0d13ae64431993027", + "thgl-github/tasks/edge-type-0-mrr.zip": "9c346bf63c2dcc36171b62884c18c4062781ca689ec861177d70f80531856e43", + "thgl-github/tasks/edge-type-1-mrr.zip": "e0db81eec7a8b12c488fe7c773f7226e7f7914ec5613b9ba11192a3bb2c19e21", + "thgl-github/tasks/edge-type-10-mrr.zip": "dcfaa9094d7d0f7944dda32075d165c5f97c4d0d2530dced6e27121163599a65", + "thgl-github/tasks/edge-type-11-mrr.zip": "7f707920f58c4fea8dabcfe0e511be1a60603ba73301f28a698a4fce87e804da", + "thgl-github/tasks/edge-type-12-mrr.zip": "58cf770c6917081865c4c8f5660670e3f821618bc9a3a03f57d5a12891254ea4", + "thgl-github/tasks/edge-type-13-mrr.zip": "cfaa8350612cc1aa99d7dcaccd32849bae294e3b2ef8e3cf2232c8b82e45650c", + "thgl-github/tasks/edge-type-2-mrr.zip": "d42be110b69c9dae1d952459118864af5701a18393a3ca11ddb96766889ea5be", + "thgl-github/tasks/edge-type-3-mrr.zip": "61e31b44e361a23c498e7aacdc271ba8b4b276b0134d1c07e0ecd694ab732665", + "thgl-github/tasks/edge-type-4-mrr.zip": "61832694d14aa45ddf393122e1530d63a0c93d902913767d29928bbadf208f0e", + "thgl-github/tasks/edge-type-5-mrr.zip": "a5378c8797a79b655b6eb4678ce2b6f5d1d39e992e6ceda08a1a6f3f7b26cd66", + "thgl-github/tasks/edge-type-6-mrr.zip": "76d606d9f7d6c69be22b39e7b1086a7d75fd47b29f682538153fdc9ccf08e1bf", + "thgl-github/tasks/edge-type-7-mrr.zip": "7c9405cc05cad4f299967aee0510eaf718a2dee1884bac790e76c81263b72f22", + "thgl-github/tasks/edge-type-8-mrr.zip": "44cad0c9358293530781d20b78cb87bd21ebf8fea36628c0875ed4558a6a9117", + "thgl-github/tasks/edge-type-9-mrr.zip": "19be77af66f1cb0cb640158f97d42a0d16431eb79bd23ed1aaace4b68519c8a9", + "thgl-myket/tasks/edge-type-0-mrr.zip": "979133caaa02ddf444e3b23ada4b5cbe11147b18adab90537740cea914e36daa", + "thgl-myket/tasks/edge-type-1-mrr.zip": "f4f84ca245d499c0773ec0c95619e5f6041461c3233cce96ad3a6302da41c12c", + "thgl-software/tasks/edge-type-0-mrr.zip": "6262267aae8309bda302c64e365d227d036f95f934ee8ba6fb038f20d28b3fed", + "thgl-software/tasks/edge-type-1-mrr.zip": "7de5c8664e9b78460cad9345a4a537e9d75641f0875605a5151e189ab568fd0e", + "thgl-software/tasks/edge-type-10-mrr.zip": "b3a1ed7410ed737fb2ca6ecd4f1de59cc0d7e0cfd29fdcb68ccd13e76e1997e3", + "thgl-software/tasks/edge-type-11-mrr.zip": "9e24b4c1365dfb989fc6c4db4cc19146e998f852205fca0fe7f1fe4a047ef44d", + "thgl-software/tasks/edge-type-12-mrr.zip": "1580ff8adde8202edae05a555b12a7fe38ce0addd360010ec925b4a7503350d2", + "thgl-software/tasks/edge-type-13-mrr.zip": "e4a51a47777c34748fec03813bb2368bf841f90479bd37ca41c653b9203491cb", + "thgl-software/tasks/edge-type-2-mrr.zip": "1b8177e304b62939b3081c054b56b99d5874300e101057593668af455558b385", + "thgl-software/tasks/edge-type-3-mrr.zip": "e0837e81655cfe953b805af162e7904dfe457a93783701d50a26d819dfec848f", + "thgl-software/tasks/edge-type-4-mrr.zip": "0d15eeaf8d9e456912518aedd2e6a239b6fddc6941095b95240ef41fb34f14cc", + "thgl-software/tasks/edge-type-5-mrr.zip": "1df8a435f4e978c976288c3d0a53f7430e759d27cd91ad966a1af6c3814da480", + "thgl-software/tasks/edge-type-6-mrr.zip": "51fa5917363476973c048ef19306188d0d8cf6c218e316f4949d08d10c9092f0", + "thgl-software/tasks/edge-type-7-mrr.zip": "ee8a0e231255e15fe58283b9f9243ca438695e1760db8cb24168b6bfbb6c078a", + "thgl-software/tasks/edge-type-8-mrr.zip": "26e37301204380e27934ed7046201ebd4edf72ac7b430c8ef792ad3defc9e638", + "thgl-software/tasks/edge-type-9-mrr.zip": "d68f2bcafee7a7d7761e4b4eadec36c500cb3e8711d1b23196ecb160962b5b9e" }