Add Lance reader stage

[VibhuJawa]

Split from #2106. This is PR 1 of 3 in the Lance IO stack.

What changed:

• Generalizes text IO readers with a task-oriented BaseReader plus BaseFileReader for existing file-path readers.
• Adds Lance fragment partitioning and LanceReaderStage on top of BaseReader.
• Emits Lance row/fragment metadata for downstream annotation writers and preserves Lance schema metadata for later writes.
• Adds a reader-only lance extra using pylance.

Stack:

1. LanceReader (this PR)
2. LanceWriter on top of this branch
3. LanceAnnotationWriter on top of LanceWriter

Validation:

• uv run --extra lance --group test pytest -q tests/stages/text/io/reader/test_lance.py tests/stages/text/io/reader/test_jsonl.py tests/stages/text/io/reader/test_parquet.py -k "not IdGenerator and not id_generation"
• uv run ruff check nemo_curator/stages/text/io/reader/base.py nemo_curator/stages/text/io/reader/jsonl.py nemo_curator/stages/text/io/reader/parquet.py nemo_curator/stages/text/io/reader/lance.py nemo_curator/utils/lance.py tests/stages/text/io/reader/test_lance.py

Note: the full JSONL ID-generator reader test path was not used for this focused validation because the local untracked outputs/ benchmark artifacts make Ray runtime-env packaging exceed 512 MiB.

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

[VibhuJawa]

Please fix these things

[greptile-apps]

Greptile Summary

This PR introduces a Lance reader stack by generalizing the existing text IO reader hierarchy (BaseReader → BaseFileReader for file-path readers; BaseReader → LanceReaderStage for fragment-based reads) and adding the new LancePartitioningStage / LanceReaderStage / LanceReader composite. Fragment ID partitioning, Blob v2 column materialization, Lance schema metadata preservation, and version pinning are all well-implemented. Test coverage is thorough across the key scenarios.

• LancePartitioningStage.process() opens the dataset once, resolves the version, and emits sorted fragment-group LanceReadTask objects; LanceReaderStage consumes those tasks, materialises blob columns by row address, and injects __lance_rowid / __lance_rowaddr / __lance_fragid columns when requested.
• pylance is added as an optional extra (lance = ["pylance>=7"]) with lazy imports in reader/__init__.py so non-Lance users pay no import cost.
• BaseFileReader is introduced as a thin adapter preserving backward compatibility for the existing JSONL and Parquet reader stages.

Confidence Score: 5/5

The change is safe to merge; all new code is additive and non-breaking, the existing JSONL/Parquet readers are only re-parented to BaseFileReader, and the two findings are both minor edge cases with no impact on the default pipeline path.

The Lance reader stack is well-structured and the core read path—fragment partitioning, blob materialization, metadata injection, version pinning—is correct. The only findings are a hash ordering sensitivity in LanceReadTask.get_deterministic_id (only affects direct callers who pass unsorted fragment lists outside the normal pipeline) and a silent no-op for the private _generate_ids/_assign_ids flags inherited from BaseReader (these flags are not exposed in the LanceReaderStage public API). Neither affects the default pipeline behaviour.

No files require special attention; nemo_curator/tasks/lance.py has the hash ordering nit and nemo_curator/stages/text/io/reader/lance.py has the inherited-flag silent-skip, but both are low-risk.

Important Files Changed

