petastorm

What it is and what it does

Petastorm is a data access library that bridges Apache Parquet storage and popular Python machine learning frameworks. It was developed at Uber ATG to enable efficient, distributed training of deep learning models directly from Parquet datasets without requiring intermediate format conversions. The library handles the schema mapping between Parquet and framework-specific types (TensorFlow, PyTorch, PySpark), and provides a unified reader interface that supports selective column access, row filtering, shuffling, and partitioning for multi-GPU training.

You use Petastorm in two phases: first, generate a Parquet dataset using PySpark with the Unischema API to define field types, shapes, and compression codecs; then, read from that dataset using the Reader class or framework-specific adapters (tf_tensors for TensorFlow, DataLoader for PyTorch). The library handles parallelism internally via threads, processes, or single-threaded modes, and supports local caching to reduce repeated I/O.

Use it for

• Train TensorFlow or PyTorch models on large Parquet datasets stored on HDFS or local filesystems without loading entire datasets into memory.
• Generate Parquet datasets from raw data using PySpark, then iterate over them in Python ML training loops with automatic schema validation.
• Distribute training across multiple GPUs by partitioning Petastorm datasets and reading different partitions on different workers.
• Apply row-level filtering and selective column reads to reduce I/O when training on subsets of large datasets.
• Compress image and array data using standard codecs (JPEG, PNG) or custom codecs within a single Parquet-backed dataset.

Worth the install?

AI-flagged interpretation of the facts on this page. Verify before relying on it.

Install

petastorm on PyPI

Before you install

Low friction install with a pure-wheel distribution and 13 runtime dependencies already packaged. Maintenance is aging—last release was 224 days ago—but the repository remains active and unarchived with 1890 stars.

Requires a Parquet dataset already created in Petastorm format; generating one requires PySpark and the Unischema API shown in the documentation.

License in practice

Apache License 2.0 is permissive, allowing commercial and private use with minimal restrictions; you must retain license notices in distributions.

Quickstart

pip install petastorm

from petastorm import make_reader

with make_reader('file:///path/to/dataset') as reader:
    for row in reader:
        print(row)

Package facts

Evidence: petastorm-0.13.1-py2.py3-none-any.whl

Tags

More Artificial Intelligence packages

See also litdata · datasets · webdataset · spark-sklearn · mosaicml-streaming · torch · tensorflow-datasets · pyspark-huggingface · raydp · keras-nightly