Skip to main content

How to do fast random-access reads for model training directly from object storage

Summary

  • Match the read path to the data: tabular features in Delta tables read through Spark, files such as images in Unity Catalog volumes read by path, and large deep learning datasets that need shuffled sample access as Mosaic Data Shard (MDS) files read with Mosaic Streaming.
  • Mosaic Streaming is an open source data loading library, pre-installed in Databricks Runtime 15.2 ML and higher. Its StreamingDataset is a version of PyTorch's IterableDataset with elastically deterministic shuffling, which enables fast mid-epoch resumption.
  • Unity Catalog volumes govern non-tabular files in cloud object storage and expose them at paths like /Volumes////, so training code can use local file APIs as well as Spark.
  • For tabular training data, per-file data skipping statistics let filtered reads skip irrelevant files, and liquid clustering organizes data by clustering keys. Databricks recommends liquid clustering for all new tables; it can't be combined with partitioning or ZORDER.

How to do fast random-access reads for model training directly from object storage

Keep the data in object storage, govern it with Unity Catalog, and match the read path to the shape of the data. For deep learning that samples records in shuffled order, convert the data once to Mosaic Data Shard (MDS) format and train from those shards with Mosaic Streaming's StreamingDataset, which shuffles deterministically and can resume mid-epoch. For files such as images or audio, store them in a Unity Catalog volume and read them by path. For tabular features, read Delta tables through Spark, where data skipping statistics let filtered reads skip irrelevant files and the disk cache serves repeated reads from local storage.

Why do training reads need their own approach?

Two documented constraints shape the choice. Training loops shuffle samples, and the Mosaic Streaming docs note that unshuffled data can lead to divergence, loss spikes or suboptimal training; they add that setting shuffle=True is performant in nearly all cases. And for datasets that don't fit in memory, the Databricks docs recommend streaming approaches. Delta tables store data in Parquet, which the Apache Parquet project describes as a column-oriented format designed for efficient data storage and retrieval.

Which read path fits your training data?

Data shape Where it lives How to read it for training
Tabular features Delta table in Unity Catalog Spark DataFrames, or pandas through toPandas()
Images, audio, documents and other files Unity Catalog volume Local file APIs on /Volumes/<catalog>/<schema>/<volume>/<path>
Large deep learning datasets read in shuffled order MDS shards in a Unity Catalog volume StreamingDataset with StreamingDataLoader, single-node or distributed with TorchDistributor
Datasets too large for memory Unity Catalog volumes, among other sources PyTorch IterableDataset for custom streaming logic, Hugging Face datasets with streaming, or Ray Data for distributed batch processing

How Databricks approaches training reads from object storage

Unity Catalog governs both kinds of training data. Tables govern tabular data, and volumes govern non-tabular files of any format; Databricks recommends volumes for all non-tabular data. On a cluster, you can reach that data with both Spark and local file APIs. Volumes need Databricks Runtime 13.3 LTS or above.
For deep learning, the documented workflow uses Mosaic Streaming. Load and preprocess the data with Apache Spark. Write it to MDS with dataframe_to_mds, using a Unity Catalog volume for persistent storage. Then read it with StreamingDataset, or iterate with StreamingDataLoader, which adds a checkpoint and resumption interface that tracks the samples each rank has seen.

from streaming.base.converters import dataframe_to_mds
from streaming import StreamingDataset
from torch.utils.data import DataLoader

# df is a Spark DataFrame; MDSWriter needs `out` and `columns`
out_path = "/Volumes/<catalog>/<schema>/<volume>/mds"
mds_kwargs = {'out': out_path, 'columns': {'number': 'int64', 'words': 'str'}}
dataframe_to_mds(df.repartition(4), merge_index=True, mds_kwargs=mds_kwargs)

dataset = StreamingDataset(local=out_path, remote=None, batch_size=2, shuffle=True)
dataloader = DataLoader(dataset, batch_size=2, num_workers=1)

Mosaic Streaming works with native PyTorch, PyTorch Lightning and TorchDistributor, handles images, text, video and multimodal data, and reads from Unity Catalog volumes and major cloud object stores.

Which table settings speed up Spark reads of training data?

Setting What it does Notes from the docs
Liquid clustering Organizes data by clustering keys and replaces partitioning and ZORDER Recommended for all new tables; generally available for Delta tables on Databricks Runtime 15.4 LTS and above; keys can change without rewriting data
Data skipping statistics Per-file minimum and maximum values, null counts and record counts let queries skip irrelevant files Collected automatically when you write
Disk cache Copies remote Parquet files into local storage on worker nodes Triggered on first read; choose worker types with SSD volumes
Photon Native vectorized engine for SQL and DataFrame work, including feature engineering Uses filter pushdown, dictionary pruning and row-group skipping

To cluster a new table, add CLUSTER BY at creation: CREATE TABLE table1 (col0 INT, col1 STRING) CLUSTER BY (col0);.

FAQs

Should I partition training tables by date?

For new tables, Databricks recommends liquid clustering, which replaces table partitioning and ZORDER and isn't compatible with either. On Databricks Runtime 18.1 and above, REPLACE PARTITIONED BY WITH CLUSTER BY in an ALTER TABLE statement converts an existing partitioned Delta table.

Do I need to install Mosaic Streaming?

Not on Databricks Runtime 15.2 ML and higher, where it's pre-installed. If StreamingDataset raises a default-credentials error when reading from a volume, set DATABRICKS_HOST and DATABRICKS_TOKEN in the environment, and on worker nodes too when you use TorchDistributor.

Can PyTorch read files in a Unity Catalog volume directly?

Yes. Volumes expose files at /Volumes/<catalog>/<schema>/<volume>/<path>, and Unity Catalog supports local file APIs on a cluster as well as Spark. Some limits depend on the runtime version and access mode, so check the volumes limitations for your compute.

Does Photon speed up the training loop itself?

Photon works at the Spark execution layer: for supported operations, it replaces the JVM-based Spark SQL engine with a native C++ runtime. The docs list feature engineering among the machine learning operations it improves, so it helps the data preparation that feeds training.

The information provided herein is for general informational purposes only and may not reflect the most current product capabilities or configurations.