relbench 2.0.2__tar.gz → 2.1.1__tar.gz

{relbench-2.0.2 → relbench-2.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: relbench
-Version: 2.0.2
+Version: 2.1.1
 Summary: RelBench: Relational Deep Learning Benchmark
 Keywords: 
 Author-email: RelBench Team <relbench@cs.stanford.edu>
@@ -18,6 +18,7 @@ Requires-Dist: typing-extensions
 Requires-Dist: datasets
 Requires-Dist: redelex ; extra == "ctu"
 Requires-Dist: dbinfer-relbench-adapter ; extra == "dbinfer"
+Requires-Dist: db-dtypes ; extra == "dbinfer"
 Requires-Dist: pre-commit ; extra == "dev"
 Requires-Dist: sentence-transformers ; extra == "example"
 Requires-Dist: pytorch_frame[full] ; extra == "example"
@@ -48,14 +49,20 @@ Provides-Extra: test
 
 
 <!-- [<img align="center" src="https://relbench.stanford.edu/img/favicon.png" width="20px" />   -->
-[**Website**](https://relbench.stanford.edu) | [**Position Paper**](https://proceedings.mlr.press/v235/fey24a.html) |  [**Benchmark Paper**](https://arxiv.org/abs/2407.20060) | [**Mailing List**](https://groups.google.com/forum/#!forum/relbench/join)
+[**Website**](https://relbench.stanford.edu) | [**Position Paper**](https://proceedings.mlr.press/v235/fey24a.html) |  [**RelBench v1 Paper**](https://arxiv.org/abs/2407.20060) | [**RelBench v2 Paper**](https://arxiv.org/abs/2602.12606) | [**Mailing List**](https://groups.google.com/forum/#!forum/relbench/join)
 
 # News
 
+**February 13, 2026: RelBench v2 paper + Temporal Graph Benchmark integration**
+
+The RelBench v2 paper is now accessible as a preprint! Please see the paper on [arXiv](https://arxiv.org/abs/2602.12606).
+
+Alongside our paper, we also integrate the [Temporal Graph Benchmark](https://tgb.complexdatalab.com/) (TGB) into RelBench. TGB integration includes translating time-stamped event streams into normalized relational schemas, which enables direct comparison between temporal graph models and relational deep learning models.
+
 **January 12, 2026: RelBench v2 is now released!**
 - Introducing Autocomplete tasks: new task paradigm to predict existing columns in the database.
 - 4 new databases: [SALT](https://relbench.stanford.edu/datasets/rel-salt), [RateBeer](https://relbench.stanford.edu/datasets/rel-ratebeer), [arXiv](https://relbench.stanford.edu/datasets/rel-arxiv), and [MIMIC-IV](https://relbench.stanford.edu/datasets/rel-mimic).
-- 40 new predictive tasks, including 28 Autocomplete tasks across new and existing databases.
+- 36 new predictive tasks, including 23 Autocomplete tasks across new and existing databases.
 - CTU integration: 70+ relational datasets from the CTU repository via [ReDeLEx](https://github.com/jakubpeleska/redelex).
 - Direct SQL database connectivity via [ReDeLEx](https://github.com/jakubpeleska/redelex).
 - 4DBInfer integration: 7 relational datasets from the [4DBInfer](https://github.com/awslabs/multi-table-benchmark) repository in RelBench format.
@@ -79,7 +86,7 @@ Provides-Extra: test
 
 Relational Deep Learning is a new approach for end-to-end representation learning on data spread across multiple tables, such as in a _relational database_ (see our [position paper](https://relbench.stanford.edu/paper.pdf)). Relational databases are the world's most widely used data management system, and are used for industrial and scientific purposes across many domains. RelBench is a benchmark designed to facilitate efficient, robust and reproducible research on end-to-end deep learning over relational databases.
 
-RelBench v1 contains 7 realistic, large-scale, and diverse relational databases spanning domains including medical, social networks, e-commerce and sport. RelBench v2 adds 4 more, now totaling 11 databases. Each database has multiple predictive tasks (70 in total) defined, each carefully scoped to be both challenging and of domain-specific importance. It provides full support for data downloading, task specification and standardized evaluation in an ML-framework-agnostic manner.
+RelBench v1 contains 7 realistic, large-scale, and diverse relational databases spanning domains including medical, social networks, e-commerce and sport. RelBench v2 adds 4 more, now totaling 11 databases. Each database has multiple predictive tasks (66 in total) defined, each carefully scoped to be both challenging and of domain-specific importance. It provides full support for data downloading, task specification and standardized evaluation in an ML-framework-agnostic manner.
 
 Additionally, RelBench provides a first open-source implementation of a Graph Neural Network based approach to relational deep learning. This implementation uses [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric) to load the data as a graph and train GNN models, and [PyTorch Frame](https://github.com/pyg-team/pytorch-frame) for modeling tabular data. Finally, there is an open [leaderboard](https://huggingface.co/relbench) for tracking progress.
 
@@ -102,7 +109,7 @@ This paper outlines our proposal for how to do end-to-end deep learning on relat
 <p align="center"><img src="https://relbench.stanford.edu/img/relbench-fig.png" alt="logo" width="900px" /></p>
 
 RelBench has the following main components:
-1. 11 databases with a total of 70 tasks; both of these automatically downloadable for ease of use
+1. 11 databases with a total of 66 tasks; both of these automatically downloadable for ease of use
 2. Easy data loading, and graph construction from pkey-fkey links
 3. Your own model, which can use any deep learning stack since RelBench is framework-agnostic. We provide a first model implementation using PyTorch Geometric and PyTorch Frame.
 4. Standardized evaluators - all you need to do is produce a list of predictions for test samples, and RelBench computes metrics to ensure standardized evaluation
@@ -118,7 +125,30 @@ pip install relbench
 
 This will allow usage of the core RelBench data and task loading functionality.
 
-<details markdown="1"><summary>Including CTU datasets</summary>
+
+To additionally use `relbench.modeling`, which requires [PyTorch](https://pytorch.org/), [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric) and [PyTorch Frame](https://github.com/pyg-team/pytorch-frame), install these dependencies manually or do:
+
+```bash
+pip install relbench[full]
+```
+
+
+For the scripts in the `examples` directory, use:
+```bash
+pip install relbench[example]
+```
+
+Then, to run a script:
+```bash
+git clone https://github.com/snap-stanford/relbench
+cd relbench/examples
+python gnn_entity.py --dataset rel-f1 --task driver-position
+```
+
+
+## Using External Integrations
+
+**Using CTU datasets**
 
 To use datasets from the CTU repository, use:
 ```bash
@@ -127,7 +157,7 @@ pip install relbench[ctu]
 
 If you use the CTU datasets in your work, please cite [ReDeLEx](https://github.com/jakubpeleska/redelex) as below:
 
-```
+```bibtex
 @misc{peleska2025redelex,
   title={REDELEX: A Framework for Relational Deep Learning Exploration},
   author={Jakub Peleška and Gustav Šír},
@@ -139,9 +169,9 @@ If you use the CTU datasets in your work, please cite [ReDeLEx](https://github.c
 }
 ```
 
-</details>
 
-<details markdown="1"><summary>Including 4DBInfer datasets</summary>
+
+**Using 4DBInfer datasets**
 
 To use datasets from the 4DBInfer repository, use:
 ```bash
@@ -149,7 +179,7 @@ pip install relbench[dbinfer]
 ```
 
 If you use the 4DBInfer datasets in your work, please cite [4DBInfer](https://github.com/awslabs/multi-table-benchmark) as below:
-```
+```bibtex
 @article{dbinfer,
   title={4DBInfer: A 4D Benchmarking Toolbox for Graph-Centric Predictive Modeling on Relational DBs},
   author={Wang, Minjie and Gan, Quan and Wipf, David and Cai, Zhenkun and Li, Ning and Tang, Jianheng and Zhang, Yanlin and Zhang, Zizhao and Mao, Zunyao and Song, Yakun and Wang, Yanbo and Li, Jiahang and Zhang, Han and Yang, Guang and Qin, Xiao and Lei, Chuan and Zhang, Muhan and Zhang, Weinan and Faloutsos, Christos and Zhang, Zheng},
@@ -158,28 +188,29 @@ If you use the 4DBInfer datasets in your work, please cite [4DBInfer](https://gi
 }
 ```
 
-</details>
-<br>
+**Using TGB datasets**
 
+[TGB](https://tgb.complexdatalab.com/) datasets and tasks can be loaded in the RelBench format, and they cover bipartite link prediction (`tgbl-*`), heterogeneous link prediction (`thgl-*`), and node property prediction (`tgbn-*`). See `relbench/tasks/tgb.py` for the official TGB evaluation protocols (one-vs-many MRR / Hits@10 for link prediction, NDCG@10 for node property prediction). Standard RelBench style evaluation works as it is.
 
-To additionally use `relbench.modeling`, which requires [PyTorch](https://pytorch.org/), [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric) and [PyTorch Frame](https://github.com/pyg-team/pytorch-frame), install these dependencies manually or do:
-
-```bash
-pip install relbench[full]
+If you use the TGB datasets in your work, please cite [TGB](https://tgb.complexdatalab.com/) as below:
+```bibtex
+@article{huang2023temporal,
+  title={Temporal graph benchmark for machine learning on temporal graphs},
+  author={Huang, Shenyang and Poursafaei, Farimah and Danovitch, Jacob and Fey, Matthias and Hu, Weihua and Rossi, Emanuele and Leskovec, Jure and Bronstein, Michael and Rabusseau, Guillaume and Rabbany, Reihaneh},
+  journal={Advances in Neural Information Processing Systems},
+  year={2023}
+}
 ```
 
-
-For the scripts in the `examples` directory, use:
-```bash
-pip install relbench[example]
+```bibtex
+@article{huang2024tgb2,
+  title={TGB 2.0: A Benchmark for Learning on Temporal Knowledge Graphs and Heterogeneous Graphs},
+  author={Gastinger, Julia and Huang, Shenyang and Galkin, Mikhail and Loghmani, Erfan and Parviz, Ali and Poursafaei, Farimah and Danovitch, Jacob and Rossi, Emanuele and Koutis, Ioannis and Stuckenschmidt, Heiner and      Rabbany, Reihaneh and Rabusseau, Guillaume},
+  journal={Advances in Neural Information Processing Systems},
+  year={2024}
+}
 ```
 
-Then, to run a script:
-```bash
-git clone https://github.com/snap-stanford/relbench
-cd relbench/examples
-python gnn_entity.py --dataset rel-f1 --task driver-position
-```
 
 
 # Package Usage
@@ -298,3 +329,16 @@ If you use RelBench in your work, please cite our position and benchmark papers:
 }
 ```
 
+If you use RelBench v2 in your work, please cite:
+```bibtex
+@misc{gu2026relbenchv2,
+      title={{RelBench} v2: A Large-Scale Benchmark and Repository for Relational Data},
+      author={Justin Gu and Rishabh Ranjan and Charilaos Kanatsoulis and Haiming Tang and Martin Jurkovic and Valter Hudovernik and Mark Znidar and Pranshu Chaturvedi and Parth Shroff and Fengyu Li and Jure Leskovec},
+      year={2026},
+      eprint={2602.12606},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2602.12606},
+}
+```
+

{relbench-2.0.2 → relbench-2.1.1}/README.md

@@ -12,14 +12,20 @@
 
 
 <!-- [<img align="center" src="https://relbench.stanford.edu/img/favicon.png" width="20px" />   -->
-[**Website**](https://relbench.stanford.edu) | [**Position Paper**](https://proceedings.mlr.press/v235/fey24a.html) |  [**Benchmark Paper**](https://arxiv.org/abs/2407.20060) | [**Mailing List**](https://groups.google.com/forum/#!forum/relbench/join)
+[**Website**](https://relbench.stanford.edu) | [**Position Paper**](https://proceedings.mlr.press/v235/fey24a.html) |  [**RelBench v1 Paper**](https://arxiv.org/abs/2407.20060) | [**RelBench v2 Paper**](https://arxiv.org/abs/2602.12606) | [**Mailing List**](https://groups.google.com/forum/#!forum/relbench/join)
 
 # News
 
+**February 13, 2026: RelBench v2 paper + Temporal Graph Benchmark integration**
+
+The RelBench v2 paper is now accessible as a preprint! Please see the paper on [arXiv](https://arxiv.org/abs/2602.12606).
+
+Alongside our paper, we also integrate the [Temporal Graph Benchmark](https://tgb.complexdatalab.com/) (TGB) into RelBench. TGB integration includes translating time-stamped event streams into normalized relational schemas, which enables direct comparison between temporal graph models and relational deep learning models.
+
 **January 12, 2026: RelBench v2 is now released!**
 - Introducing Autocomplete tasks: new task paradigm to predict existing columns in the database.
 - 4 new databases: [SALT](https://relbench.stanford.edu/datasets/rel-salt), [RateBeer](https://relbench.stanford.edu/datasets/rel-ratebeer), [arXiv](https://relbench.stanford.edu/datasets/rel-arxiv), and [MIMIC-IV](https://relbench.stanford.edu/datasets/rel-mimic).
-- 40 new predictive tasks, including 28 Autocomplete tasks across new and existing databases.
+- 36 new predictive tasks, including 23 Autocomplete tasks across new and existing databases.
 - CTU integration: 70+ relational datasets from the CTU repository via [ReDeLEx](https://github.com/jakubpeleska/redelex).
 - Direct SQL database connectivity via [ReDeLEx](https://github.com/jakubpeleska/redelex).
 - 4DBInfer integration: 7 relational datasets from the [4DBInfer](https://github.com/awslabs/multi-table-benchmark) repository in RelBench format.
@@ -43,7 +49,7 @@
 
 Relational Deep Learning is a new approach for end-to-end representation learning on data spread across multiple tables, such as in a _relational database_ (see our [position paper](https://relbench.stanford.edu/paper.pdf)). Relational databases are the world's most widely used data management system, and are used for industrial and scientific purposes across many domains. RelBench is a benchmark designed to facilitate efficient, robust and reproducible research on end-to-end deep learning over relational databases.
 
-RelBench v1 contains 7 realistic, large-scale, and diverse relational databases spanning domains including medical, social networks, e-commerce and sport. RelBench v2 adds 4 more, now totaling 11 databases. Each database has multiple predictive tasks (70 in total) defined, each carefully scoped to be both challenging and of domain-specific importance. It provides full support for data downloading, task specification and standardized evaluation in an ML-framework-agnostic manner.
+RelBench v1 contains 7 realistic, large-scale, and diverse relational databases spanning domains including medical, social networks, e-commerce and sport. RelBench v2 adds 4 more, now totaling 11 databases. Each database has multiple predictive tasks (66 in total) defined, each carefully scoped to be both challenging and of domain-specific importance. It provides full support for data downloading, task specification and standardized evaluation in an ML-framework-agnostic manner.
 
 Additionally, RelBench provides a first open-source implementation of a Graph Neural Network based approach to relational deep learning. This implementation uses [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric) to load the data as a graph and train GNN models, and [PyTorch Frame](https://github.com/pyg-team/pytorch-frame) for modeling tabular data. Finally, there is an open [leaderboard](https://huggingface.co/relbench) for tracking progress.
 
@@ -66,7 +72,7 @@ This paper outlines our proposal for how to do end-to-end deep learning on relat
 <p align="center"><img src="https://relbench.stanford.edu/img/relbench-fig.png" alt="logo" width="900px" /></p>
 
 RelBench has the following main components:
-1. 11 databases with a total of 70 tasks; both of these automatically downloadable for ease of use
+1. 11 databases with a total of 66 tasks; both of these automatically downloadable for ease of use
 2. Easy data loading, and graph construction from pkey-fkey links
 3. Your own model, which can use any deep learning stack since RelBench is framework-agnostic. We provide a first model implementation using PyTorch Geometric and PyTorch Frame.
 4. Standardized evaluators - all you need to do is produce a list of predictions for test samples, and RelBench computes metrics to ensure standardized evaluation
@@ -82,7 +88,30 @@ pip install relbench
 
 This will allow usage of the core RelBench data and task loading functionality.
 
-<details markdown="1"><summary>Including CTU datasets</summary>
+
+To additionally use `relbench.modeling`, which requires [PyTorch](https://pytorch.org/), [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric) and [PyTorch Frame](https://github.com/pyg-team/pytorch-frame), install these dependencies manually or do:
+
+```bash
+pip install relbench[full]
+```
+
+
+For the scripts in the `examples` directory, use:
+```bash
+pip install relbench[example]
+```
+
+Then, to run a script:
+```bash
+git clone https://github.com/snap-stanford/relbench
+cd relbench/examples
+python gnn_entity.py --dataset rel-f1 --task driver-position
+```
+
+
+## Using External Integrations
+
+**Using CTU datasets**
 
 To use datasets from the CTU repository, use:
 ```bash
@@ -91,7 +120,7 @@ pip install relbench[ctu]
 
 If you use the CTU datasets in your work, please cite [ReDeLEx](https://github.com/jakubpeleska/redelex) as below:
 
-```
+```bibtex
 @misc{peleska2025redelex,
   title={REDELEX: A Framework for Relational Deep Learning Exploration},
   author={Jakub Peleška and Gustav Šír},
@@ -103,9 +132,9 @@ If you use the CTU datasets in your work, please cite [ReDeLEx](https://github.c
 }
 ```
 
-</details>
 
-<details markdown="1"><summary>Including 4DBInfer datasets</summary>
+
+**Using 4DBInfer datasets**
 
 To use datasets from the 4DBInfer repository, use:
 ```bash
@@ -113,7 +142,7 @@ pip install relbench[dbinfer]
 ```
 
 If you use the 4DBInfer datasets in your work, please cite [4DBInfer](https://github.com/awslabs/multi-table-benchmark) as below:
-```
+```bibtex
 @article{dbinfer,
   title={4DBInfer: A 4D Benchmarking Toolbox for Graph-Centric Predictive Modeling on Relational DBs},
   author={Wang, Minjie and Gan, Quan and Wipf, David and Cai, Zhenkun and Li, Ning and Tang, Jianheng and Zhang, Yanlin and Zhang, Zizhao and Mao, Zunyao and Song, Yakun and Wang, Yanbo and Li, Jiahang and Zhang, Han and Yang, Guang and Qin, Xiao and Lei, Chuan and Zhang, Muhan and Zhang, Weinan and Faloutsos, Christos and Zhang, Zheng},
@@ -122,28 +151,29 @@ If you use the 4DBInfer datasets in your work, please cite [4DBInfer](https://gi
 }
 ```
 
-</details>
-<br>
+**Using TGB datasets**
 
+[TGB](https://tgb.complexdatalab.com/) datasets and tasks can be loaded in the RelBench format, and they cover bipartite link prediction (`tgbl-*`), heterogeneous link prediction (`thgl-*`), and node property prediction (`tgbn-*`). See `relbench/tasks/tgb.py` for the official TGB evaluation protocols (one-vs-many MRR / Hits@10 for link prediction, NDCG@10 for node property prediction). Standard RelBench style evaluation works as it is.
 
-To additionally use `relbench.modeling`, which requires [PyTorch](https://pytorch.org/), [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric) and [PyTorch Frame](https://github.com/pyg-team/pytorch-frame), install these dependencies manually or do:
-
-```bash
-pip install relbench[full]
+If you use the TGB datasets in your work, please cite [TGB](https://tgb.complexdatalab.com/) as below:
+```bibtex
+@article{huang2023temporal,
+  title={Temporal graph benchmark for machine learning on temporal graphs},
+  author={Huang, Shenyang and Poursafaei, Farimah and Danovitch, Jacob and Fey, Matthias and Hu, Weihua and Rossi, Emanuele and Leskovec, Jure and Bronstein, Michael and Rabusseau, Guillaume and Rabbany, Reihaneh},
+  journal={Advances in Neural Information Processing Systems},
+  year={2023}
+}
 ```
 
-
-For the scripts in the `examples` directory, use:
-```bash
-pip install relbench[example]
+```bibtex
+@article{huang2024tgb2,
+  title={TGB 2.0: A Benchmark for Learning on Temporal Knowledge Graphs and Heterogeneous Graphs},
+  author={Gastinger, Julia and Huang, Shenyang and Galkin, Mikhail and Loghmani, Erfan and Parviz, Ali and Poursafaei, Farimah and Danovitch, Jacob and Rossi, Emanuele and Koutis, Ioannis and Stuckenschmidt, Heiner and      Rabbany, Reihaneh and Rabusseau, Guillaume},
+  journal={Advances in Neural Information Processing Systems},
+  year={2024}
+}
 ```
 
-Then, to run a script:
-```bash
-git clone https://github.com/snap-stanford/relbench
-cd relbench/examples
-python gnn_entity.py --dataset rel-f1 --task driver-position
-```
 
 
 # Package Usage
@@ -261,3 +291,16 @@ If you use RelBench in your work, please cite our position and benchmark papers:
       url={https://arxiv.org/abs/2407.20060},
 }
 ```
+
+If you use RelBench v2 in your work, please cite:
+```bibtex
+@misc{gu2026relbenchv2,
+      title={{RelBench} v2: A Large-Scale Benchmark and Repository for Relational Data},
+      author={Justin Gu and Rishabh Ranjan and Charilaos Kanatsoulis and Haiming Tang and Martin Jurkovic and Valter Hudovernik and Mark Znidar and Pranshu Chaturvedi and Parth Shroff and Fengyu Li and Jure Leskovec},
+      year={2026},
+      eprint={2602.12606},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2602.12606},
+}
+```

{relbench-2.0.2 → relbench-2.1.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "relbench"
-version = "2.0.2"
+version = "2.1.1"
 description = "RelBench: Relational Deep Learning Benchmark"
 authors = [{name = "RelBench Team", email = "relbench@cs.stanford.edu"}]
 readme = "README.md"
@@ -44,6 +44,7 @@ dev=[
 ]
 dbinfer=[
 	"dbinfer-relbench-adapter",
+	"db-dtypes",
 ]
 ctu=[
 	"redelex",

{relbench-2.0.2 → relbench-2.1.1}/relbench/base/table.py

@@ -84,6 +84,11 @@ class Table:
         # Read the Parquet file using pyarrow
         table = pa.parquet.read_table(path)
         df = table.to_pandas()
+        # DuckDB currently cannot ingest pandas StringDtype columns reliably.
+        # Convert them to plain object dtype on load.
+        for col in df.columns:
+            if isinstance(df[col].dtype, pd.StringDtype):
+                df[col] = df[col].astype(object)
 
         # Extract metadata
         metadata_bytes = table.schema.metadata

{relbench-2.0.2 → relbench-2.1.1}/relbench/base/task_recommendation.py

@@ -83,6 +83,9 @@ class RecommendationTask(BaseTask):
                 f"{pred.shape} given."
             )
 
+        if len(target_table) == 0:
+            return {fn.__name__: float("nan") for fn in metrics}
+
         pred_isin_list = []
         dst_count_list = []
         for true_dst_nodes, pred_dst_nodes in zip(

{relbench-2.0.2 → relbench-2.1.1}/relbench/datasets/__init__.py

@@ -18,6 +18,7 @@ from relbench.datasets import (
     ratebeer,
     salt,
     stack,
+    tgb,
     trial,
 )
 from relbench.utils import get_relbench_cache_dir
@@ -70,15 +71,8 @@ def download_dataset(name: str) -> None:
     The downloaded database will be automatically picked up by the dataset object, when
     `dataset.get_db()` is called.
     """
-    if name.startswith("dbinfer-"):
-        print(
-            f"Dataset '{name}' is derived from 4DBInfer and must be generated "
-            "locally; skipping download."
-        )
-        return
 
     if name == "rel-mimic":
-        print("Downloading Mimic dataset...")
         from relbench.datasets.mimic import verify_mimic_access
 
         verify_mimic_access()
@@ -126,6 +120,8 @@ def get_dataset(name: str, download=True) -> Dataset:
     if name == "rel-mimic":
         from relbench.datasets import mimic
 
+        mimic.verify_mimic_access()
+
         cls, args, kwargs = (
             mimic.MimicDataset,
             (),
@@ -150,10 +146,29 @@ register_dataset("rel-arxiv", arxiv.ArxivDataset)
 register_dataset("rel-salt", salt.SALTDataset)
 register_dataset("rel-ratebeer", ratebeer.RateBeerDataset)
 register_dataset("dbinfer-avs", dbinfer.DBInferAVSDataset)
-register_dataset("dbinfer-mag", dbinfer.DBInferMAGDataset)
+# register_dataset("dbinfer-mag", dbinfer.DBInferMAGDataset)    # Experimental dataset: dbinfer-mag
 register_dataset("dbinfer-diginetica", dbinfer.DBInferDigineticaDataset)
 register_dataset("dbinfer-retailrocket", dbinfer.DBInferRetailRocketDataset)
 register_dataset("dbinfer-seznam", dbinfer.DBInferSeznamDataset)
 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")

{relbench-2.0.2 → relbench-2.1.1}/relbench/datasets/dbinfer.py

@@ -38,7 +38,18 @@ class DBInferDatasetBase(Dataset):
         self._adapter_cache_dir = adapter_cache_dir
 
     def _load_dataset_adapter(self):
-        from dbinfer_relbench_adapter.loader import load_dbinfer_data
+        try:
+            from dbinfer_relbench_adapter.loader import load_dbinfer_data
+        except ImportError:
+            raise ImportError(
+                "dbinfer-relbench-adapter is not installed. "
+                "Please install it with:\n"
+                "  pip install dbinfer-relbench-adapter\n"
+                "  pip install dgl --force-reinstall -f https://data.dgl.ai/wheels/torch-2.3/cu121/repo.html\n"
+                "Note: dgl may conflict with other dependencies. We recommend building "
+                "dbinfer datasets in a separate virtual environment and using the cached "
+                "datasets in your main working environment."
+            ) from None
 
         dataset_adapter, _ = load_dbinfer_data(
             dataset_name=self.dbinfer_name,

relbench-2.1.1/relbench/datasets/hashes.json

@@ -0,0 +1,38 @@
+{
+  "rel-trial/db.zip": "9fb5ba14f7cbca8115f3dfe0800415f98d6ddc15561e56c35ee614da6b89552a",
+  "rel-avito/db.zip": "274e692295027a753063b9201815a9d2dea94d4cda968be81be936f546d77058",
+  "rel-stack/db.zip": "f1374fda2828f62ac98e360f969dbf76a8a6404f6691b0a27a521d0730e22ec3",
+  "rel-f1/db.zip": "ec31a4e1bc2b2f9c36c05fcd3dfe2a40a506f335dc51ce79c3ec8bb40feb1482",
+  "rel-hm/db.zip": "3fb2aa3c978cc388aa81f43f9404fed329b502104ff58f3ee435959ab95e7029",
+  "rel-event/db.zip": "2593477065e1090af0258c07769ec22f40dad01f035613ccf35258b924d8b066",
+  "rel-amazon/db.zip": "db71c7701b892a4eb7481ff04d14d25465795501dba3a5931aabee9930805efe",
+  "rel-salt/db.zip": "fca91ab7d9e37646dcf1cb0007cc4229e9b23ef3c85f3c9e578d0f3fcb167001",
+  "rel-mimic/db.zip": "b41c6fa39e9d37b8e1e66f8311eb44622c79ea1644ee2cce280fc731d8366bbb",
+  "rel-arxiv/db.zip": "b7db1bc1aa17a6d79e0cf102fe3f258920a610f58ad8f8891169676fc120bc5c",
+  "rel-ratebeer/db.zip": "49afdf23a8b464dc0e54731a3782b51e53f3acee93dae0aa6d968cfbd9caa4be",
+  "dbinfer-amazon/db.zip": "0d9291bf880b5afa232b711a462ed7522a94a0ef5c141699ee8a2687dbbf0d08",
+  "dbinfer-avs/db.zip": "84295768a283c3c9ee0dbb25eb1c47d6d3f2c1956b523e7361f086fcf6664a93",
+  "dbinfer-diginetica/db.zip": "396e2aeb88cc616672b3b0d88a02dc13829c81b855edc3e668f7f2dbad3b96f0",
+  "dbinfer-outbrain-small/db.zip": "d186a71fc534bcac4299c616110805a997fa56fd16c15110d02e1c8eb3975210",
+  "dbinfer-retailrocket/db.zip": "6dbe83488e11c0d159b4592aa4cff57f6fa22b9f4bbf4ae38d0388d536897d75",
+  "dbinfer-seznam/db.zip": "77314bf874dc495e8a4a61f2dc5f12982bbec3c5b6b7af5555e9a2bb587154d9",
+  "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"
+}

{relbench-2.0.2 → relbench-2.1.1}/relbench/datasets/mimic.py

@@ -26,6 +26,15 @@ def verify_mimic_access() -> None:
     except Exception as e:
         raise RuntimeError(
             f"\nACCESS FAILED - BigQuery credential check encountered an error: {e}"
+            "\n\nTo gain access to MIMIC-IV via RelBench:"
+            "\n  1. Request access at https://physionet.org/content/mimiciv/"
+            "\n     and complete the 'Request access using Google BigQuery' step."
+            "\n  2. Install the Google Cloud SDK: https://cloud.google.com/sdk/docs/install"
+            "\n  3. Authenticate:  gcloud auth application-default login"
+            "\n  4. Set quota project:  gcloud auth application-default set-quota-project YOUR_PROJECT_ID"
+            "\n  5. Export your project ID:  export PROJECT_ID=YOUR_PROJECT_ID"
+            "\n  6. If credentials are not at the default path, set:"
+            "\n       export GOOGLE_APPLICATION_CREDENTIALS=/path/to/application_default_credentials.json"
         )
 
 
@@ -74,7 +83,9 @@ def convert_dtypes(df: pd.DataFrame) -> pd.DataFrame:
             or "time" in col.lower()
         ):
             try:
-                df[col] = pd.to_datetime(df[col], errors="coerce")
+                df[col] = pd.to_datetime(df[col], errors="coerce").astype(
+                    "datetime64[ns]"
+                )
             except Exception:
                 print("WARNING ")
                 pass
@@ -111,6 +122,10 @@ class MimicDataset(Dataset):
     The dataset is anchored around ICU stays (`icustays`) and patients (`patients`) and constructs all other
     tables with respect to these entities.
 
+    Note: MIMIC-IV applies date shifting for patient privacy. Timestamps are offset by a
+    random amount (typically decades) while preserving relative ordering. Expect dates
+    in the 2110-2200 range rather than the original 2008-2012 collection period.
+
     If certain parameters are not provided, default values will be used.
 
     Parameters:

{relbench-2.0.2 → relbench-2.1.1}/relbench/datasets/stack.py

@@ -51,6 +51,7 @@ class StackDataset(Dataset):
                 "Score",
                 "LastEditorDisplayName",
                 "LastEditorUserId",
+                "AcceptedAnswerId",
             ],
             inplace=True,
         )
@@ -123,7 +124,6 @@ class StackDataset(Dataset):
             fkey_col_to_pkey_table={
                 "OwnerUserId": "users",
                 "ParentId": "posts",  # notice the self-reference
-                "AcceptedAnswerId": "posts",
             },
             pkey_col="Id",
             time_col="CreationDate",

relbench-2.1.1/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."
+        )