Load Unity Catalog tables into PyTorch with UCDataset

Important

This feature is in Beta.

UCDataset is a PyTorch IterableDataset that streams data from Unity Catalog tables into a training loop on AI Runtime. You describe what to read as a DataFrame: choose a table, then select columns and filter rows. The dataset reads the matching rows directly from the table's data files in cloud storage, so you don't need a Spark session or a SQL warehouse. It applies the table's deletion vectors as it reads, splits the matching rows evenly across GPUs and data loader workers, and yields one dictionary per batch. Numeric and boolean columns become tensors. Other columns become Python lists.

Choose a dataset based on where your training data lives:

Training data Dataset
Rows in a Unity Catalog table, such as feature and label columns in a Delta Lake table UCDataset (Beta)
Files in a Unity Catalog volume, such as images, audio, or documents UCVolumeDataset

The two datasets work side by side. For example, a training job can read labels from a table with UCDataset and the matching images from a volume with UCVolumeDataset.

For other data, use these options instead:

Requirements

  • An AI Runtime workload. UCDataset is available wherever AI Runtime is. See AI Runtime availability by Geo.
  • Linux on x86_64 with Python 3.10 or above. The package is built for x86_64 only. AI Runtime GPU compute meets this requirement.
  • A Delta Lake table registered in Unity Catalog.
  • The USE CATALOG privilege on the table's catalog, the USE SCHEMA privilege on its schema, and the SELECT privilege on the table.
  • To train on more than one GPU, a launcher that sets the RANK and WORLD_SIZE environment variables for every process, such as torchrun. If neither variable is set, such as in an AI Runtime notebook, UCDataset reads the table as a single shard. If only one of them is set, UCDataset raises an error.

Important

When you submit a workload with databricks air run, start the training script with torchrun, even on a single GPU. AI Runtime sets WORLD_SIZE for every workload but doesn't set RANK, and UCDataset raises an error when only one of the two is set. torchrun sets both for each process. See Run a multi-node job on AI Runtime.

Load a table

  1. Install the databricks-uc-dataset package. For a workload that you submit with databricks air run, add the package to environment.dependencies. See Run a multi-node job on AI Runtime. In a notebook, run the following:

    %pip install databricks-uc-dataset==0.0.1b1

  2. Select the table, columns, and rows to read. table() takes a three-level table name, and the optional version_as_of argument reads a specific table version:

    from databricks.uc_dataset.beta.dataframe import table, col

df = (
    table("main.ml.clicks")
    .select("features", "label")
    .filter(col("label").isNotNull() & col("features").isNotNull())
)

  3. Create the dataset and wrap it in a PyTorch DataLoader. The dataset already returns batches, so set batch_size=None on the loader. If you train with PyTorch DistributedDataParallel (DDP), call torch.distributed.init_process_group() before you create the dataset. Rank 0 then plans the read one time and shares the plan with the other ranks:

    import torch
from torch.utils.data import DataLoader
from databricks.uc_dataset.beta import UCDataset

dataset = UCDataset(df=df, batch_size=256)
loader = DataLoader(dataset, batch_size=None, num_workers=4)

for epoch in range(num_epochs):
    for batch in loader:
        # Array columns arrive as Python lists, and a NULL array arrives as None.
        # Convert the lists before you train.
        features = torch.as_tensor(batch["features"], dtype=torch.float32)
        label = batch["label"]
        train_step(features, label)

  4. To train on more than one GPU, save the code as train.py and launch it with torchrun. In a notebook, skip this step. For example, to train on 8 GPUs on one node, run the following:

    torchrun --nproc_per_node=8 train.py

Each rank and worker reads its own shard of the matching rows. Each batch is a dictionary that maps each selected column name to a tensor or a list.

Run a multi-node job on AI Runtime

To train on more than one node, describe the workload in workload.yaml and submit it with databricks air run:

  • Set environment.version to databricks_ai_v6, which includes PyTorch, and add databricks-uc-dataset==0.0.1b1 to environment.dependencies.

  • Start the training script with torchrun. AI Runtime runs the command on every node and sets NUM_NODES, NODE_RANK, LOCAL_WORLD_SIZE, MASTER_ADDR, and MASTER_PORT. torchrun uses them to set RANK for each process:

    command: |
  torchrun --nnodes=$NUM_NODES --node_rank=$NODE_RANK --nproc_per_node=$LOCAL_WORLD_SIZE \
    --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT "$CODE_SOURCE_PATH/train.py"

For a complete workload with train.py, workload.yaml, and a sample table, see Click-through-rate model on a Unity Catalog table in the AI Runtime Cookbook. To train on more GPUs, change num_accelerators in workload.yaml. train.py doesn't change.

Filter rows

UCDataset applies filters while it reads the files, so rows that don't match don't reach your training loop. Build a filter with col() and the following expressions:

Expression Example
Comparison col("amount") > 50. Also ==, !=, <, <=, and >=.
Null check col("label").isNull(), col("label").isNotNull()
Null-safe equality col("region").eqNullSafe("us")
Boolean logic (col("a") > 1) & (col("b") == "x"). Also \| (or) and ~ (not).

