feat(interleaved): Foundation IO improvements — schema utilities, base-class fixes, and bug fixes by VibhuJawa · Pull Request #1652 · NVIDIA-NeMo/Curator

VibhuJawa · 2026-03-24T13:41:33Z

Summary

This is Part 1 of 2 splitting PR #1570. It contains only the foundation layer improvements (base classes, schema utilities, bug fixes) with no new reader/writer stages. Part 2 (new `InterleavedParquetReader` + `InterleavedWebdatasetWriter`) will follow as a separate PR building on top of this.

What's included

New: `utils/schema.py`

Centralized schema utilities shared by all arrow-based readers and writers:

`reconcile_schema()` — canonical types for reserved columns; passthrough columns preserved as-is; avoids unsafe large↔small downcasts; unwraps Parquet dictionary encoding from passthrough columns
`align_table()` — pad, reorder, and cast an Arrow table to a target schema; reserved columns use `safe=False` for large↔small, passthrough columns use `safe=True` to surface overflow rather than silently corrupt data; does not re-reconcile the user-provided target (which would silently override intentional type choices)
`resolve_schema()` — merges `schema=` and `schema_overrides=` with `INTERLEAVED_SCHEMA`; warns when both are provided; returns `None` when neither is set

Reader base (`readers/base.py`)

Add `schema=` and `schema_overrides=` parameters for strict output alignment
Warn when both `schema=` and `schema_overrides=` are provided (instead of silently ignoring overrides)
Preserve `nullable=False` from `INTERLEAVED_SCHEMA` when applying `schema_overrides`

Writer base (`writers/base.py`)

Same `schema=` / `schema_overrides=` API as the reader
New `on_materialize_error=` parameter: `"error"` | `"warn"` | `"drop_row"` | `"drop_sample"`
`_materialize_dataframe()` applies the error policy after fetch, covering errors set by both the fetch step and upstream stages
`_write_dataframe()` changed from `@abstractmethod` to a default `NotImplementedError` body, allowing WDS-style subclasses that override `write_data()` directly without a dead stub
Pre-compute `_effective_write_kwargs` in `post_init` (filters `storage_options`, forces `index=False`) instead of copying the dict on every task

Materialization (`utils/materialization.py`)

Extract `_build_global_range_index()`: group paths by filesystem object so mixed-backend batches (e.g. S3 + local) work correctly instead of failing silently; restore `list[tuple[str, int, int]]` type hints on `_scatter_range_blobs`
Fix `_extract_tiff_frame()`: copy frame before closing the PIL context to avoid use-after-close; do not reuse source JPEG compression which corrupts RGBA frames

WebDataset reader (`readers/webdataset.py`)

Remove `source_id_field` — it only excluded a JSON key from passthrough columns, achievable via `fields=` already. Having both `source_id_field` and `sample_id_field` with overlapping names was confusing with no benefit
Rename class from `WebdatasetReaderStage` → `InterleavedWebdatasetReaderStage` for consistency
Rename tiff result variable from `extracted` to `tiff_frame` to eliminate variable shadowing
Fix `source_files` metadata per split: when `max_batch_bytes` causes multiple output batches, each batch's `source_files` now lists only the tar files that actually contributed rows to that split (instead of all source tars). Tracks `sample_id → tar` during reading; deduplicates with `unique()` in Arrow before resolving.

`pyproject.toml`

Add `multimodal_cpu` extras group: `Pillow` and `s3fs>=2024.12.0`

Tests added / updated

`test_base_writer.py` (new, 9 tests): `on_materialize_error` modes, schema alignment, write orchestration
`test_materialization.py`: 4 new tests for multi-backend path grouping
`test_multimodal_core.py`: 4 new `align_table` tests for safe-cast behaviour
`test_multimodal_reader.py`: `schema_overrides` and `nullable` preservation tests; remove all `source_id_field=` usage; new test `test_reader_source_files_per_split_only_contributing_tars` verifying per-split source file accuracy

149 / 149 tests pass. Ruff clean.

What's NOT included (Part 2)

`InterleavedParquetReader` / `InterleavedParquetReaderStage` — new composite + stage for reading Parquet
`InterleavedWebdatasetWriter` / `InterleavedWebdatasetWriterStage` — new composite + stage for writing WebDataset tar archives

Test plan

