$ show uber/petastorm
Petastorm is a library for loading Apache Parquet datasets into machine learning training frameworks, primarily TensorFlow and PyTorch. It provides a reader API that supports parallel I/O, row-group shuffling, predicate filtering, and on-the-fly transformations. The library also includes a Spark Dataset Converter to materialize Spark DataFrames as cached Parquet for training. Architecture: modular with worker pools (process/thread), filesystem abstraction (local, HDFS, S3, GCS), schema system (Unischema with codecs), and shuffling buffers. It features CLI tools for benchmarking, metadata generation, and dataset copying. The codebase demonstrates strong engineering practices with clear separation of concerns, but shows signs of aging: reliance on TensorFlow 1.x, Python 2 compatibility, and pinned dependencies that may cause drift. The project has a large test suite with legacy dataset fixtures for backward compatibility, but maintenance has slowed significantly with the last release (v0.13.1) and minimal recent commits.
The LLM identified 3 files worth highlighting for reuse:
Implements a process-based worker pool using ZMQ for inter-process communication. Provides start, ventilate (distribute work items), get_results (collect results), and stop lifecycle. The ventilator pattern decouples work distribution from result collection, enabling deterministic ordering and throttling. This module is self-contained with minimal internal dependencies (imports from workers_pool.ventilator and workers_pool.worker_base) and could be extracted as a standalone parallel task executor.
Provides an abstract shuffling buffer base class (ShufflingBufferBase) and concrete implementations: RandomShufflingBuffer and NoopShufflingBuffer. The pattern enables decorrelation of row groups during training data loading, with configurable buffer size and support for both row-level and batch-level shuffling (BatchedRandomShufflingBuffer). The abstraction is clean, isolated, and can be reused in any streaming data pipeline requiring controlled shuffling.
Implements FilesystemResolver, a factory that creates a PyArrow filesystem from a dataset URL, supporting local, HDFS, S3, GCS, and ADLS via fsspec. Handles path resolution, file globbing, and serialization for distributed workers. The resolver pattern abstracts storage backend details behind a simple interface, making it reusable in any project that needs to read from multiple filesystems with a consistent API.