Filters don't support SQL strings, isin or IN, BETWEEN, string matching such as like or contains, arithmetic on columns, cast, or user-defined functions (UDFs). Pass each condition as a Column expression, and combine conditions with &, |, and ~. A comparison takes a column and a literal value, in either order. You can't compare two columns. To match a list of values, combine == conditions with |. Very long filters can fail. See Limitations.

The following example reads purchase events from a fixed table version. The version must exist:

df = (
    table("main.ml.events", version_as_of=120)
    .filter((col("event_type") == "purchase") & (col("amount") > 50))
    .select("user_id", "amount", "label")
)

Output format

Each batch is a Python dictionary that maps each selected column name to its values:

Column type Batch value
Integer, with no nulls in the batch 1-D tensor with the column's integer type, such as int32
Integer, with nulls in the batch (except 64-bit) 1-D float64 tensor, with nulls as NaN
Floating point 1-D tensor, with nulls as NaN
Boolean, with no nulls Boolean tensor
Binary List of bytes objects
Other types, such as strings, arrays, structs, and timestamps Python list

Because an integer column's type depends on whether the batch contains a null, the column can change type from one batch to the next. If your model needs a fixed integer type, such as for embedding lookups or class labels, filter out the nulls or replace them in a transform function.

If a column can't be converted without losing data, UCDataset raises an error. This happens for a 64-bit integer column that contains nulls and for a boolean column that contains nulls. The error appears when a batch with a null is read, so it can happen partway through an epoch. To avoid it, filter out the nulls with isNotNull(), or replace them in a transform function.

To change the batch layout or handle nulls, pass a transform function. UCDataset calls it on each batch as a PyArrow RecordBatch, before any conversion to tensors, and yields whatever the function returns, so convert every column that you use. The following example returns a tuple of tensors:

import torch

def to_features_and_label(batch):
    features = torch.as_tensor(batch.column("features").to_pylist(), dtype=torch.float32)
    label = torch.from_numpy(batch.column("label").to_numpy(zero_copy_only=False))
    return features, label

dataset = UCDataset(df=df, transform=to_features_and_label)

The following example replaces nulls in a 64-bit integer column with 0, so the column stays int64 in every batch:

import pyarrow.compute as pc
import torch

def fill_nulls(batch):
    user_ids = pc.fill_null(batch.column("user_id"), 0)
    return {"user_id": torch.tensor(user_ids.to_numpy())}

dataset = UCDataset(df=df, transform=fill_nulls)

Distributed training and shuffling

UCDataset splits the rows that match your filter across every rank and every data loader worker. You don't need a distributed sampler. By default, shard sizes differ by at most one row.

Equal shard sizes don't guarantee equal batch counts. One rank can get one more batch than another, for example when some shards hold an exact multiple of batch_size rows and others hold one row more, or when a filter keeps fewer rows than there are shards. With DDP, the rank with the extra batch then waits on a gradient all-reduce that the other ranks never join, and the job fails. To let ranks finish an epoch at different times, wrap each epoch in the PyTorch Join context manager, where model is the DDP-wrapped model:

from torch.distributed.algorithms.join import Join

for epoch in range(num_epochs):
    with Join([model]):
        for batch in loader:
            train_step(batch)

dataset.num_rows() returns the number of rows that match the filter, across all ranks and workers. len(dataset) and len(loader) raise a TypeError. If a library needs the number of steps in an epoch, such as a learning rate scheduler, estimate it from num_rows(), the world size, and batch_size. Each rank can read one more batch than the estimate, and each additional data loader worker can add another partial batch.

UCDataset reads the table's files in a fixed, shuffled order. Every epoch uses the same order, and so does a new dataset on an unchanged table. You can't set a seed, and the files aren't reshuffled between epochs.

Authentication

UCDataset reads the table as the identity that runs the workload, and Unity Catalog permissions on the table apply.

Unity Catalog doesn't issue storage credentials for tables with row filters or column masks, so UCDataset can't read them. Creating the dataset fails with an error that includes the message from Unity Catalog.

Limitations

  • UCDataset reads only Delta Lake tables in Unity Catalog. Views, materialized views, and streaming tables aren't supported.
  • Tables with row filters or column masks can't be read, because Unity Catalog doesn't issue storage credentials for them.
  • Shallow clones can't be read, because their data files are stored with the source table. To pin the rows for a training run, read the source table with version_as_of instead. A clone starts its own version history, so it doesn't have the source table's earlier versions.
  • Tables with a VARIANT column can't be read, even if you don't select that column.
  • version_as_of takes a table version number. Timestamps aren't supported: version_as_of="2026-01-01" raises a versionAsOf must be a non-negative integer error.
  • Select at least one column that isn't a partition column. Selecting only partition columns raises a Delta reads require at least one non-partition column error.
  • A filter with a long chain of conditions, such as many == comparisons combined with |, can fail with a failed to decode Spark Connect plan error. If it does, use fewer conditions.
  • Filters support only the expressions in Filter rows.
  • Array columns, such as embeddings, are returned as Python lists, which is much slower than numeric columns.
  • UCDataset doesn't checkpoint its read position. If training restarts, the dataset starts again from the beginning of the table, and you can't resume in the middle of an epoch.

To report an issue or give feedback during the Beta, contact your Azure Databricks account team.

Additional resources

  • Last updated on