RelBench: Relational Deep Learning Benchmark

RelBench uses make_pkey_fkey_graph to transform a relational database into a HeteroData object. This process involves:

1. Determining Column Types: Use get_stype_proposal(db) to automatically detect the stype (modality) for each column. Always verify this dictionary for accuracy.
2. Configuring Text Encoders: If your data contains text, provide a TextEmbedderConfig containing a text embedding model (e.g., using sentence-transformers).
3. Materializing the Graph: Call make_pkey_fkey_graph with the database, column types, and text encoder configuration. You can specify a cache_dir to store the materialized graph for faster subsequent loads.

from relbench.modeling.utils import get_stype_proposal
from relbench.modeling.graph import make_pkey_fkey_graph
from torch_frame.config.text_embedder import TextEmbedderConfig

# 1. Get column types
db = dataset.get_db()
col_to_stype_dict = get_stype_proposal(db)

# 2. Setup text embedding (example using SentenceTransformer)
# ... define GloveTextEmbedding class ...
text_embedder_cfg = TextEmbedderConfig(
    text_embedder=GloveTextEmbedding(device=device), 
    batch_size=256
)

# 3. Create the graph
data, col_stats_dict = make_pkey_fkey_graph(
    db,
    col_to_stype_dict=col_to_stype_dict,
    text_embedder_cfg=text_embedder_cfg,
    cache_dir="./data/rel-f1_materialized_cache"
)

The make_db function is responsible for preprocessing raw data (from URLs or local files) into the RelBench format. It should return a relbench.base.Database object, which is a collection of relbench.base.Table objects.

Creating a Table object

Each table in the database is instantiated using the relbench.base.Table class with the following parameters:

• df: A pd.DataFrame containing the table data.
• fkey_col_to_pkey_table: A dictionary mapping foreign-key columns to their corresponding primary-key table names.
• pkey_col: The name of the primary key column (or None).
• time_col: The name of the column representing the row's creation time (or None). If None, the row is treated as created at -inf.

Pkey/Fkey Reindexing

When you use dataset.get_db(), RelBench automatically calls db.reindex_pkeys_and_fkeys(). This converts primary and foreign keys into consecutive integers starting from 0. To preserve original ID values for features or cross-referencing, add a duplicate column that is not marked as the pkey_col.

from relbench.base import Database, Table
import pandas as pd

# Inside make_db...
tables = {
    "users": Table(
        df=users_df,
        fkey_col_to_pkey_table={},
        pkey_col="user_id",
        time_col="created_at"
    ),
    "orders": Table(
        df=orders_df,
        fkey_col_to_pkey_table={"user_id": "users"},
        pkey_col="order_id",
        time_col="order_date"
    )
}
return Database(tables)

Once a Dataset subclass is defined, you can instantiate it directly or register it for use with the global get_dataset API.

Direct Usage

Instantiate the class with a cache_dir. Use .get_db() to retrieve the database with temporal splitting applied, or .make_db() to retrieve the full database (useful for debugging).

Registration

Use register_dataset to make your dataset available via get_dataset. Note that the registry is not persistent across Python processes; you must re-register the dataset every time you start a new session (e.g., when running baseline scripts).

Note: When registering, you can pass additional args and kwargs that will be used during the dataset's initialization.

# Direct usage
my_dataset = MyCustomDataset(cache_dir="./cache")
db = my_dataset.get_db()

# Registration
from relbench.datasets import register_dataset, get_dataset

register_dataset("my-custom-id", MyCustomDataset, arg1="value")
new_dataset = get_dataset("my-custom-id")

To define a custom dataset in RelBench, you must subclass relbench.base.Dataset. A valid subclass requires implementing three components:

1. val_timestamp: A pd.Timestamp defining the validation split point.
2. test_timestamp: A pd.Timestamp defining the test split point.
3. make_db function: A method that returns a relbench.base.Database object.

Temporal Splitting Logic

RelBench uses these timestamps to prevent temporal leakage:

• Validation Set: Only database rows up to val_timestamp are used for prediction. Ground truth labels for the validation set can only use rows up to test_timestamp.
• Test Set: Only database rows up to test_timestamp are used for prediction. Ground truth labels for the test set use rows after test_timestamp.
• Training Set: Ground truth labels for the training set can only use rows up to val_timestamp.

from relbench.base import Dataset
import pandas as pd

class MyCustomDataset(Dataset):
    val_timestamp = pd.Timestamp("2023-01-01")
    test_timestamp = pd.Timestamp("2024-01-01")

    def make_db(self) -> Database:
        # Implementation logic here
        pass