Align interleaved source locators with Parquet FILE

[VibhuJawa]

Summary

Align interleaved source locators with the Parquet FILE logical-type shape, simplify materialization, and improve actual Parquet reader/writer throughput.

Previous field	New location
source_ref.path	source_ref.uri
source_ref.byte_offset	source_ref.offset
source_ref.byte_size	source_ref.size
source_ref.member	adjacent source_member column
source_ref.frame_index	adjacent source_frame_index column

source_ref has all six optional FILE children from the specification: uri, offset, size, content_type, checksum, and inline.

What changed

• WebDataset emits a typed FILE-compatible struct, preserving canonical Arrow child order.
• FILE_REFERENCE_TYPE lives in shared nemo_curator.utils.storage_utils for reuse by text, image, video, and interleaved stages.
• Existing JSON-string locators are migrated once at the shared schema-alignment boundary.
• Legacy Nemotron {page, bbox} provenance migrates to typed adjacent page_number and source_bbox columns instead of being discarded.
• uri + size now correctly resolves [0, size) when offset is absent.
• Whole-object and ranged references share one deduplicated fs.cat_ranges() path, with one call per filesystem backend.
• Parquet output is opened through the configured fsspec filesystem, preserving remote protocols and storage_options.
• Arrow inputs bypass redundant Arrow → pandas → Arrow conversion when materialization/error filtering is unnecessary.
• Dictionary encoding defaults to low-cardinality interleaved fields; callers can still override use_dictionary.

Archive-member and TIFF-frame metadata remain adjacent because FILE is a closed group. source_member is only needed for tar extraction when a usable FILE byte range is unavailable; ranged and whole-object FILE reads do not depend on it.

Compatibility and scope

The change keeps binary_content unchanged. The inline FILE child is present in the canonical schema, but moving existing payload storage or materialization to inline is outside this PR.

PyArrow 19 can write the exact physical group but cannot emit the new LogicalType.FILE footer annotation. The FILE specification merged on 2026-07-26, and the cited parquet-java and arrow-rs implementations are still open/unreleased. This PR therefore avoids a custom Thrift footer rewriter and can adopt native Arrow support when it lands.

• Parquet FILE specification
• parquet-java proof of concept
• arrow-rs proof of concept
• cuDF FILE tracking
• cuDF binary/inline prerequisite

Performance validation

Compared latest upstream main (5732455) with this branch on the same CPU host. These are real InterleavedParquetWriterStage and InterleavedParquetReaderStage runs over 200,000 typed Arrow rows; values are medians of seven warm runs.

Metric	Upstream main	This PR	Change
Parquet write throughput	4.78M rows/s	6.24M rows/s	+30.7%
Parquet read throughput	3.43M rows/s	4.03M rows/s	+17.6%
Parquet file size	3,128,791 bytes	1,592,172 bytes	−49.1%

Remote HTTP materialization used a pinned, range-capable 466 KB object and verified every returned payload:

• 32 rows / 8 unique byte ranges: combined alternating-run median was 0.109 s here versus 0.115 s upstream (network-sensitive, no regression).
• 32 duplicate whole-object references: combined median was 0.174 s here versus 0.273 s upstream; individual runs were network-noisy, but the unified path showed no consistent slowdown.

Validation

• Ruff passes across nemo_curator and interleaved tests.
• 208 top-level interleaved CPU tests pass, covering locator and PDF-provenance migration, FILE semantics, size-only ranges, materialization, fsspec output, Parquet/WebDataset/Lance readers and writers, and round trips.
• Latest fork main, upstream main, and the PR base all resolve to 5732455; the branch is based directly on current main.

Diff

• 19 files changed
• 356 additions / 359 deletions
• net reduction of 3 lines

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

The PR replaces JSON-string interleaved source locators with a Parquet FILE-compatible struct while moving archive, frame, and PDF provenance into adjacent typed columns. It also consolidates legacy migration at schema alignment and streamlines materialization and Parquet I/O.

• Emits canonical FILE-shaped source references from WebDataset ingestion.
• Migrates legacy locator and Nemotron provenance fields during shared schema alignment.
• Batches and deduplicates whole-object and ranged materialization requests.
• Writes Arrow tables directly through the configured fsspec filesystem.

Confidence Score: 5/5

The PR appears safe to merge because the previously reported compatibility and provenance failures are resolved at the current head.

No blocking failure remains.

Important Files Changed

