tutorials/custom_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a85f346-d7a7-4c32-9c30-3351e964e293",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install relbench\n",
    "\n",
    "import relbench\n",
    "\n",
    "relbench.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "098e9452-03b7-4d33-acf2-72d507f3f96d",
   "metadata": {},
   "source": [
    "You can work with your own data in the RelBench framework. This tutorial walks you through the steps to create custom datasets. Code in this notebook has been adapted from `relbench/datasets/f1.py` and `relbench/datasets/amazon.py`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e840a5f-a6e4-4aff-b6b8-4a777dc2045a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-22T22:51:35.258604Z",
     "iopub.status.busy": "2024-07-22T22:51:35.258167Z",
     "iopub.status.idle": "2024-07-22T22:51:35.269009Z",
     "shell.execute_reply": "2024-07-22T22:51:35.268587Z",
     "shell.execute_reply.started": "2024-07-22T22:51:35.258581Z"
    }
   },
   "source": [
    "# Custom Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cd3d3e7-e810-483e-beeb-e74c66c8ad48",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e78e8039-1bb8-4265-ac9e-65520deab3a5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:19.877872Z",
     "iopub.status.busy": "2024-07-26T00:25:19.877382Z",
     "iopub.status.idle": "2024-07-26T00:25:20.526819Z",
     "shell.execute_reply": "2024-07-26T00:25:20.525898Z",
     "shell.execute_reply.started": "2024-07-26T00:25:19.877832Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pooch\n",
    "import pyarrow as pa\n",
    "import pyarrow.json\n",
    "\n",
    "from relbench.base import Database, Dataset, Table\n",
    "from relbench.datasets import get_dataset, get_dataset_names, register_dataset\n",
    "from relbench.utils import unzip_processor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dff46a3c-a202-4f27-8533-1dfdc6d41dbf",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10240d58-21f8-42ba-a936-031114aec2fc",
   "metadata": {},
   "source": [
    "To define a custom dataset, we subclass the `relbench.base.Dataset` class. This requires specifying 3 things:\n",
    "1. `val_timestamp` of type `pd.Timestamp`\n",
    "2. `test_timestamp` of type `pd.Timestamp`\n",
    "3. a `make_db` function which returns a `relbench.base.Database` object\n",
    "\n",
    "These are described in further detail below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee17a117-c8f0-4f43-9f9d-7b41c0886acb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-22T23:11:18.430211Z",
     "iopub.status.busy": "2024-07-22T23:11:18.429861Z",
     "iopub.status.idle": "2024-07-22T23:11:18.446220Z",
     "shell.execute_reply": "2024-07-22T23:11:18.445784Z",
     "shell.execute_reply.started": "2024-07-22T23:11:18.430194Z"
    }
   },
   "source": [
    "### Temporal splitting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7deac8d1-bdf1-4f96-9cb2-93cd665ca71d",
   "metadata": {},
   "source": [
    "`val_timestamp` and `test_timestamp` define a unified temporal splitting for the dataset. All tasks defined on the dataset will provide train/val/test sets based on this splitting."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00e29e72-420f-481b-a986-e2369fc3c88a",
   "metadata": {},
   "source": [
    "Only database rows upto `val_timestamp` should be used by the model to predict over the val set. Similarly, only database rows upto `test_timestamp` should be used by the model to predict over the test set. Rows after `test_timestamp` are only used for computing ground truth labels for the test set. Importantly, only rows upto `val_timestamp` can be used to obtain ground truth labels for the train set; and only rows upto `test_timestamp` can be used to obtain ground truth labels for the val set."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "533a2345-bf6f-4922-8454-446af46b6e3f",
   "metadata": {},
   "source": [
    "This is important to prevent temporal leakage of information, which we will see is an important consideration in many aspects of Relational Deep Learning."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0714ecc7-7eb6-41da-a1ce-8ddfa868ac07",
   "metadata": {},
   "source": [
    "### The make_db function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b47dcff3-60ab-43b7-80c3-a61e45493b40",
   "metadata": {},
   "source": [
    "The raw data for your dataset can be in any form. You will need to preprocess your dataset into the RelBench format to work with it in RelBench. The `make_db` function is the place to do this."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ddb7d3e-55e0-49f2-ac41-43f9454b4f3e",
   "metadata": {},
   "source": [
    "Inside the `make_db` function we first download the raw files (or read from the local filesystem) and then create a `relbench.base.Database` object out of those. Thus, the `make_db` functions serves as documentation for your pre-processing steps, while also conveniently allowing you to develop and debug them within the RelBench framework."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9552931-ef4a-4c2b-b63e-6ec56738db31",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-08-12T05:14:46.141436Z",
     "iopub.status.busy": "2024-08-12T05:14:46.140960Z",
     "iopub.status.idle": "2024-08-12T05:14:46.160386Z",
     "shell.execute_reply": "2024-08-12T05:14:46.159281Z",
     "shell.execute_reply.started": "2024-08-12T05:14:46.141396Z"
    }
   },
   "source": [
    "#### Pkey/Fkey Reindexing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9845d78-7943-47be-a1ce-da5b0b01a3d4",
   "metadata": {},
   "source": [
    "The intended usage is not to call the `make_db` function directly but to use the `get_db` function which internally calls `make_db` and adds a layer of other functionality such as caching."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c67714b3-f3f0-4841-b66f-3623242ff033",
   "metadata": {},
   "source": [
    "Another important thing that `get_db` does is that it calls `db.reindex_pkeys_and_fkeys()` on the database `db` returned by `make_db`. This reindexes the primary- and foreign- key columns so that the primary keys columns are consecutive integers starting from 0. This makes some downstream logic in RelBench convenient to implement, as it can work under the unified assumption that the pkeys and fkeys are integers, and that too sequential."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcc2dc99-128f-43dd-bb31-20752ee469d8",
   "metadata": {},
   "source": [
    "If you want to preserve the original pkey values, either because you believe they can be used as features for predictive tasks, or because you would like to cross-reference the prediction results with the original data source, simply add a duplicate column without marking it as pkey_col. The model designer is free to decide whether to include this duplicate column as input to the model or not."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fec0b02e-ce3e-426b-8fe2-38b8fa8d20c8",
   "metadata": {},
   "source": [
    "### The Database object"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "241e9f15-5dd4-4f31-be2f-28382aa8cf4b",
   "metadata": {},
   "source": [
    "The `relbench.base.Database` object is simply a collection of named `relbench.base.Table` objects. A `relbench.base.Table` object is instantiated by providing:\n",
    "1. `df` of type `pd.DataFrame` representing the table content.\n",
    "2. `fkey_col_to_pkey_table` dict to map foreign-key columns in this table to their primary-key table names.\n",
    "3. `pkey_col`, the name of the primary key column, if any (else `None`).\n",
    "4. `time_col`, the name of the time column, if available (else `None`)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53711365-a256-4d78-a8f4-1b06ac0d7600",
   "metadata": {},
   "source": [
    "The `time_col` denotes the creation time of a row in the database. If absent or `None`, the row is treated as if it was created at time `-inf`. Note that there can be other columns of type `pd.Timestamp` which are not necessarily the creation time. For example, there can be a `CreationDate` column in a `Users` table which would be suitable for `time_col`, but also a `DateOfBirth` column, which is better treated as an ordinary feature column."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d382e5d-8e33-45b0-a676-7b2fe9e2f1cd",
   "metadata": {},
   "source": [
    "## Annotated Sample Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "165d2698-ae4c-44f9-be75-24cf55028bde",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:22.150312Z",
     "iopub.status.busy": "2024-07-26T00:25:22.149838Z",
     "iopub.status.idle": "2024-07-26T00:25:22.180445Z",
     "shell.execute_reply": "2024-07-26T00:25:22.179540Z",
     "shell.execute_reply.started": "2024-07-26T00:25:22.150280Z"
    }
   },
   "outputs": [],
   "source": [
    "class F1Dataset(Dataset):\n",
    "    ################################################################################\n",
    "    # Choose the val_timestamp and test_timestamp carefully\n",
    "    ################################################################################\n",
    "    val_timestamp = pd.Timestamp(\"2005-01-01\")\n",
    "    test_timestamp = pd.Timestamp(\"2010-01-01\")\n",
    "\n",
    "    def make_db(self) -> Database:\n",
    "        r\"\"\"Process the raw files into a database.\"\"\"\n",
    "        ################################################################################\n",
    "        # The raw files are at this URL. You can use any URL or even local files.\n",
    "        ################################################################################\n",
    "        url = \"https://relbench.stanford.edu/data/relbench-f1-raw.zip\"\n",
    "\n",
    "        path = pooch.retrieve(\n",
    "            url,\n",
    "            known_hash=\"2933348953b30aa9723b4831fea8071b336b74977bbcf1fb059da63a04f06eba\",\n",
    "            progressbar=True,\n",
    "            processor=unzip_processor,\n",
    "        )\n",
    "\n",
    "        path = os.path.join(path, \"raw\")\n",
    "\n",
    "        ################################################################################\n",
    "        # Here, we read from the raw CSV files\n",
    "        ################################################################################\n",
    "        circuits = pd.read_csv(os.path.join(path, \"circuits.csv\"))\n",
    "        drivers = pd.read_csv(os.path.join(path, \"drivers.csv\"))\n",
    "        results = pd.read_csv(os.path.join(path, \"results.csv\"))\n",
    "        races = pd.read_csv(os.path.join(path, \"races.csv\"))\n",
    "        standings = pd.read_csv(os.path.join(path, \"driver_standings.csv\"))\n",
    "        constructors = pd.read_csv(os.path.join(path, \"constructors.csv\"))\n",
    "        constructor_results = pd.read_csv(os.path.join(path, \"constructor_results.csv\"))\n",
    "        constructor_standings = pd.read_csv(\n",
    "            os.path.join(path, \"constructor_standings.csv\")\n",
    "        )\n",
    "        qualifying = pd.read_csv(os.path.join(path, \"qualifying.csv\"))\n",
    "\n",
    "        ################################################################################\n",
    "        # It is important to understand the data.\n",
    "        # This can point to columns which should be removed.\n",
    "        # The most important of these is temporal leakage columns, which we will\n",
    "        # discuss in detail later.\n",
    "        ################################################################################\n",
    "\n",
    "        # Remove columns that are irrelevant, leak time,\n",
    "        # or have too many missing values\n",
    "\n",
    "        # Drop the Wikipedia URL and some time columns with many missing values\n",
    "        races.drop(\n",
    "            columns=[\n",
    "                \"url\",\n",
    "                \"fp1_date\",\n",
    "                \"fp1_time\",\n",
    "                \"fp2_date\",\n",
    "                \"fp2_time\",\n",
    "                \"fp3_date\",\n",
    "                \"fp3_time\",\n",
    "                \"quali_date\",\n",
    "                \"quali_time\",\n",
    "                \"sprint_date\",\n",
    "                \"sprint_time\",\n",
    "            ],\n",
    "            inplace=True,\n",
    "        )\n",
    "\n",
    "        # Drop the Wikipedia URL as it is unique for each row\n",
    "        circuits.drop(\n",
    "            columns=[\"url\"],\n",
    "            inplace=True,\n",
    "        )\n",
    "\n",
    "        # Drop the Wikipedia URL (unique) and number (803 / 857 are nulls)\n",
    "        drivers.drop(\n",
    "            columns=[\"number\", \"url\"],\n",
    "            inplace=True,\n",
    "        )\n",
    "\n",
    "        # Drop the positionText, time, fastestLapTime and fastestLapSpeed\n",
    "        results.drop(\n",
    "            columns=[\n",
    "                \"positionText\",\n",
    "                \"time\",\n",
    "                \"fastestLapTime\",\n",
    "                \"fastestLapSpeed\",\n",
    "            ],\n",
    "            inplace=True,\n",
    "        )\n",
    "\n",
    "        # Drop the positionText\n",
    "        standings.drop(\n",
    "            columns=[\"positionText\"],\n",
    "            inplace=True,\n",
    "        )\n",
    "\n",
    "        # Drop the Wikipedia URL\n",
    "        constructors.drop(\n",
    "            columns=[\"url\"],\n",
    "            inplace=True,\n",
    "        )\n",
    "\n",
    "        # Drop the positionText\n",
    "        constructor_standings.drop(\n",
    "            columns=[\"positionText\"],\n",
    "            inplace=True,\n",
    "        )\n",
    "\n",
    "        # Drop the status as it only contains two categories, and\n",
    "        # only 17 rows have value 'D' (0.138%)\n",
    "        constructor_results.drop(\n",
    "            columns=[\"status\"],\n",
    "            inplace=True,\n",
    "        )\n",
    "\n",
    "        # Drop the time in qualifying 1, 2, and 3\n",
    "        qualifying.drop(\n",
    "            columns=[\"q1\", \"q2\", \"q3\"],\n",
    "            inplace=True,\n",
    "        )\n",
    "\n",
    "        ################################################################################\n",
    "        # Make sure to properly process time columns into pd.Timestamp datatype.\n",
    "        # Sometimes, you might need to handle timezone information carefully to do\n",
    "        # this correctly.\n",
    "        # If time information can be inferred for other tables, it might help to add\n",
    "        # as it makes temporal sampling more effective.\n",
    "        ################################################################################\n",
    "\n",
    "        # replase missing data and combine date and time columns\n",
    "        races[\"time\"] = races[\"time\"].replace(r\"^\\\\N$\", \"00:00:00\", regex=True)\n",
    "        races[\"date\"] = races[\"date\"] + \" \" + races[\"time\"]\n",
    "        # Convert date column to pd.Timestamp\n",
    "        races[\"date\"] = pd.to_datetime(races[\"date\"])\n",
    "\n",
    "        # add time column to other tables\n",
    "        results = results.merge(races[[\"raceId\", \"date\"]], on=\"raceId\", how=\"left\")\n",
    "        standings = standings.merge(races[[\"raceId\", \"date\"]], on=\"raceId\", how=\"left\")\n",
    "        constructor_results = constructor_results.merge(\n",
    "            races[[\"raceId\", \"date\"]], on=\"raceId\", how=\"left\"\n",
    "        )\n",
    "        constructor_standings = constructor_standings.merge(\n",
    "            races[[\"raceId\", \"date\"]], on=\"raceId\", how=\"left\"\n",
    "        )\n",
    "\n",
    "        qualifying = qualifying.merge(\n",
    "            races[[\"raceId\", \"date\"]], on=\"raceId\", how=\"left\"\n",
    "        )\n",
    "\n",
    "        # Subtract a day from the date to account for the fact\n",
    "        # that the qualifying time is the day before the main race\n",
    "        qualifying[\"date\"] = qualifying[\"date\"] - pd.Timedelta(days=1)\n",
    "\n",
    "        ################################################################################\n",
    "        # Make sure that the missing data has been parsed properly.\n",
    "        # Following Pandas, we represent missing values with NaNs in the dataframe.\n",
    "        ################################################################################\n",
    "\n",
    "        # Replace \"\\N\" with NaN in results tables\n",
    "        results = results.replace(r\"^\\\\N$\", np.nan, regex=True)\n",
    "\n",
    "        # Replace \"\\N\" with NaN in circuits tables, especially\n",
    "        # for the column `alt` which has 3 rows of \"\\N\"\n",
    "        circuits = circuits.replace(r\"^\\\\N$\", np.nan, regex=True)\n",
    "        # Convert alt from string to float\n",
    "        circuits[\"alt\"] = circuits[\"alt\"].astype(float)\n",
    "\n",
    "        # Convert non-numeric values to NaN in the specified column\n",
    "        results[\"rank\"] = pd.to_numeric(results[\"rank\"], errors=\"coerce\")\n",
    "        results[\"number\"] = pd.to_numeric(results[\"number\"], errors=\"coerce\")\n",
    "        results[\"grid\"] = pd.to_numeric(results[\"grid\"], errors=\"coerce\")\n",
    "        results[\"position\"] = pd.to_numeric(results[\"position\"], errors=\"coerce\")\n",
    "        results[\"points\"] = pd.to_numeric(results[\"points\"], errors=\"coerce\")\n",
    "        results[\"laps\"] = pd.to_numeric(results[\"laps\"], errors=\"coerce\")\n",
    "        results[\"milliseconds\"] = pd.to_numeric(\n",
    "            results[\"milliseconds\"], errors=\"coerce\"\n",
    "        )\n",
    "        results[\"fastestLap\"] = pd.to_numeric(results[\"fastestLap\"], errors=\"coerce\")\n",
    "\n",
    "        # Convert drivers date of birth to datetime\n",
    "        drivers[\"dob\"] = pd.to_datetime(drivers[\"dob\"])\n",
    "\n",
    "        ################################################################################\n",
    "        # Here, we collect all tables in the database as relbench.base.Table objects.\n",
    "        ################################################################################\n",
    "\n",
    "        tables = {}\n",
    "\n",
    "        tables[\"races\"] = Table(\n",
    "            df=pd.DataFrame(races),\n",
    "            fkey_col_to_pkey_table={\n",
    "                \"circuitId\": \"circuits\",\n",
    "            },\n",
    "            pkey_col=\"raceId\",\n",
    "            time_col=\"date\",\n",
    "        )\n",
    "\n",
    "        tables[\"circuits\"] = Table(\n",
    "            df=pd.DataFrame(circuits),\n",
    "            fkey_col_to_pkey_table={},\n",
    "            pkey_col=\"circuitId\",\n",
    "            time_col=None,\n",
    "        )\n",
    "\n",
    "        tables[\"drivers\"] = Table(\n",
    "            df=pd.DataFrame(drivers),\n",
    "            fkey_col_to_pkey_table={},\n",
    "            pkey_col=\"driverId\",\n",
    "            time_col=None,\n",
    "        )\n",
    "\n",
    "        tables[\"results\"] = Table(\n",
    "            df=pd.DataFrame(results),\n",
    "            fkey_col_to_pkey_table={\n",
    "                \"raceId\": \"races\",\n",
    "                \"driverId\": \"drivers\",\n",
    "                \"constructorId\": \"constructors\",\n",
    "            },\n",
    "            pkey_col=\"resultId\",\n",
    "            time_col=\"date\",\n",
    "        )\n",
    "\n",
    "        tables[\"standings\"] = Table(\n",
    "            df=pd.DataFrame(standings),\n",
    "            fkey_col_to_pkey_table={\"raceId\": \"races\", \"driverId\": \"drivers\"},\n",
    "            pkey_col=\"driverStandingsId\",\n",
    "            time_col=\"date\",\n",
    "        )\n",
    "\n",
    "        tables[\"constructors\"] = Table(\n",
    "            df=pd.DataFrame(constructors),\n",
    "            fkey_col_to_pkey_table={},\n",
    "            pkey_col=\"constructorId\",\n",
    "            time_col=None,\n",
    "        )\n",
    "\n",
    "        tables[\"constructor_results\"] = Table(\n",
    "            df=pd.DataFrame(constructor_results),\n",
    "            fkey_col_to_pkey_table={\"raceId\": \"races\", \"constructorId\": \"constructors\"},\n",
    "            pkey_col=\"constructorResultsId\",\n",
    "            time_col=\"date\",\n",
    "        )\n",
    "\n",
    "        tables[\"constructor_standings\"] = Table(\n",
    "            df=pd.DataFrame(constructor_standings),\n",
    "            fkey_col_to_pkey_table={\"raceId\": \"races\", \"constructorId\": \"constructors\"},\n",
    "            pkey_col=\"constructorStandingsId\",\n",
    "            time_col=\"date\",\n",
    "        )\n",
    "\n",
    "        tables[\"qualifying\"] = Table(\n",
    "            df=pd.DataFrame(qualifying),\n",
    "            fkey_col_to_pkey_table={\n",
    "                \"raceId\": \"races\",\n",
    "                \"driverId\": \"drivers\",\n",
    "                \"constructorId\": \"constructors\",\n",
    "            },\n",
    "            pkey_col=\"qualifyId\",\n",
    "            time_col=\"date\",\n",
    "        )\n",
    "\n",
    "        return Database(tables)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aac9d740-e2db-4d7f-b230-b5c06d892e29",
   "metadata": {},
   "source": [
    "## Using the custom dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31d15944-4cb0-4ce8-872d-349d2d654950",
   "metadata": {},
   "source": [
    "You can use the custom `Dataset` subclass directly as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "65c18113-7067-41cf-81da-bb159f17e581",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:23.545040Z",
     "iopub.status.busy": "2024-07-26T00:25:23.544759Z",
     "iopub.status.idle": "2024-07-26T00:25:23.575626Z",
     "shell.execute_reply": "2024-07-26T00:25:23.574925Z",
     "shell.execute_reply.started": "2024-07-26T00:25:23.545014Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "F1Dataset()"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f1_dataset = F1Dataset(cache_dir=\"./cache/f1\")\n",
    "f1_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7d9de641-13ae-442e-bfa0-08d5efa5a1c1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:24.277335Z",
     "iopub.status.busy": "2024-07-26T00:25:24.276960Z",
     "iopub.status.idle": "2024-07-26T00:25:24.603887Z",
     "shell.execute_reply": "2024-07-26T00:25:24.603241Z",
     "shell.execute_reply.started": "2024-07-26T00:25:24.277301Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Making Database object from scratch...\n",
      "(You can also use `get_dataset(..., download=True)` for datasets prepared by the RelBench team.)\n",
      "Done in 0.24 seconds.\n",
      "Caching Database object to ./cache/f1/db...\n",
      "Done in 0.05 seconds.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Database()"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f1_db = f1_dataset.get_db()\n",
    "f1_db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6ab7a185-fc58-48fd-a8a9-7b38b265acb3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:24.813771Z",
     "iopub.status.busy": "2024-07-26T00:25:24.813564Z",
     "iopub.status.idle": "2024-07-26T00:25:24.864693Z",
     "shell.execute_reply": "2024-07-26T00:25:24.863584Z",
     "shell.execute_reply.started": "2024-07-26T00:25:24.813754Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Table(df=\n",
       "     raceId  year  round  circuitId                  name                date  \\\n",
       "0         0  1950      1          8    British Grand Prix 1950-05-13 00:00:00   \n",
       "1         1  1950      2          5     Monaco Grand Prix 1950-05-21 00:00:00   \n",
       "2         2  1950      3         18      Indianapolis 500 1950-05-30 00:00:00   \n",
       "3         3  1950      4         65      Swiss Grand Prix 1950-06-04 00:00:00   \n",
       "4         4  1950      5         12    Belgian Grand Prix 1950-06-18 00:00:00   \n",
       "..      ...   ...    ...        ...                   ...                 ...   \n",
       "815     815  2009     13         13    Italian Grand Prix 2009-09-13 12:00:00   \n",
       "816     816  2009     14         14  Singapore Grand Prix 2009-09-27 12:00:00   \n",
       "817     817  2009     15         21   Japanese Grand Prix 2009-10-04 05:00:00   \n",
       "818     818  2009     16         17  Brazilian Grand Prix 2009-10-18 16:00:00   \n",
       "819     819  2009     17         23  Abu Dhabi Grand Prix 2009-11-01 11:00:00   \n",
       "\n",
       "         time  \n",
       "0    00:00:00  \n",
       "1    00:00:00  \n",
       "2    00:00:00  \n",
       "3    00:00:00  \n",
       "4    00:00:00  \n",
       "..        ...  \n",
       "815  12:00:00  \n",
       "816  12:00:00  \n",
       "817  05:00:00  \n",
       "818  16:00:00  \n",
       "819  11:00:00  \n",
       "\n",
       "[820 rows x 7 columns],\n",
       "  fkey_col_to_pkey_table={'circuitId': 'circuits'},\n",
       "  pkey_col=raceId,\n",
       "  time_col=date)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f1_db.table_dict[\"races\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fd40abb-c7b4-42b3-af69-78d10f01b12b",
   "metadata": {},
   "source": [
    "### Development advice"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "818a1fd5-0c84-4eba-b3d6-7b9f2da55d45",
   "metadata": {},
   "source": [
    "While developing `make_db` code, it is suggested to call `make_db` directly to avoid caching artifacts while debugging. Note that this will give the full database, unlike `get_db()` which removes rows after `test_timestamp`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "53f7fa8e-ec89-4bc5-b8ef-1037e0c25917",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:26.487123Z",
     "iopub.status.busy": "2024-07-26T00:25:26.486671Z",
     "iopub.status.idle": "2024-07-26T00:25:26.720974Z",
     "shell.execute_reply": "2024-07-26T00:25:26.719924Z",
     "shell.execute_reply.started": "2024-07-26T00:25:26.487085Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Database()"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f1_full_db = f1_dataset.make_db()\n",
    "f1_full_db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "942649f9-4527-4dba-817f-52eb44851189",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:27.079812Z",
     "iopub.status.busy": "2024-07-26T00:25:27.079571Z",
     "iopub.status.idle": "2024-07-26T00:25:27.106195Z",
     "shell.execute_reply": "2024-07-26T00:25:27.105351Z",
     "shell.execute_reply.started": "2024-07-26T00:25:27.079795Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Table(df=\n",
       "      raceId  year  round  circuitId                      name  \\\n",
       "0          1  2009      1          1     Australian Grand Prix   \n",
       "1          2  2009      2          2      Malaysian Grand Prix   \n",
       "2          3  2009      3         17        Chinese Grand Prix   \n",
       "3          4  2009      4          3        Bahrain Grand Prix   \n",
       "4          5  2009      5          4        Spanish Grand Prix   \n",
       "...      ...   ...    ...        ...                       ...   \n",
       "1096    1116  2023     18         69  United States Grand Prix   \n",
       "1097    1117  2023     19         32    Mexico City Grand Prix   \n",
       "1098    1118  2023     20         18      São Paulo Grand Prix   \n",
       "1099    1119  2023     21         80      Las Vegas Grand Prix   \n",
       "1100    1120  2023     22         24      Abu Dhabi Grand Prix   \n",
       "\n",
       "                    date      time  \n",
       "0    2009-03-29 06:00:00  06:00:00  \n",
       "1    2009-04-05 09:00:00  09:00:00  \n",
       "2    2009-04-19 07:00:00  07:00:00  \n",
       "3    2009-04-26 12:00:00  12:00:00  \n",
       "4    2009-05-10 12:00:00  12:00:00  \n",
       "...                  ...       ...  \n",
       "1096 2023-10-22 19:00:00  19:00:00  \n",
       "1097 2023-10-29 20:00:00  20:00:00  \n",
       "1098 2023-11-05 17:00:00  17:00:00  \n",
       "1099 2023-11-19 06:00:00  06:00:00  \n",
       "1100 2023-11-26 13:00:00  13:00:00  \n",
       "\n",
       "[1101 rows x 7 columns],\n",
       "  fkey_col_to_pkey_table={'circuitId': 'circuits'},\n",
       "  pkey_col=raceId,\n",
       "  time_col=date)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f1_full_db.table_dict[\"races\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab7ee822-8fa0-4cfe-98d4-5040d289967d",
   "metadata": {},
   "source": [
    "### Registering your custom dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cb1297b-f06e-4317-a27c-49abcc0f3d0c",
   "metadata": {},
   "source": [
    "You can also register your dataset to make it available to `relbench.datasets.get_dataset` and use standardized caching locations (`~/.cache/relbench/<dataset-name>`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3df1be1a-3024-4cec-9ef7-17b645af6aa5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:28.421274Z",
     "iopub.status.busy": "2024-07-26T00:25:28.420889Z",
     "iopub.status.idle": "2024-07-26T00:25:28.448827Z",
     "shell.execute_reply": "2024-07-26T00:25:28.447787Z",
     "shell.execute_reply.started": "2024-07-26T00:25:28.421239Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['rel-amazon',\n",
       " 'rel-avito',\n",
       " 'rel-event',\n",
       " 'rel-f1',\n",
       " 'rel-hm',\n",
       " 'rel-stack',\n",
       " 'rel-trial',\n",
       " 'rel-custom_f1']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "register_dataset(\"rel-custom_f1\", F1Dataset)\n",
    "get_dataset_names()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ca0fe893-b439-4f7f-aff1-f7e19200ecf9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:28.943803Z",
     "iopub.status.busy": "2024-07-26T00:25:28.943378Z",
     "iopub.status.idle": "2024-07-26T00:25:28.973952Z",
     "shell.execute_reply": "2024-07-26T00:25:28.973012Z",
     "shell.execute_reply.started": "2024-07-26T00:25:28.943766Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "F1Dataset()"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reg_f1_dataset = get_dataset(\"rel-custom_f1\")\n",
    "reg_f1_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c95941be-00f4-4aef-b532-a58f2c7f76a9",
   "metadata": {},
   "source": [
    "## Advanced Customization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "655c3a2f-803b-4c76-92f5-3b2235c0bc97",
   "metadata": {},
   "source": [
    "You can also add an `__init__` function to the `Dataset` subclass. This can allow customizing the returned `Dataset` object with `args` and `kwargs`. An example of this can be seen in `relbench/datasets/amazon.py`, where the same `AmazonDataset` class can be used for the Books subset of the raw Amazon Reviews dataset or the Fashion subset (this can easily be extended to the other subsets too)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15912e34-5412-40fb-9107-8d093836e92c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-25T04:53:22.440077Z",
     "iopub.status.busy": "2024-07-25T04:53:22.439629Z",
     "iopub.status.idle": "2024-07-25T04:53:22.499489Z",
     "shell.execute_reply": "2024-07-25T04:53:22.498485Z",
     "shell.execute_reply.started": "2024-07-25T04:53:22.440041Z"
    }
   },
   "source": [
    "We reproduce the code snippet below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "851458c3-c9bf-48b5-a09e-f5a802ef2e62",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:30.982702Z",
     "iopub.status.busy": "2024-07-26T00:25:30.982290Z",
     "iopub.status.idle": "2024-07-26T00:25:31.037729Z",
     "shell.execute_reply": "2024-07-26T00:25:31.036603Z",
     "shell.execute_reply.started": "2024-07-26T00:25:30.982666Z"
    }
   },
   "outputs": [],
   "source": [
    "class AmazonDataset(Dataset):\n",
    "    val_timestamp = pd.Timestamp(\"2015-10-01\")\n",
    "    test_timestamp = pd.Timestamp(\"2016-01-01\")\n",
    "\n",
    "    url_prefix = \"https://datarepo.eng.ucsd.edu/mcauley_group/data/amazon_v2\"\n",
    "    _category_to_url_key = {\"books\": \"Books\", \"fashion\": \"AMAZON_FASHION\"}\n",
    "\n",
    "    known_hashes = {\n",
    "        \"meta_Books.json.gz\": \"80ed7ac64f5967a140401e8d7bf0587d2e5087492de9e94077a7f554ef6b18f0\",\n",
    "        \"Books_5.json.gz\": \"ded924d1d1a22bae499f1a1c2b39397104304bfdb24232a2dd0aa50e89cd37bb\",\n",
    "    }\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        category: str = \"books\",\n",
    "        use_5_core: bool = True,\n",
    "        cache_dir: str = None,\n",
    "    ):\n",
    "        self.category = category\n",
    "        self.use_5_core = use_5_core\n",
    "        super().__init__(cache_dir=cache_dir)\n",
    "\n",
    "    def make_db(self) -> Database:\n",
    "        r\"\"\"Process the raw files into a database.\"\"\"\n",
    "\n",
    "        ### product table ###\n",
    "\n",
    "        url_key = self._category_to_url_key[self.category]\n",
    "        url = f\"{self.url_prefix}/metaFiles2/meta_{url_key}.json.gz\"\n",
    "        path = pooch.retrieve(\n",
    "            url,\n",
    "            known_hash=self.known_hashes.get(url.split(\"/\")[-1], None),\n",
    "            progressbar=True,\n",
    "            processor=pooch.Decompress(),\n",
    "        )\n",
    "        print(f\"reading product info from {path}...\")\n",
    "        tic = time.time()\n",
    "        ptable = pa.json.read_json(\n",
    "            path,\n",
    "            parse_options=pa.json.ParseOptions(\n",
    "                explicit_schema=pa.schema(\n",
    "                    [\n",
    "                        (\"asin\", pa.string()),\n",
    "                        (\"category\", pa.list_(pa.string())),\n",
    "                        (\"brand\", pa.string()),\n",
    "                        (\"title\", pa.string()),\n",
    "                        (\"description\", pa.list_(pa.string())),\n",
    "                        (\"price\", pa.string()),\n",
    "                    ]\n",
    "                ),\n",
    "                unexpected_field_behavior=\"ignore\",\n",
    "            ),\n",
    "        )\n",
    "        toc = time.time()\n",
    "        print(f\"done in {toc - tic:.2f} seconds.\")\n",
    "\n",
    "        print(\"converting to pandas dataframe...\")\n",
    "        tic = time.time()\n",
    "        pdf = ptable.to_pandas()\n",
    "        toc = time.time()\n",
    "        print(f\"done in {toc - tic:.2f} seconds.\")\n",
    "\n",
    "        print(\"processing product info...\")\n",
    "        tic = time.time()\n",
    "\n",
    "        # asin is not intuitive / recognizable\n",
    "        pdf.rename(columns={\"asin\": \"product_id\"}, inplace=True)\n",
    "\n",
    "        # somehow the raw data has duplicate product_id's\n",
    "        pdf.drop_duplicates(subset=[\"product_id\"], inplace=True)\n",
    "\n",
    "        # price is like \"$x,xxx.xx\", \"$xx.xx\", or \"$xx.xx - $xx.xx\", or garbage html\n",
    "        # if it's a range, we take the first value\n",
    "        pdf.loc[:, \"price\"] = pdf[\"price\"].apply(\n",
    "            lambda x: (\n",
    "                None\n",
    "                if x is None or x == \"\" or x[0] != \"$\"\n",
    "                else float(x.split(\" \")[0][1:].replace(\",\", \"\"))\n",
    "            )\n",
    "        )\n",
    "\n",
    "        # remove products with missing price\n",
    "        pdf = pdf.dropna(subset=[\"price\"])\n",
    "\n",
    "        pdf.loc[:, \"category\"] = pdf[\"category\"].apply(\n",
    "            lambda x: None if x is None or len(x) == 0 else x\n",
    "        )\n",
    "\n",
    "        # some rows are stored as ['cat1' 'cat2' 'cat3' ...]\n",
    "        # this function maps them to ['cat1', 'cat2', 'cat3', ...] (list of strings)\n",
    "        # since otherwise pytorch-frame breaks\n",
    "        def fix_column(value):\n",
    "            if isinstance(value, str):\n",
    "                return value  # Already a string\n",
    "            elif value is None:\n",
    "                return None\n",
    "            else:\n",
    "                return list(value)\n",
    "\n",
    "        pdf[\"category\"] = pdf[\"category\"].apply(fix_column)\n",
    "\n",
    "        # description is either [] or [\"some description\"]\n",
    "        pdf.loc[:, \"description\"] = pdf[\"description\"].apply(\n",
    "            lambda x: None if x is None or len(x) == 0 else x[0]\n",
    "        )\n",
    "\n",
    "        toc = time.time()\n",
    "        print(f\"done in {toc - tic:.2f} seconds.\")\n",
    "\n",
    "        ### review table ###\n",
    "\n",
    "        if self.use_5_core:\n",
    "            url = f\"{self.url_prefix}/categoryFilesSmall/{url_key}_5.json.gz\"\n",
    "        else:\n",
    "            url = f\"{self.url_prefix}/categoryFiles/{url_key}.json.gz\"\n",
    "        path = pooch.retrieve(\n",
    "            url,\n",
    "            known_hash=self.known_hashes.get(url.split(\"/\")[-1], None),\n",
    "            progressbar=True,\n",
    "            processor=pooch.Decompress(),\n",
    "        )\n",
    "        print(f\"reading review and customer info from {path}...\")\n",
    "        tic = time.time()\n",
    "        rtable = pa.json.read_json(\n",
    "            path,\n",
    "            parse_options=pa.json.ParseOptions(\n",
    "                explicit_schema=pa.schema(\n",
    "                    [\n",
    "                        (\"unixReviewTime\", pa.int32()),\n",
    "                        (\"reviewerID\", pa.string()),\n",
    "                        (\"reviewerName\", pa.string()),\n",
    "                        (\"asin\", pa.string()),\n",
    "                        (\"overall\", pa.float32()),\n",
    "                        (\"verified\", pa.bool_()),\n",
    "                        (\"reviewText\", pa.string()),\n",
    "                        (\"summary\", pa.string()),\n",
    "                    ]\n",
    "                ),\n",
    "                unexpected_field_behavior=\"ignore\",\n",
    "            ),\n",
    "        )\n",
    "        toc = time.time()\n",
    "        print(f\"done in {toc - tic:.2f} seconds.\")\n",
    "\n",
    "        print(\"converting to pandas dataframe...\")\n",
    "        tic = time.time()\n",
    "        rdf = rtable.to_pandas()\n",
    "        toc = time.time()\n",
    "        print(f\"done in {toc - tic:.2f} seconds.\")\n",
    "\n",
    "        print(\"processing review and customer info...\")\n",
    "        tic = time.time()\n",
    "\n",
    "        rdf.rename(\n",
    "            columns={\n",
    "                \"unixReviewTime\": \"review_time\",\n",
    "                \"reviewerID\": \"customer_id\",\n",
    "                \"reviewerName\": \"customer_name\",\n",
    "                \"asin\": \"product_id\",\n",
    "                \"overall\": \"rating\",\n",
    "                \"reviewText\": \"review_text\",\n",
    "            },\n",
    "            inplace=True,\n",
    "        )\n",
    "\n",
    "        rdf.loc[:, \"review_time\"] = pd.to_datetime(rdf[\"review_time\"], unit=\"s\")\n",
    "\n",
    "        toc = time.time()\n",
    "        print(f\"done in {toc - tic:.2f} seconds.\")\n",
    "\n",
    "        print(\"keeping only products common to product and review tables...\")\n",
    "        tic = time.time()\n",
    "        plist = list(set(pdf[\"product_id\"]) & set(rdf[\"product_id\"]))\n",
    "        pdf.query(\"product_id == @plist\", inplace=True)\n",
    "        rdf.query(\"product_id == @plist\", inplace=True)\n",
    "        toc = time.time()\n",
    "        print(f\"done in {toc - tic:.2f} seconds.\")\n",
    "\n",
    "        print(\"extracting customer table...\")\n",
    "        tic = time.time()\n",
    "        cdf = (\n",
    "            rdf[[\"customer_id\", \"customer_name\"]]\n",
    "            .drop_duplicates(subset=[\"customer_id\"])\n",
    "            .copy()\n",
    "        )\n",
    "        rdf.drop(columns=[\"customer_name\"], inplace=True)\n",
    "        toc = time.time()\n",
    "        print(f\"done in {toc - tic:.2f} seconds.\")\n",
    "\n",
    "        db = Database(\n",
    "            table_dict={\n",
    "                \"product\": Table(\n",
    "                    df=pdf,\n",
    "                    fkey_col_to_pkey_table={},\n",
    "                    pkey_col=\"product_id\",\n",
    "                    time_col=None,\n",
    "                ),\n",
    "                \"customer\": Table(\n",
    "                    df=cdf,\n",
    "                    fkey_col_to_pkey_table={},\n",
    "                    pkey_col=\"customer_id\",\n",
    "                    time_col=None,\n",
    "                ),\n",
    "                \"review\": Table(\n",
    "                    df=rdf,\n",
    "                    fkey_col_to_pkey_table={\n",
    "                        \"customer_id\": \"customer\",\n",
    "                        \"product_id\": \"product\",\n",
    "                    },\n",
    "                    pkey_col=None,\n",
    "                    time_col=\"review_time\",\n",
    "                ),\n",
    "            }\n",
    "        )\n",
    "\n",
    "        db = db.from_(pd.Timestamp(\"2008-01-01\"))\n",
    "\n",
    "        return db"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "631fa5e1-ec21-4e23-b9ca-7e455459f49d",
   "metadata": {},
   "source": [
    "Now to make the fashion dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c315c96a-5450-47d8-b974-c1f61510366b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:31.624395Z",
     "iopub.status.busy": "2024-07-26T00:25:31.624017Z",
     "iopub.status.idle": "2024-07-26T00:25:31.654247Z",
     "shell.execute_reply": "2024-07-26T00:25:31.653223Z",
     "shell.execute_reply.started": "2024-07-26T00:25:31.624360Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AmazonDataset()"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "amazon_fashion_dataset = AmazonDataset(category=\"fashion\", use_5_core=True)\n",
    "amazon_fashion_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "bc68cd58-d4fa-4433-a95b-6b912167b761",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:31.797241Z",
     "iopub.status.busy": "2024-07-26T00:25:31.796793Z",
     "iopub.status.idle": "2024-07-26T00:25:32.211913Z",
     "shell.execute_reply": "2024-07-26T00:25:32.210878Z",
     "shell.execute_reply.started": "2024-07-26T00:25:31.797192Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Making Database object from scratch...\n",
      "(You can also use `get_dataset(..., download=True)` for datasets prepared by the RelBench team.)\n",
      "reading product info from /lfs/ampere4/0/ranjanr/.cache/pooch/b70e8d295f37e2465ea17803b6d1e11d-meta_AMAZON_FASHION.json.gz.decomp...\n",
      "done in 0.07 seconds.\n",
      "converting to pandas dataframe...\n",
      "done in 0.16 seconds.\n",
      "processing product info...\n",
      "done in 0.11 seconds.\n",
      "reading review and customer info from /lfs/ampere4/0/ranjanr/.cache/pooch/26323778935ec86761e2c260cea27601-AMAZON_FASHION_5.json.gz.decomp...\n",
      "done in 0.01 seconds.\n",
      "converting to pandas dataframe...\n",
      "done in 0.00 seconds.\n",
      "processing review and customer info...\n",
      "done in 0.00 seconds.\n",
      "keeping only products common to product and review tables...\n",
      "done in 0.01 seconds.\n",
      "extracting customer table...\n",
      "done in 0.00 seconds.\n",
      "Done in 0.38 seconds.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Database()"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "amazon_fashion_db = amazon_fashion_dataset.get_db()\n",
    "amazon_fashion_db"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fd697b8-46e1-4d69-a70a-1687700d63c1",
   "metadata": {},
   "source": [
    "To register this dataset, pass in the `args` and `kwargs` too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "92e0a135-3fb2-4ff4-8e9a-9964e8f037e5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:32.213317Z",
     "iopub.status.busy": "2024-07-26T00:25:32.213145Z",
     "iopub.status.idle": "2024-07-26T00:25:32.231430Z",
     "shell.execute_reply": "2024-07-26T00:25:32.230569Z",
     "shell.execute_reply.started": "2024-07-26T00:25:32.213302Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['rel-amazon',\n",
       " 'rel-avito',\n",
       " 'rel-event',\n",
       " 'rel-f1',\n",
       " 'rel-hm',\n",
       " 'rel-stack',\n",
       " 'rel-trial',\n",
       " 'rel-custom_f1',\n",
       " 'rel-amazon_fashion']"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "register_dataset(\n",
    "    \"rel-amazon_fashion\", AmazonDataset, category=\"fashion\", use_5_core=True\n",
    ")\n",
    "get_dataset_names()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9dda56b3-43d6-43fb-b28a-cfc172db26eb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:25:32.375618Z",
     "iopub.status.busy": "2024-07-26T00:25:32.375348Z",
     "iopub.status.idle": "2024-07-26T00:25:32.392944Z",
     "shell.execute_reply": "2024-07-26T00:25:32.392057Z",
     "shell.execute_reply.started": "2024-07-26T00:25:32.375595Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AmazonDataset()"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "registered_amazon_fashion_dataset = get_dataset(\"rel-amazon_fashion\")\n",
    "registered_amazon_fashion_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e734a81e-fad5-470a-83a6-ee12132d53e3",
   "metadata": {},
   "source": [
    "Note that the registry does not persist beyond the running Python process. This means that to run the baseline scripts at `examples/` in the RelBench repo, you will first have to modify the script to register your own dataset before `get_dataset` is called in the script."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "244716a8-ef85-4531-95a1-e45aea78cc60",
   "metadata": {},
   "source": [
    "## Next steps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10a22967-f376-4157-be22-d4443c04d108",
   "metadata": {},
   "source": [
    "To use the custom dataset within the RelBench framework, you would want to define custom tasks on it. See this tutorial for how to do that: [custom_task.ipynb](custom_task.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a2cf7b6-1eeb-4b38-a03c-01b4d159d3b7",
   "metadata": {},
   "source": [
    "Please also consider sharing your dataset with the community by getting it added to the RelBench dataset repository. Check out our [CONTRIBUTING.md](https://github.com/snap-stanford/relbench/blob/main/CONTRIBUTING.md) for how to do this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45ccdb7c-bc92-47ae-a0df-90ce8ef1fdfd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}