Interleaved Lance Reader

[suiyoubi]

Summary

Adds an InterleavedLanceReader for row-wise interleaved multimodal Lance datasets.

The reader builds on Curator's existing LancePartitioningStage and
LanceReaderStage, then splits each resulting Arrow table into validated
InterleavedBatch outputs while preserving row order and sample_id
document boundaries.

Changes

• Added nemo_curator.stages.interleaved.lance.InterleavedLanceReader
• Added InterleavedLanceReaderStage
• Reuses LanceReaderStage.read_task() for:
  • Dataset and scanner configuration
  • Pinned Lance versions
  • Blob materialization
  • Optional Lance metadata columns
• Supports document-aware output batching with:
  • max_batch_bytes
  • max_batch_rows
• Supports:
  • fragment_ids
  • fragments_per_partition
  • Field projection and Lance read options
• Preserves source row order and never splits consecutive rows belonging
  to the same sample_id
• Validates required InterleavedBatch fields and rejects null sample_id
  values
• Exports the reader through the interleaved stage package

Tests

Added unit coverage for:

• Required field validation
• Reader decomposition/config propagation
• Invalid batch/output limits
• Row-count based splitting without splitting sample_id
• Streaming reads where one sample spans scanner batches
• max_output_rows behavior without emitting partial samples
• Lance metadata columns
• Rejection of unsupported streaming blob columns

[greptile-apps]

Greptile Summary

This PR adds InterleavedLanceReader and InterleavedLanceReaderStage for reading row-wise interleaved multimodal Lance datasets, refactors split_table_by_group_max_bytes into split_table_by_group (adding max_batch_rows support and null-value rejection), extracts reusable helper methods from LanceReaderStage, and adds align_interleaved_table to the schema utilities.

• New Lance interleaved reader: InterleavedLanceReaderStage delegates to the existing LanceReaderStage.read_task() for dataset open/scan, then calls split_table_by_group to partition the resulting Arrow table into InterleavedBatch outputs without splitting same-sample_id rows across batches.
• split_table_by_group refactor: The renamed utility now validates null sample_id values and enforces row-count limits in addition to byte limits; crucially, null validation now fires unconditionally (even when no batch limits are set), which is a behavioral change that also affects the existing Parquet and WebDataset readers.
• LanceReaderStage refactor: _requested_blob_columns, _dataset_and_scanner_kwargs, and _metadata_for_task are extracted as reusable helper methods, making the logic available to the new interleaved stage without duplication.

Confidence Score: 5/5

Safe to merge — the new Lance interleaved reader correctly delegates to the existing read path, and null sample_id values are now rejected at the split boundary rather than silently merged.

The core read-and-split logic is straightforward, reusing the well-tested LanceReaderStage.read_task path and the refactored split_table_by_group utility. The null-validation order change is intentional and consistent with the new documented contract. The only genuine finding is a dead-code metadata fallback that can never fire. No data-correctness or runtime-crash risks were identified.

nemo_curator/core/utils.py deserves a second look for callers of split_table_by_group with no batch limits that may now surface null-sample_id errors they previously swallowed silently.

Important Files Changed

Filename	Overview
nemo_curator/stages/interleaved/lance/reader.py	New InterleavedLanceReaderStage and InterleavedLanceReader; dead-code fallback on line 72 (output.metadata is never None); otherwise logic is clean.
nemo_curator/core/utils.py	split_table_by_group adds max_batch_rows and null-rejection; null check now runs before the no-limit early return, silently changing behavior for all callers that pass max_batch_bytes=None.
nemo_curator/stages/text/io/reader/lance.py	Clean extraction of _requested_blob_columns, _dataset_and_scanner_kwargs, and _metadata_for_task; read_task logic is unchanged, just reorganized.
nemo_curator/stages/interleaved/utils/schema.py	New align_interleaved_table helper correctly unifies the two call sites in base.py and the new lance reader.
nemo_curator/stages/interleaved/init.py	Lazy getattr import avoids loading lance at package import time; pattern is consistent with io/init.py.
tests/stages/interleaved/test_interleaved_lance_reader.py	Good coverage of required fields, decomposition, null rejection, splits, and version pinning; missing coverage for the streaming blob-column rejection mentioned in the PR description.
tests/stages/interleaved/test_multimodal_core.py	All split_table_by_group tests correctly updated to the renamed function with keyword-only arguments; new tests for max_batch_rows and null rejection added.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant ILR as InterleavedLanceReader
    participant LPS as LancePartitioningStage
    participant ILRS as InterleavedLanceReaderStage
    participant LRS as LanceReaderStage.read_task
    participant STG as split_table_by_group

    Caller->>ILR: decompose()
    ILR-->>Caller: [LancePartitioningStage, InterleavedLanceReaderStage]

    Caller->>LPS: process(EmptyTask)
    LPS->>LPS: open dataset, pin version
    LPS-->>Caller: list[LanceReadTask] (fragment_ids + version)

    Caller->>ILRS: process(LanceReadTask)
    ILRS->>LRS: read_task(task, read_kwargs, fields)
    LRS->>LRS: _dataset_and_scanner_kwargs()
    LRS->>LRS: scanner.to_table() + blob materialization
    LRS-->>ILRS: ReaderOutput(table, metadata)

    ILRS->>ILRS: align_interleaved_table(table)
    ILRS->>STG: split_table_by_group(table, sample_id, max_batch_bytes, max_batch_rows)
    STG->>STG: validate no null sample_ids
    STG->>STG: compute group boundaries
    STG-->>ILRS: list[pa.Table] (splits)

    ILRS->>ILRS: build InterleavedBatch per split
    ILRS->>ILRS: batch.validate() per split
    ILRS-->>Caller: "InterleavedBatch | list[InterleavedBatch]"

