Datasets

The Chalk Dataset class governs metadata related to offline queries, supports revisions to queries over time, and enables the easy retrieval of data from the cloud.

​

Datasets from offline query

Dataset instances are obtained by calling ChalkClient.offline_query() which computes feature values from the offline store. If inputs are given, the method returns the values corresponding to those inputs. Otherwise, the method returns a random sample according to the parameter max_samples, or features from within timebounds specified by lower_bound and upper_bound.

from chalk.client import ChalkClient, Dataset
uids = [1, 2, 3, 4]
at = datetime.now()
dataset: Dataset = ChalkClient().offline_query(
     input={
         User.id: uids,
     },
     input_times=[at] * len(uids),
     output=[
         User.id,
         User.fullname,
         User.email,
         User.name_email_match_score,
     ],
     dataset_name='my_dataset'
)

sample_dataset: Dataset = ChalkClient().offline_query(
     output=[
         User.id,
         User.fullname,
         User.email,
         User.name_email_match_score,
     ],
     max_samples=10,
     lower_bound=datetime.now() - timedelta(days=7),
     upper_bound=datetime.now(),
     dataset_name='my_sample'
)

Here, we attach a unique name to the Dataset. Whenever we send additional queries with the same name, a new DatasetRevision instance will be created and attached to the existing dataset. If a dataset_name is not given, the output data won’t be retrievable beyond the current session.

A Dataset’s revisions can be inspected in Dataset.revisions: they hold useful metadata relating to the offline query job and the data itself. Be sure to check out Dataset.errors for any errors upon submitting the query.

​

Retrieving output data

Since offline queries are not realtime, the Dataset instance returned is not guaranteed to have the outputs of the query instantaneously. Thus, loading the data may take some time.

The data can be accessed programmatically by calling Dataset.get_data_as_pandas(), Dataset.get_data_as_polars(), or Dataset.get_data_as_dataframe(). If the offline query job is still running, the Dataset will poll the engine until the results are completed.

from chalk.client import ChalkClient, Dataset
uids = [1, 2, 3, 4]
at = datetime.now()
dataset: Dataset = ChalkClient().offline_query(
     input={
         User.id: uids,
     },
     input_times=[at] * len(uids),
     output=[
         User.id,
         User.fullname,
         User.email,
         User.name_email_match_score,
     ],
     dataset_name='my_dataset'
)

pandas_df: pd.DataFrame = dataset.get_data_as_pandas()
polars_df: pl.LazyFrame = dataset.get_data_as_polars()
chalk_df: chalk.features.DataFrame = dataset.get_data_as_dataframe()

The file outputs of the query themselves can also be downloaded to a specified directory.

from chalk.client import ChalkClient, Dataset
uids = [1, 2, 3, 4]
at = datetime.now()
dataset: Dataset = ChalkClient().offline_query(
     input={
         User.id: uids,
     },
     input_times=[at] * len(uids),
     output=[
         User.id,
         User.fullname,
         User.email,
         User.name_email_match_score,
     ],
     dataset_name='my_dataset'
)
dataset.download_data('my_directory')

By default, Dataset instances fetch the output data from their most recent revision. A specific DatasetRevision’s output data can be fetched using the same methods.

from chalk.client import ChalkClient, Dataset
uids = [1, 2, 3, 4]
at = datetime.now()
dataset: Dataset = ChalkClient().offline_query(
     input={
         User.id: uids,
     },
     input_times=[at] * len(uids),
     output=[
         User.id,
         User.fullname,
         User.email,
         User.name_email_match_score,
     ],
     dataset_name='my_dataset'
)
for revision in dataset.revisions:
    print(revision.get_data_as_pandas())

​

Dataset Inputs

Dataset objects also store the inputs for each revision.

from chalk.client import ChalkClient, Dataset
uids = [1, 2, 3, 4]
at = datetime.now()
dataset: Dataset = ChalkClient().offline_query(
     input={
         User.id: uids,
     },
     input_times=[at] * len(uids),
     output=[
         User.id,
         User.fullname,
         User.email,
         User.name_email_match_score,
     ],
     dataset_name='my_dataset'
)
df = dataset.get_input_dataframe()

​

Managing Datasets

​

Renaming Datasets

You can rename a dataset from the Chalk dashboard to better organize your datasets or reflect changes in your workflow.

To rename a dataset:

1. From the Datasets page under the Offline section in the sidebar, click on a dataset to open its detail page
2. Click the Edit button in the page header next to the dataset name
3. Enter the new name and click Save, or press Enter

​

