feat(interleaved): add InterleavedParquetReader and InterleavedWebdatasetWriter - #1657
Merged
VibhuJawa merged 25 commits intoApr 3, 2026
Merged
Conversation
…e-class fixes, and bug fixes Splits out the foundation layer from PR NVIDIA-NeMo#1570 as a standalone PR: Schema utilities (utils/schema.py — new): - reconcile_schema(): canonical types for reserved columns, passthrough columns preserved - align_table(): pad, reorder, cast tables to a target schema; safe=True for passthrough columns to surface overflow errors instead of silently corrupting data (B5 fix) Reader base (readers/base.py): - Add schema= and schema_overrides= parameters for strict schema alignment - _resolve_schema() helper: schema > schema_overrides > INTERLEAVED_SCHEMA priority - Warn (not silently ignore) when both schema= and schema_overrides= are provided (C5) - Fix return type annotation: -> pa.Schema | None (C4) - Preserve nullable=False from INTERLEAVED_SCHEMA when applying schema_overrides (C6) Writer base (writers/base.py): - Add schema= and schema_overrides= for output alignment (same API as reader) - Add on_materialize_error= parameter: "error"|"warn"|"drop_row"|"drop_sample" - _materialize_dataframe(): apply error policy after fetch; guard binary_content column - Make _write_dataframe() non-abstract (default pass body) to allow WDS-style overrides (D3) Materialization (utils/materialization.py): - _build_global_range_index(): group paths by filesystem object so mixed-backend batches work correctly (B3 fix) WebDataset reader (readers/webdataset.py): - Rename tiff result variable to tiff_frame to remove variable shadowing (D5 fix) Parquet reader (readers/parquet.py): - Guard pq.read_schema() call behind if self.fields is not None to avoid N+1 schema reads in the common case (C1 fix) Tests: - test_base_writer.py (new): 9 tests covering on_materialize_error modes, schema alignment, and write_data orchestration - test_materialization.py: 4 new tests for multi-backend path grouping (B3) - test_multimodal_core.py: 4 new align_table tests for safe-cast behaviour (B5) - test_multimodal_reader.py: schema_overrides and nullable preservation tests (C5/C6) - conftest.py: additional helpers for new test patterns Note: InterleavedParquetReader, InterleavedParquetReaderStage, and InterleavedWebdatasetWriterStage are intentionally excluded from this branch; they will be added in the follow-up PR on feat/interleaved-io-readers-writers. Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
…ad code, pre-compute write_kwargs Move resolve_schema to utils/schema.py: - Removes cross-module readers->writers private import (_resolve_schema was defined in readers/base.py and imported by writers/base.py) - Renamed _resolve_schema to resolve_schema (public shared utility) - Returns None when both args are None instead of unreachable ValueError writers/base.py: - Pre-compute _effective_write_kwargs in __post_init__ (filters storage_options, forces index=False) so write_data() no longer copies the dict on every task - Remove noisy section-header comments Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
…oundation Signed-off-by: Vibhu Jawa <vjawa@nvidia.com> # Conflicts: # .github/workflows/config/.secrets.baseline
This override was added to work around ray[llm]'s unconditional nixl dep on ARM but is not part of the interleaved IO foundation work. Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
Reverts the removal — nixl-cu12 override is needed to work around ray[llm]'s unconditional nixl dep on ARM. Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
source_id_field was only used as a passthrough exclusion hint — the same effect is achievable via the fields= parameter. Having both source_id_field and sample_id_field was confusing with overlapping names. Removing source_id_field simplifies the API with no loss of functionality. Update all tests, benchmarking scripts, tutorials, and README to remove all references to source_id_field. Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
- align_table: remove redundant reconcile_schema(target) call; the safe/unsafe cast logic already handles large<->small conversions for reserved columns, and re-reconciling a user-provided schema would silently override intentional type overrides - _scatter_range_blobs: restore specific tuple type hints (list[tuple[str,int,int]]) lost when range key shape changed from (offset,size) to (fs_path,offset,size) - reconcile_schema: add comment explaining why dictionary-encoded passthrough columns are unwrapped to their value type Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
…der metadata - Fix interleaved_data_quickstart.ipynb: - max_aspect_ratio 1.0 -> 2.0 (was dropping 180 rows instead of 22) - write filtered_batch instead of unfiltered batch in Step 5 - Fix InterleavedWebdatasetReaderStage.process(): populate source_files in output metadata from task.data so the writer can use a deterministic hash-based filename instead of a UUID Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Vibhu Jawa <vibhujawa@gmail.com>
The interleaved_data_quickstart.ipynb outputs contain base64-encoded PNG images and hex sha256 hashes flagged by detect-secrets. These are all false positives from inline cell output, not real credentials. Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
When max_batch_bytes splits a single tar into N batches, each batch
previously got the same source_files=[tar_path], making splits
indistinguishable by source_files alone.
Append '::split_{idx:05d}' to each path when there are multiple splits,
so each batch's source_files uniquely identifies both the source tar and
its position within the split sequence. Single-split output is unchanged.
Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
…t reader When max_batch_bytes causes the reader to emit multiple InterleavedBatch splits, each split's source_files metadata now lists only the tar files that actually contributed rows to that split, instead of all source tars. Also optimises the hot path: sample_id is recorded once per member (from the first/metadata row) rather than iterating every row, and unique() deduplicates the sample_id column in Arrow before conversion. Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
Contributor
Author
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
- Add test_schema_utils.py: 8 tests for resolve_schema, reconcile_schema, align_table - Add 5 tests in test_multimodal_core.py: aspect ratio edge cases, iter_materialized_bytes paths - Add inputs/outputs contract test in test_base_writer.py - Remove dead metadata-propagation code from utils/schema.py Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
…oundation Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
…ebdatasetWriterStage Completes the interleaved IO round-trip by adding the two missing directions: - InterleavedParquetReaderStage: reads Parquet → InterleavedBatch with fields= tuple for passthrough column selection, null-fill for missing columns, max_batch_bytes splitting, and per-split source_files tracking. - InterleavedParquetReader: composite stage wrapping FilePartitioningStage + InterleavedParquetReaderStage. - InterleavedWebdatasetWriterStage: writes InterleavedBatch → MINT-1T-style tar shards, reconstructing texts/images lists from row-based data with percent-encoded tar member keys for safe roundtrip via sample_id_field. - _source_files_for_split() moved to BaseInterleavedReader so both WDS and Parquet readers share the per-split source tracking logic. Tests: 15 new tests covering roundtrip, missing columns, fields null-fill, max_batch_bytes splits, key escaping, passthrough columns, and all on_materialize_error modes. Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
…etReader Aligns naming with the Interleaved* prefix convention used by all other stages in this module. InterleavedParquetReader was already correctly named; this makes the WebDataset composite reader consistent. Also updates README.md architecture diagram to show both composite readers (InterleavedWebdatasetReader and InterleavedParquetReader) side-by-side, reflecting the full round-trip pipeline. Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
…chmark - Add collect_parquet_input_metrics and collect_wds_input_metrics to utils.py - Add shared _accumulate_modality_counts helper to eliminate duplicated loop - Collect input metrics before pipeline.run() in multimodal_mint1t_benchmark.py - Add run_mint1t_benchmark_all.sh runner for all 4 paths to /raid/vjawa/benchamrk_ouput_dir Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
…uet-reader-wds-writer Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
Add _count_wds_samples helper that counts .json members across tars (one per sample). Used by collect_wds_input_metrics and collect_wds_output_metrics so WDS paths report sample counts and throughput alongside parquet paths. Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
abhinavg4
reviewed
Apr 2, 2026
abhinavg4
reviewed
Apr 2, 2026
abhinavg4
reviewed
Apr 2, 2026
- Replace 4 narrow input/output collect functions with 2 format-scoped functions (collect_interleaved_parquet_metrics, collect_interleaved_wds_metrics) using neutral keys; caller adds input_/output_ prefix via dict comprehension - Add _resolve_paths helper to unify single-file vs directory resolution - Track num_samples (distinct sample_ids), num_metadata, num_texts, num_images, materialize_error_count in parquet metrics; text/image counts in WDS metrics - Add validate_wds_ordering(tar_path) single-file spot-checker with bunching detection, position collision, and missing image validation - Fix _check_wds_sample return to tuple[int, bool] | None (None replaces -1) - Add _validate_output helper with clean bool | None semantics - Add wds_valid and throughput_samples_per_sec to benchmark metrics output Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
… iterrows with itertuples - InterleavedParquetReaderStage and InterleavedWebdatasetReaderStage now cache resolve_storage_options() in __post_init__ instead of calling it on every process() invocation; read_kwargs is immutable after construction so the result never changes across tasks - Replace iterrows() with itertuples(index=False) in WDS writer _write_sample; all accessed columns (position, modality, text_content, content_type, binary_content) are RESERVED_COLUMNS guaranteed present after _align_output(); ~7%% process-time improvement measured on single-shard benchmark Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
feat(interleaved): add InterleavedParquetReader and InterleavedWebdatasetWriter
Overview
Completes the interleaved IO round-trip by adding the two missing directions:
With the existing WDS reader and Parquet writer on
main, all four IO combinations are now supported:New classes
InterleavedParquetReaderStageio/readers/parquet.pyInterleavedBatchInterleavedParquetReaderio/reader.pyInterleavedWebdatasetWriterStageio/writers/webdataset.pyInterleavedBatch→ MINT-1T-style tar shardsKey design decisions
InterleavedParquetReaderStagefields=tuple selects passthrough columns — consistent with WDS reader; reserved columns always read; missing columns null-filled in a single passpq.read_schema()per file for push-down column projection (avoids reading unnecessary data from disk/S3)max_batch_bytessplitting viasplit_table_by_group_max_bytes; each split'ssource_fileslists only its contributing parquet files_source_files_for_split()moved toBaseInterleavedReaderso both WDS and Parquet readers share per-split lineage trackingInterleavedWebdatasetWriterStage_write_dataframe(); base class handles deterministic naming, materialization, schema alignment, andprocess()orchestrationurllib.parse.quote(sample_id, safe="")key escaping — injective, roundtrip-safe viasample_id_field="sample_id"on readdf.groupby("sample_id", sort=False)— O(n) single pass instead of O(n×m) per-sample filtermax_pos + 1withNoneat gaps (WDS reader skipsNone)metadata,text,image— any other raisesValueErrorat write timeBenchmarks
All four IO paths benchmarked end-to-end on 80 local NVMe shards (6,818 samples, MINT-1T PDF data) with an aspect-ratio filter applied. Full results and reproduction instructions: benchmark gist.
Key takeaway: Parquet-sourced paths are ~5× faster than WDS-sourced paths. The bottleneck is entirely in the WDS reader (396 s cumulative stage time across 80 shards vs 14 s for PQ) due to sequential tar scanning and network image fetches. Filter and writer costs are near-identical across formats.
Stage-level breakdown (cumulative across 80 tasks, 16 workers)
Environment: Python 3.10.12 · Ray 2.54.0 · PyArrow 19.0.1 · 16-core CPU · NVMe ·
materialize_on_read=TrueTests (15 new)
Parquet reader (
test_multimodal_reader.py)test_parquet_reader_roundtrip— write then read back; sample_id, modality, text_content matchtest_parquet_reader_missing_columns_filled_with_null— 3-column file; all reserved cols null-filledtest_parquet_reader_fields_subset—fields=selects extra cols; reserved always presenttest_parquet_reader_fields_null_fill_missing—fields=("nonexistent",)→ null column, no errortest_parquet_reader_max_batch_bytes_splits— 2 files → 2 splits; each lists only its contributing file insource_filestest_parquet_reader_empty_file— 0-row parquet → empty batch, correct schema, no crashtest_parquet_reader_composite_decompose—decompose()→[FilePartitioningStage, InterleavedParquetReaderStage]WDS writer (
test_multimodal_writer.py)test_escape_key_encodes_special_chars—a/b:c→a%2Fb%3Actest_ext_from_content_type_known/_fallback— MIME → extension mappingtest_wds_writer_roundtrip— write then read back viaInterleavedWebdatasetReaderStagetest_wds_writer_text_only_sample— no image rows; JSON has all-None images listtest_wds_writer_unsupported_modality_raises—modality="video"→ValueErrortest_wds_writer_key_escaping— special-char sample_id; roundtrip recovers original idtest_wds_writer_passthrough_columns_in_json— extraurlcol in metadata row → appears in JSON payloadtest_wds_writer_null_binary_skips_member— nullbinary_content→ no image tar member writtentest_wds_writer_deterministic_filename— samesource_files+task_id→ same output filename