`python -m pytest tests/stages/interleaved/ -v` → 149 passed
`python -m ruff check nemo_curator/stages/interleaved/ tests/stages/interleaved/` → all checks passed
Pre-commit hooks pass (ruff, signed-off-by, large-file, private-key)

…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>

…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>

- 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>

abhinavg4 · 2026-04-01T19:09:52Z

    )
    pipeline.add_stage(
        WebdatasetReader(
-            source_id_field="pdf_name",


I think this change will affect pdf PR too. Why was this removed? Is the reading happening via schema now ?

I just removed it , because it was excessive to have around. It now is just a passthrough column and not given special treatment.

I think if you remove it and its in the dataset, it will just work. Happy to help with your PR too.

We will revisit this discussion in 1657

abhinavg4

Looks good. left some minor comments.

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

copy-pr-bot Bot temporarily deployed to test March 24, 2026 13:42 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci March 24, 2026 13:42 Failure

copy-pr-bot Bot temporarily deployed to nemo-ci March 24, 2026 13:42 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci March 24, 2026 14:38 Inactive

VibhuJawa added 2 commits March 24, 2026 07:40

Merge remote-tracking branch 'origin/main' into feat/interleaved-io-f…

a1eeccc

…oundation Signed-off-by: Vibhu Jawa <vjawa@nvidia.com> # Conflicts: # .github/workflows/config/.secrets.baseline

chore: remove unrelated nixl-cu12 from override-dependencies

33088b1

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>

copy-pr-bot Bot temporarily deployed to test March 24, 2026 14:48 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci March 24, 2026 14:48 Failure

VibhuJawa and others added 8 commits March 24, 2026 07:53

chore: restore nixl-cu12 override-dependency

6f0e655

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>

Apply suggestion from @claude[bot]

575e505

Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Vibhu Jawa <vibhujawa@gmail.com>

This was referenced Mar 24, 2026

feat(interleaved): add InterleavedParquetReader and InterleavedWebdatasetWriter VibhuJawa/NeMo-Curator#5

Closed

feat(interleaved): add InterleavedParquetReader and InterleavedWebdatasetWriter #1657

Merged

VibhuJawa added 3 commits March 30, 2026 15:34

Merge remote-tracking branch 'origin/main' into feat/interleaved-io-f…

9c35387

…oundation Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

abhinavg4 reviewed Apr 1, 2026

View reviewed changes

Comment thread .github/workflows/config/.secrets.baseline

abhinavg4 reviewed Apr 1, 2026

View reviewed changes

abhinavg4 approved these changes Apr 1, 2026

View reviewed changes

Merge branch 'main' into feat/interleaved-io-foundation

802ab23

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

thomasdhc reviewed Apr 1, 2026

View reviewed changes

Comment thread pyproject.toml

VibhuJawa and others added 2 commits April 1, 2026 14:28

feat: add interleaved_cpu to all extras in pyproject.toml

5a82f0e

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

Merge branch 'main' into feat/interleaved-io-foundation

54d1378

lbliii mentioned this pull request Apr 3, 2026

docs: add fern docs for 26.04 PRs (#1160, #1575, #1576, #1603, #1612, #1652) #1730

Closed

3 tasks

Navigation Menu

Uh oh!

feat(interleaved): Foundation IO improvements — schema utilities, base-class fixes, and bug fixes - #1652

feat(interleaved): Foundation IO improvements — schema utilities, base-class fixes, and bug fixes#1652
VibhuJawa merged 18 commits into
NVIDIA-NeMo:mainfrom
VibhuJawa:feat/interleaved-io-foundation

VibhuJawa commented Mar 24, 2026 •

edited

Loading

Uh oh!

Uh oh!

abhinavg4 Apr 1, 2026

Uh oh!

VibhuJawa Apr 1, 2026

Uh oh!

VibhuJawa Apr 1, 2026

Uh oh!

VibhuJawa Apr 1, 2026

Uh oh!

abhinavg4 left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

VibhuJawa commented Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What's included

New: `utils/schema.py`

Reader base (`readers/base.py`)

Writer base (`writers/base.py`)

Materialization (`utils/materialization.py`)

WebDataset reader (`readers/webdataset.py`)

`pyproject.toml`

Tests added / updated

What's NOT included (Part 2)

Test plan

Uh oh!

Uh oh!

abhinavg4 Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

VibhuJawa Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

VibhuJawa Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

VibhuJawa Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

abhinavg4 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

VibhuJawa commented Mar 24, 2026 •

edited

Loading