Filename	Overview
nemo_curator/stages/interleaved/utils/schema.py	Centralizes migration of legacy JSON locators and PDF provenance before schema reconciliation; the previously reported compatibility and provenance failures are addressed.
nemo_curator/tasks/interleaved.py	Defines the canonical FILE-compatible source_ref shape and adjacent archive/frame metadata columns, with current callers updated to the replacement API.
nemo_curator/stages/interleaved/pdf/nemotron_parse/utils.py	Preserves Nemotron page and bounding-box provenance in typed adjacent columns rather than overloading source_ref.
nemo_curator/stages/interleaved/utils/materialization.py	Updates materialization to consume FILE locators and adjacent metadata while sharing deduplicated filesystem range requests.
nemo_curator/stages/interleaved/io/writers/base.py	Adds a direct Arrow-table preparation path while retaining materialization-error handling when conversion is necessary.
nemo_curator/stages/interleaved/io/writers/tabular.py	Writes Parquet through the configured fsspec filesystem and applies default dictionary encoding to selected low-cardinality fields.
nemo_curator/stages/interleaved/io/readers/webdataset.py	Emits typed FILE references and stores tar-member and TIFF-frame metadata in adjacent canonical columns.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Legacy JSON source_ref] --> B[align_interleaved_table]
  C[Typed WebDataset locator] --> B
  B --> D[FILE-compatible source_ref]
  B --> E[source_member]
  B --> F[source_frame_index]
  B --> G[page_number and source_bbox]
  D --> H[Materialization]
  E --> H
  F --> H
  D --> I[Parquet writer via fsspec]
  E --> I
  F --> I
  G --> I

Loading

_{Reviews (9): Last reviewed commit: "refactor(schema): share Parquet FILE typ..." | Re-trigger Greptile}

[VibhuJawa]

/okay to test 7a5a8e9

[VibhuJawa]

/claude review

[Copilot]

Pull request overview

This PR migrates interleaved provenance from a JSON-string source_ref locator to a Parquet FILE-shaped source_ref struct (with adjacent source_member / source_frame_index), and updates schema alignment, materialization, and I/O stages to improve correctness and throughput while preserving backward compatibility via a one-time migration at the schema boundary.

Changes:

• Introduces a canonical Arrow FILE_REFERENCE_TYPE and updates INTERLEAVED_SCHEMA to store source_ref as a FILE-compatible struct plus adjacent metadata columns.
• Adds align_interleaved_table() migration/casting logic and updates materialization + readers/writers to operate on typed Arrow tables and fsspec-backed output.
• Updates tests and docs to reflect the new locator shape and semantics (including PDF provenance migration to adjacent columns).

Reviewed changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/stages/interleaved/test_schema_utils.py	Adds coverage for legacy JSON source_ref migration into typed FILE struct + adjacent columns.
tests/stages/interleaved/test_multimodal_writer.py	Updates writer tests for typed source_ref, adjacent member/frame fields, and fsspec-backed Parquet output.
tests/stages/interleaved/test_multimodal_reader.py	Updates reader assertions to validate source_member / source_frame_index instead of parsing JSON.
tests/stages/interleaved/test_multimodal_core.py	Adjusts core materialization/classification tests to the FILE + adjacent metadata model and unified range-read path.
tests/stages/interleaved/test_materialization.py	Extends edge-case coverage (e.g., size-only reads; empty-range acceptance) under new range-read semantics.
tests/stages/interleaved/test_interleaved_task.py	Verifies canonical FILE child ordering and validates range constraints in build_source_ref.
tests/stages/interleaved/test_base_writer.py	Updates base-writer tests to _write_table and Arrow-table alignment outputs.
tests/stages/interleaved/pdf/nemotron_parse/test_utils.py	Adds assertions for migrated PDF provenance fields (page_number, source_bbox) with source_ref=None.
tests/stages/interleaved/conftest.py	Updates fixtures/helpers to construct typed FILE source_ref and adjacent metadata columns.
nemo_curator/tasks/interleaved.py	Defines FILE_REFERENCE_TYPE, updates INTERLEAVED_SCHEMA, and replaces JSON helpers with FILE struct helpers/expansion.
nemo_curator/stages/interleaved/utils/schema.py	Adds align_interleaved_table() to migrate/cast legacy locator shapes to canonical FILE struct and append adjacent columns.
nemo_curator/stages/interleaved/utils/materialization.py	Consolidates whole-file and ranged reads into deduped fs.cat_ranges() flow; updates classification inputs/columns.
nemo_curator/stages/interleaved/README.md	Documents new FILE-shaped source_ref and adjacent metadata strategy; updates materialization strategy description.
nemo_curator/stages/interleaved/pdf/nemotron_parse/utils.py	Migrates Nemotron provenance from embedded JSON source_ref to adjacent page_number / source_bbox.
nemo_curator/stages/interleaved/io/writers/webdataset.py	Moves writer implementation to _write_table and converts Arrow → pandas only at the boundary.
nemo_curator/stages/interleaved/io/writers/tabular.py	Switches Parquet writing to pq.write_table through the configured fsspec filesystem and sets dictionary defaults.
nemo_curator/stages/interleaved/io/writers/base.py	Refactors base writer to prepare/align Arrow tables directly and fast-path Arrow inputs when materialization is unnecessary.
nemo_curator/stages/interleaved/io/readers/webdataset.py	Emits typed FILE-compatible source_ref scalars and populates adjacent source_member / source_frame_index.

[claude]

This .tar.gz/.tgz guard is a subtle correctness fix — info.offset_data/info.size are positions in the decompressed stream, so emitting them as a FILE byte range would make materialization cat_ranges() read garbage from the compressed file. Good catch, but it's new behavior with no direct test. Consider adding a case that reads image members from a gzipped tar and asserts source_ref carries no offset/size (falls back to tar-extract), so a future refactor doesn't silently reintroduce the range read.