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

Petastorm enables direct training of deep learning models from Apache Parquet datasets using TensorFlow, PyTorch, PySpark, or pure Python, with support for selective column access, shuffling, filtering, and distributed parallelism.

Yes, if you are already using Parquet for data storage and training with TensorFlow, PyTorch, or PySpark. The low install friction and permissive license make it a straightforward addition. However, note the aging maintenance status (224 days since last release)—verify compatibility with your specific framework versions before committing to production use. No known security vulnerabilities.

Install

petastorm on PyPI

pip install petastorm

uv add petastorm

poetry add petastorm

Installing petastorm

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.

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)

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

Verify before relying

• Whether the aging maintenance status (224 days since last release) affects compatibility with recent TensorFlow or PyTorch versions.
• Performance characteristics and scalability limits for very large distributed datasets.
• Current state of optional dependencies (tf, tf_gpu, torch, opencv) and their version compatibility.

Package facts

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

Tags

More Artificial Intelligence packages