Filename	Overview
nemo_curator/stages/text/io/reader/base.py	Generalizes BaseReader to accept both FileGroupTask and LanceReadTask; adds BaseFileReader adapter and allow_empty flag; clean refactor with no regressions
nemo_curator/stages/text/io/reader/lance.py	Adds LancePartitioningStage, LanceReaderStage, and composite LanceReader; logic for blob v2 materialization and metadata injection is solid, but _generate_ids/_assign_ids inherited from BaseReader are silently no-ops for pa.Table data
nemo_curator/tasks/lance.py	Defines LanceReadTask with fragment-list data; get_deterministic_id is order-sensitive for fragment IDs which could produce non-deterministic hashes if callers don't pre-sort
nemo_curator/utils/lance.py	Vectorized fragment-ID extraction via pyarrow.compute.shift_right; blob column materialization still uses to_pylist() for row addresses (noted in prior review)
tests/stages/text/io/reader/test_lance.py	Good coverage of partitioning, blob materialization, stable row IDs, column filtering, and version pinning; one test passes the EmptyTask class instead of an instance (flagged in prior review)
pyproject.toml	Adds lance optional extra (pylance>=7) and includes it in the all aggregate; CPU and GPU test matrices updated accordingly

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Pipeline
    participant LanceReader
    participant LancePartitioningStage
    participant LanceReaderStage
    participant LanceDataset

    Pipeline->>LanceReader: decompose()
    LanceReader-->>Pipeline: [LancePartitioningStage, LanceReaderStage]

    Pipeline->>LancePartitioningStage: process(EmptyTask)
    LancePartitioningStage->>LanceDataset: "lance.dataset(path, **dataset_kwargs)"
    LanceDataset-->>LancePartitioningStage: dataset (version resolved)
    LancePartitioningStage->>LanceDataset: get_fragments()
    LanceDataset-->>LancePartitioningStage: available fragment IDs
    LancePartitioningStage-->>Pipeline: "[LanceReadTask(fragment_ids=[0,1]), ...]"

    loop for each LanceReadTask
        Pipeline->>LanceReaderStage: process(LanceReadTask)
        LanceReaderStage->>LanceDataset: "lance.dataset(task.path, version=task.version)"
        LanceReaderStage->>LanceDataset: "scanner(fragments=..., columns=..., filter=...).to_table()"
        LanceDataset-->>LanceReaderStage: pa.Table (with _rowid, _rowaddr)
        alt has Blob v2 columns
            LanceReaderStage->>LanceDataset: "read_blobs(field_name, addresses=row_addresses)"
            LanceDataset-->>LanceReaderStage: binary payloads
        end
        LanceReaderStage->>LanceReaderStage: add_lance_metadata_columns
        LanceReaderStage-->>Pipeline: "DocumentBatch(data=pa.Table, _metadata={lance schema, version, fragment_ids})"
    end

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Pipeline
    participant LanceReader
    participant LancePartitioningStage
    participant LanceReaderStage
    participant LanceDataset

    Pipeline->>LanceReader: decompose()
    LanceReader-->>Pipeline: [LancePartitioningStage, LanceReaderStage]

    Pipeline->>LancePartitioningStage: process(EmptyTask)
    LancePartitioningStage->>LanceDataset: "lance.dataset(path, **dataset_kwargs)"
    LanceDataset-->>LancePartitioningStage: dataset (version resolved)
    LancePartitioningStage->>LanceDataset: get_fragments()
    LanceDataset-->>LancePartitioningStage: available fragment IDs
    LancePartitioningStage-->>Pipeline: "[LanceReadTask(fragment_ids=[0,1]), ...]"

    loop for each LanceReadTask
        Pipeline->>LanceReaderStage: process(LanceReadTask)
        LanceReaderStage->>LanceDataset: "lance.dataset(task.path, version=task.version)"
        LanceReaderStage->>LanceDataset: "scanner(fragments=..., columns=..., filter=...).to_table()"
        LanceDataset-->>LanceReaderStage: pa.Table (with _rowid, _rowaddr)
        alt has Blob v2 columns
            LanceReaderStage->>LanceDataset: "read_blobs(field_name, addresses=row_addresses)"
            LanceDataset-->>LanceReaderStage: binary payloads
        end
        LanceReaderStage->>LanceReaderStage: add_lance_metadata_columns
        LanceReaderStage-->>Pipeline: "DocumentBatch(data=pa.Table, _metadata={lance schema, version, fragment_ids})"
    end

Loading

_{Reviews (17): Last reviewed commit: "Merge upstream main into feat/lance-read..." | Re-trigger Greptile}

[VibhuJawa]

/ok to test

[copy-pr-bot]

/ok to test

@VibhuJawa, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/

[VibhuJawa]

/ok to test 8a2eb42

[VibhuJawa]

/ok to test a8285af