rllm.datasets.RelBenchDataset

class rllm.datasets.RelBenchDataset(root: str, force_reload: bool | None = False)

Bases: Dataset

Override methods for RelBench datasets.

Subclasses need to assign the following properties after processing:

self._task_dict: Dict[str, RelBenchTask] self._table_dict: Dict[str, TableData] self._hdata: HeteroGraphData self._tabledata_stats_dict: Dict[str, Any] self._table_meta_dict: Dict[str, RelBenchTableMeta]

download()

Download and unzip raw files.

property has_download

check whether data has been downloaded

property has_process

check whether data has been processed

load_all()

Force load all cached properties.

make_pkey_fkey_graph() → Tuple[HeteroGraphData, Dict]

Make primary key - foreign key graph for the dataset.

This method lazy materializes each TableData, saves them to processed_dir, and constructs the HeteroGraphData based on pkey-fkey relations.

Returns:

Heterogeneous graph data. Dict: table_name -> TableData.metadata

Return type:

HeteroGraphData

process()

process the datasets to self.processed_dir

property processed_filenames

file names in the self.processed_dir

property raw_filenames

file names in the self.raw_dir

validate_dataset()

Validate the integrity of downloaded files. 1. validate primary keys 2. validate foreign keys (correct if necessary)