diff --git a/relbench/datasets/__init__.py b/relbench/datasets/__init__.py index 2464e03f..91dcf8f5 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 @@ -152,3 +153,22 @@ 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) + +# Temporal Graph Benchmark (TGB) +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 670e1b90..cce44f0c 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", + + "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", + + "tgbn-genre/db.zip": "e46aecc28315ca9872117817ce65bc1f4d00ed261ab6d35038a84bcf0ebab7bf", + "tgbn-reddit/db.zip": "bc444f5cbaf7004ef7c6c8f98ae12ae72e8427ccc9710e96e46b0fc156cc952e", + "tgbn-token/db.zip": "44a0f90a62642b8054ef1fbee35c862a67e5d0cec92f08b8fd44c503e0986be8", + "tgbn-trade/db.zip": "8a983ae6281ea058dc8d84f6ce339d2dacf0ae9bb98bb82cca3fed54ec3fe370", + + "thgl-forum/db.zip": "fdb20c1afc542e8026b85df850b5e8a539694c182eb4530924fad65964461256", + "thgl-github/db.zip": "f69c1d49779a4dbead101e7325447854c8890982d254efb2b722117819ba8304", + "thgl-myket/db.zip": "67a3af25553ceabe8a7eff8906c2213bd82e932cdaeb48ffa326b39d30e0f0cf", + "thgl-software/db.zip": "35816ef6a07be2291349f081814e18b498a2e46e006ed21be36fdb1a8a0eb90d" } diff --git a/relbench/datasets/tgb.py b/relbench/datasets/tgb.py new file mode 100644 index 00000000..ae9f914a --- /dev/null +++ b/relbench/datasets/tgb.py @@ -0,0 +1,75 @@ +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] = { + "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), + "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), + "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"""Temporal Graph Benchmark (TGB) datasets exported to RelBench format.""" + + 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] + 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 83b45752..cb643d95 100644 --- a/relbench/tasks/__init__.py +++ b/relbench/tasks/__init__.py @@ -20,6 +20,7 @@ mimic, ratebeer, stack, + tgb, trial, ) from relbench.utils import get_relbench_cache_dir @@ -544,3 +545,52 @@ 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) + +# Temporal Graph Benchmark (TGB) +for dataset_name in [ + "tgbl-wiki", + "tgbl-wiki-v2", + "tgbl-review", + "tgbl-review-v2", + "tgbl-coin", + "tgbl-comment", + "tgbl-flight", +]: + register_task( + dataset_name, + "src-dst-mrr", + tgb.TGBOneVsManyLinkPredTask, + spec=tgb.TGBLinkPredSpec(event_table="events"), + k_value=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, + ) + + +_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 [ + "tgbn-trade", + "tgbn-genre", + "tgbn-reddit", + "tgbn-token", +]: + 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 dfa3a5fe..c9be091f 100644 --- a/relbench/tasks/hashes.json +++ b/relbench/tasks/hashes.json @@ -80,5 +80,49 @@ "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", + + "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" } diff --git a/relbench/tasks/tgb.py b/relbench/tasks/tgb.py new file mode 100644 index 00000000..240757d7 --- /dev/null +++ b/relbench/tasks/tgb.py @@ -0,0 +1,627 @@ +from __future__ import annotations + +import pickle +import re +from dataclasses import dataclass +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 BaseTask, Dataset, RecommendationTask, Table, TaskType +from relbench.metrics import ( + link_prediction_map, + link_prediction_precision, + link_prediction_recall, +) + + +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 + ) + + +def _tgb_eval_hits_and_mrr( + y_pred_pos: np.ndarray, + y_pred_neg: np.ndarray, + *, + k_value: int, +) -> dict[str, float]: + 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) + 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 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" + time_col: str = "event_ts" + + +class TGBOneVsManyLinkPredTask(BaseTask): + r"""TGB link prediction with official one-vs-many MRR/Hits@k evaluation.""" + + task_type = TaskType.LINK_PREDICTION + timedelta = pd.Timedelta(seconds=1) + metrics = [] + num_eval_timestamps = 1 + + def __init__( + self, + dataset: Dataset, + *, + spec: TGBLinkPredSpec, + k_value: int = 10, + cache_dir: Optional[str] = None, + ) -> None: + self.spec = spec + 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 + + 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, timestamps): # pragma: no cover + raise RuntimeError( + "TGBOneVsManyLinkPredTask expects precomputed task tables " + "(train/val/test.parquet) and overrides _get_table()." + ) + + 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]) + ] + 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: + 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." + ) + + 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" + + @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]: + 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: + 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: + 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: + 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 TGBNextLinkPredTask(RecommendationTask): + r"""TGB-style "next" link prediction task in RelBench recommendation format. + + 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`) + + 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 + timedelta = pd.Timedelta(seconds=1) + num_eval_timestamps = 1 + metrics = [link_prediction_precision, link_prediction_recall, link_prediction_map] + + def __init__( + self, + dataset: Dataset, + *, + 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.eval_k = int(eval_k) + + self.time_col = str(time_col) + self.src_entity_col = str(src_entity_col) + self.dst_entity_col = str(dst_entity_col) + + 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: + 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": + 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) + + 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, 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 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).") + 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 + + 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 + + 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}." + ) + + 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: + 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: + 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)}." + ) + + 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)}