Petastorm

0.13.1 · active · verified Thu Apr 16

Petastorm is a Python library that enables single-node or distributed training of machine learning models directly from datasets stored in Parquet format. It provides data access for popular frameworks like TensorFlow, PyTorch, and Apache Spark. The current stable version is 0.13.1, with releases typically following a feature-driven cadence, often including release candidates before stable versions.

Common errors

```
ModuleNotFoundError: No module named 'petastorm.spark'
```
cause The `petastorm.spark` module is only available if Petastorm was installed with the `spark` extra and `pyspark` is installed.

fix Install Petastorm with the Spark extra: `pip install petastorm[spark]` and ensure `pyspark` is also installed in your environment.
```
FileNotFoundError: [Errno 2] No such file or directory: 'file:///path/to/my_dataset'
```
cause The specified dataset URL or path for `make_reader` or `make_writer` does not exist or is inaccessible. This often happens with incorrect paths, network drive issues, or missing data.

fix Double-check the `dataset_url` to ensure it's a valid path to an existing dataset directory (or a directory where you intend to write data). For HDFS/S3, ensure proper authentication and client setup.
```
Pyarrow.lib.ArrowInvalid: Could not convert ...
```
cause This error typically indicates a data type or schema mismatch when writing or reading data. The data being processed doesn't conform to the `Unischema` or expected Parquet types.

fix Verify that the `Unischema` definition correctly matches the actual data types and shapes you are writing. For reading, ensure the schema used by Petastorm aligns with the schema of the Parquet files.

Warnings

breaking The default `reader_pool_type` for `make_reader` changed from 'thread' to 'process' in Petastorm v0.13.0. This can cause issues if your data contains objects that are not picklable, or if you expect thread-based concurrency.
Fix: If you encounter `TypeError: cannot pickle ...` or unexpected performance, explicitly set `reader_pool_type='thread'` in your `make_reader` call: `make_reader(..., reader_pool_type='thread', ...)`.
deprecated The `PetastormDataset` class (e.g., from `petastorm.reader`) and direct instantiation of `Reader` were deprecated in favor of the `make_reader` factory function.
Fix: Always use `from petastorm import make_reader` and instantiate readers via `make_reader(...)` for future compatibility and resource management.
gotcha Using Petastorm with TensorFlow or PyTorch requires installing the corresponding 'extras' (e.g., `pip install petastorm[tensorflow]`). Without these, you might miss framework-specific utilities or experience integration issues.
Fix: Ensure you install Petastorm with the relevant extras for your ML framework: `pip install petastorm[tensorflow]` or `pip install petastorm[pytorch]`.
gotcha When using `make_reader` with Spark, ensure `pyspark` is installed and the `spark` extra is included during Petastorm installation. Otherwise, Spark-specific modules will be missing.
Fix: Install `pyspark` and Petastorm with the Spark extra: `pip install pyspark petastorm[spark]`.

Install

pip install petastorm Base installation
pip install petastorm[tensorflow] For TensorFlow integration
pip install petastorm[pytorch] For PyTorch integration
pip install petastorm[spark] For Apache Spark integration

Imports

make_reader
```
from petastorm import make_reader
```
make_writer
```
from petastorm import make_writer
```

Unischema

from petastorm.unischema import Unischema

DataLoader (PyTorch)

from petastorm.pytorch import DataLoader

SparkDatasetConverter
wrong:
```
from petastorm.etl.spark_dataset_converter import SparkDatasetConverter
```
correct:
```
from petastorm.spark import SparkDatasetConverter
```
The `spark` module was moved directly under `petastorm` for easier access as of v0.10.0.
Reader (direct import)
wrong:
```
from petastorm.reader import Reader
```
correct:
```
from petastorm import make_reader
```
`make_reader` is the recommended factory function for creating readers and managing resources, especially since `Reader` from `petastorm.reader` was deprecated/removed in v0.13.0.

Quickstart

This quickstart demonstrates how to define a data schema, write sample data to a Parquet dataset using `make_writer`, and then read it back using `make_reader`. The example cleans up the temporary directory after execution. For real-world usage, consider configuring `reader_pool_type` and `num_epochs` based on your training requirements.

import os
import shutil
import numpy as np
from petastorm import make_reader, make_writer
from petastorm.unischema import Unischema, UnischemaField, ScalarCodec
from petastorm.codecs import CompressedNdarrayCodec

# 1. Define a schema for your data
MySchema = Unischema(
    'MySchema',
    [
        UnischemaField('id', np.int32, (), ScalarCodec(np.int32), False),
        UnischemaField('value', np.float64, (), ScalarCodec(np.float64), False),
        UnischemaField('image', np.uint8, (10, 10, 3), CompressedNdarrayCodec(), False),
    ]
)

# 2. Define a dataset path (using a temporary local directory for example)
dataset_url = 'file:///tmp/petastorm_example_data'
# Clean up previous data if it exists
if os.path.exists('/tmp/petastorm_example_data'):
    shutil.rmtree('/tmp/petastorm_example_data')

# 3. Write some dummy data to the Parquet dataset
print(f"Writing dummy data to {dataset_url}...")
with make_writer(dataset_url, MySchema, row_group_size_bytes=2 * 1024 * 1024) as writer:
    for i in range(10):
        writer.write(
            MySchema.make_row(
                id=i,
                value=float(i * 10),
                image=np.random.randint(0, 256, size=(10, 10, 3), dtype=np.uint8)
            )
        )
print(f"Successfully wrote 10 rows.")

# 4. Read data using make_reader
# reader_pool_type='thread' is often suitable for local development.
# For production, 'process' might be preferred depending on data access patterns.
print("\nReading data from the dataset:")
with make_reader(dataset_url, reader_pool_type='thread', num_epochs=1) as reader:
    for i, row in enumerate(reader):
        print(f"Row {i}: id={row.id}, value={row.value}, image_shape={row.image.shape}")
        if i >= 2: # Print only a few rows for brevity
            break
print("Finished reading example data.")

# Clean up the temporary dataset
shutil.rmtree('/tmp/petastorm_example_data')

view raw JSON →