Loading

_{Reviews (6): Last reviewed commit: "Refactor Interleaved Lance Reader: Repla..." | Re-trigger Greptile}

[suiyoubi]

/ok-to-test

[VibhuJawa]

codex: Consolidation opportunities with the existing Lance and interleaved reader infrastructure.

[VibhuJawa]

codex: This overlaps with split_table_by_group_max_bytes, which is already used by the Parquet and WebDataset interleaved readers. Can we consolidate these into one shared interleaved batching utility that accepts tables or record batches, preserves row order, and supports byte, row, and output limits?

[suiyoubi]

makes sense, but it appears to me that existing utility sorts by group column; and this Lance path preserves row order and supports streaming carry-over, row limits, and output limits.
I think it needs a careful shared utility refactor with tests across Parquet/WebDataset/Lance. @VibhuJawa do you think this is necessary to merge ?

[VibhuJawa]

Actually, I am still confused about the value of the streaming path. _tables_from_record_batches() collects every output chunk in a list, and process() returns them only after scanning finishes. Therefore, it does not provide bounded-memory streaming unless max_output_rows stops the scan early.

What concrete workload requires this path instead of reading the task normally and splitting the resulting table afterward?

For example:

output = self.read_task(task, dict(self.read_kwargs or {}), self.fields)
table = align_interleaved_table(output.data)

splits = split_table_by_group(
    table,
    "sample_id",
    max_batch_bytes=self.max_batch_bytes,
    max_batch_rows=self.max_batch_rows,
)

Like, what are we trying to achieve by adding this code ?

[suiyoubi]

The original motivation for this path was to avoid read_task() calling scanner.to_table() and constructing one complete fragment-size Arrow table before any document-aware splitting occurs.

But based on my experiments with mint1t, I think this might not be necessary, so I'll refactor to use your recommendation 👍

[VibhuJawa]

codex: Optional dependency import boundary.

[VibhuJawa]

Mostly looks good, added a comment on compatibility shim.

[VibhuJawa]

This function split_table_by_group looks good, dont think we should add the compatibility shim , so can get rid of split_table_by_group_max_bytes(

[VibhuJawa]

codex: Can we test the owning stage directly and verify that the error identifies every missing field? Please rename this to `test_interleaved_lance_reader_stage_reports_missing_required_fields.

Something like:

def test_interleaved_lance_reader_stage_reports_missing_required_fields() -> None:
    fields = ["sample_id"]
    missing = sorted(InterleavedBatch.REQUIRED_COLUMNS - set(fields))

    with pytest.raises(ValueError) as exc_info:
        InterleavedLanceReaderStage(fields=fields)

    assert str(exc_info.value) == (
        f"InterleavedLanceReaderStage fields are missing required columns: {missing}"
    )

[VibhuJawa]

codex: Since task._stage_perf is empty, this would still pass if the reader discarded its contents and created unrelated empty lists. Can we seed the task with a performance entry and assert that every output preserves it while owning a distinct list?

[VibhuJawa]

codex: The requested version is also the current dataset version, so this test passes even if read_kwargs["version"] is ignored.

Can we write distinguishable old and new dataset versions, request the old version through InterleavedLanceReader, and assert both task.version == old_version and that only the old rows are returned? This would also match the version-history coverage in test_lance_reader_uses_task_version_over_read_kwargs.

[suiyoubi]

/ok-to-test

[VibhuJawa]

LGTM, thanks for the work on this