Archiving Dataset Revisions

As datasets accumulate revisions over time, you can archive older or buggy revisions to keep your dataset history organized. Some metadata about your revision will still be accessible through the dashboard but revisions will be permanently inaccessible after archiving.

To archive a revision:

1. From the Datasets page under the Offline section in the sidebar, click on a dataset to open its detail page
2. Select the Revisions tab
3. Select the checkboxes next to all revisions you want to archive
4. An Archive button will appear in the actions bar at the top of the page. Click this button and confirm the action in the dialog.

To view archived revisions, check the Include archived revisions checkbox above the revisions table.

Datasets are stored as parquet files in the storage buckets in your cloud. Archiving a revision does not delete the data there—to conserve storage, you can set an expiration policy directly on those buckets.

​

Recompute a Dataset

Datasets expose a recompute method that enables users to see the results of updates to resolvers/features in the context of this dataset. recompute takes a list of features as an argument to be recomputed, and any other required input features are sampled from the offline store.

​

Promote a Dataset to Iceberg

A dataset revision’s output is stored as Parquet. promote_to_iceberg generates standard Apache Iceberg metadata that references those Parquet files in place — no data is rewritten. The Iceberg metadata.json and manifests are written to a sibling iceberg/ prefix next to the revision’s outputs and reference the existing Parquet files.

from chalk.client import ChalkClient

dataset = ChalkClient().get_dataset(dataset_name="my_dataset")
result = dataset.promote_to_iceberg()

# gcs://<your-dataset-bucket>/job_<rev>/query_<rev>/iceberg/metadata/00001-<uuid>.metadata.json
print(result.metadata_location)
print(result.num_rows, result.num_data_files, result.snapshot_id)

The returned metadata_location points at the table’s metadata.json. The table is not registered in any Iceberg catalog — you register metadata_location in your own catalog (Snowflake, Glue, Polaris, BigLake, …) for name-based discovery. Because it is a standard Iceberg table, any Iceberg-aware client (Snowflake, Spark, DuckDB, PyIceberg, Trino) can read it.

​

Example: reading a promoted dataset from Snowflake

Snowflake reads the promoted table through an object-store catalog integration plus an external volume pointing at the dataset storage bucket.

One-time setup (Snowflake, ACCOUNTADMIN):

CREATE CATALOG INTEGRATION IF NOT EXISTS chalk_iceberg
  CATALOG_SOURCE = OBJECT_STORE TABLE_FORMAT = ICEBERG ENABLED = TRUE;

CREATE EXTERNAL VOLUME IF NOT EXISTS chalk_datasets
  ALLOW_WRITES = FALSE
  STORAGE_LOCATIONS = (
    (NAME = 'chalk' STORAGE_PROVIDER = 'GCS'
     STORAGE_BASE_URL = 'gcs://<your-dataset-bucket>/')
  );

-- Copy the STORAGE_GCP_SERVICE_ACCOUNT value from the output:
DESC EXTERNAL VOLUME chalk_datasets;

Grant that service account read access to the bucket (GCP side):

SA="<STORAGE_GCP_SERVICE_ACCOUNT from DESC>"
gcloud storage buckets add-iam-policy-binding gs://<your-dataset-bucket> \
  --member="serviceAccount:$SA" --role="roles/storage.objectViewer"
gcloud storage buckets add-iam-policy-binding gs://<your-dataset-bucket> \
  --member="serviceAccount:$SA" --role="roles/storage.legacyBucketReader"

Per promoted dataset, create an Iceberg table from the metadata file. METADATA_FILE_PATH is the metadata_location returned by promote_to_iceberg() with the gcs://<your-dataset-bucket>/ prefix stripped:

CREATE ICEBERG TABLE my_dataset
  CATALOG = 'chalk_iceberg' EXTERNAL_VOLUME = 'chalk_datasets'
  METADATA_FILE_PATH = 'job_<rev>/query_<rev>/iceberg/metadata/00001-<uuid>.metadata.json';

SELECT "USER.ID", "USER.SCORE" FROM my_dataset LIMIT 10;

Output columns are the dataset’s feature fully-qualified names, uppercased and quoted (e.g. "USER.ID"). For an S3-backed dataset bucket, use STORAGE_PROVIDER = 'S3' in the external volume and grant the vended IAM principal read access instead of a GCP service account.

​

Integrate Datases into PyTorch

Datasets expose methods create_torch_map_dataset and create_torch_iter_dataset, which create PyTorch datasets from the results of the Chalk dataset.

To learn more, see Chalk’s PyTorch Integration with Datasets.