Aot/interleaved lance materializer

[suiyoubi]

Description

Usage

# Add snippet demonstrating usage

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds a Lance image materialization stage (LanceRowIdImageMaterializationStage) that reads binary image data from a pinned Lance table using either stable row IDs or fragment/row-offset addresses, writing the fetched bytes into interleaved batch columns without any URL lookup. It also extracts the shared LanceTableConfig dataclass into its own module and fixes a relative-import bug in the interleaved.io package.

• LanceTableConfig (config.py): frozen dataclass holding the Lance URI, optional pinned version, and storage options; used by both the new materializer and the existing reader.
• fetch.py: two fetcher strategies — _LanceRowIdFetcher (calls the private dataset._take_rows API, optional sort-for-sequential-IO) and _LanceRowAddressFetcher (per-fragment take with a snapshot of fragments taken at init); both support a configurable timeout via _submit_fetches with FIRST_COMPLETED polling and cancel-on-timeout.
• materialize.py: LanceRowIdImageMaterializationStage orchestrates batch preparation, combined multi-task fetch, per-task result slicing, column projection (fill_null/overwrite/error), presence tracking, retry on _LanceFetchTimeoutError, and optional fetcher recycling via fetcher_max_batches.

Confidence Score: 4/5

The new materialization stage and fetch utilities are well-structured with thorough validation, test coverage for all major code paths, and correct ordering logic verified through both unit and integration tests. The flagged items are non-blocking quality issues.

The core logic — multi-task address collection, combined fetch, per-task _slice_fetched_tables slicing, column projection, and presence tracking — is correct and backed by tests. The TYPE_CHECKING-only import of LanceTableConfig in the dataclass will silently break any caller that invokes typing.get_type_hints() on LanceRowIdImageMaterializationStage, and dead code remains after the retry loop. Neither prevents normal operation today, but the import issue is a latent trap for serialization or introspection-based workflows.

nemo_curator/stages/interleaved/lance/materialize.py deserves a second look for the TYPE_CHECKING-only LanceTableConfig import and the dead code at the end of _fetch_requested_images.

Important Files Changed

Filename	Overview
nemo_curator/stages/interleaved/lance/config.py	New frozen dataclass LanceTableConfig for pinned Lance table URI/version/storage options; includes non-empty URI validation in __post_init__.
nemo_curator/stages/interleaved/lance/fetch.py	New fetch utilities: timeout-aware _submit_fetches, ordered-restore helpers, and two fetcher strategies (_LanceRowIdFetcher via private _take_rows API, _LanceRowAddressFetcher via fragment snapshots).
nemo_curator/stages/interleaved/lance/materialize.py	New LanceRowIdImageMaterializationStage dataclass (2 address modes, fill_null/overwrite/error column policy, retry on timeout, batch recycling). LanceTableConfig is TYPE_CHECKING-only import; unreachable dead code in retry loop.
tests/stages/interleaved/test_interleaved_lance_materializer.py	Comprehensive unit and integration tests covering row-id mode, row-address mode, fill_null skipping, timeout retry, JSON source_ref parsing, and real local Lance dataset round-trips.
nemo_curator/stages/interleaved/lance/init.py	Exports LanceTableConfig and LanceRowIdImageMaterializationStage alongside existing reader exports.
nemo_curator/stages/interleaved/io/init.py	Switches InterleavedLanceReader lazy import to relative ..lance and fixes import_module to pass package=__name__ for correct relative resolution.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant Stage as LanceRowIdImageMaterializationStage
    participant Fetcher as _LanceRowIdFetcher / _LanceRowAddressFetcher
    participant Lance as Lance Dataset

    Caller->>Stage: process_batch(tasks)
    Stage->>Stage: _ensure_fetcher() → create fetcher (opens Lance dataset)
    Stage->>Stage: _prepare_task() × N (validate, extract addresses)
    Stage->>Stage: flatten requested_addresses across all tasks
    Stage->>Fetcher: fetch(requested_addresses)
    loop retry on _LanceFetchTimeoutError
        Fetcher->>Lance: _take_rows / fragment.take (via ThreadPoolExecutor)
        Lance-->>Fetcher: pa.Table chunks
        Fetcher->>Fetcher: _restore_fetched_original_order
        Fetcher-->>Stage: _RowIdFetchResult (tables in original order)
    end
    Stage->>Stage: _slice_fetched_tables per task
    loop for each prepared task
        Stage->>Stage: _apply_projection (scatter or replace columns)
        Stage->>Stage: _apply_presence (set presence flag)
    end
    Stage->>Stage: _log_metrics
    Stage->>Stage: _maybe_recycle_fetcher_after_success
    Stage-->>Caller: list[InterleavedBatch]

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 Caller
    participant Stage as LanceRowIdImageMaterializationStage
    participant Fetcher as _LanceRowIdFetcher / _LanceRowAddressFetcher
    participant Lance as Lance Dataset

    Caller->>Stage: process_batch(tasks)
    Stage->>Stage: _ensure_fetcher() → create fetcher (opens Lance dataset)
    Stage->>Stage: _prepare_task() × N (validate, extract addresses)
    Stage->>Stage: flatten requested_addresses across all tasks
    Stage->>Fetcher: fetch(requested_addresses)
    loop retry on _LanceFetchTimeoutError
        Fetcher->>Lance: _take_rows / fragment.take (via ThreadPoolExecutor)
        Lance-->>Fetcher: pa.Table chunks
        Fetcher->>Fetcher: _restore_fetched_original_order
        Fetcher-->>Stage: _RowIdFetchResult (tables in original order)
    end
    Stage->>Stage: _slice_fetched_tables per task
    loop for each prepared task
        Stage->>Stage: _apply_projection (scatter or replace columns)
        Stage->>Stage: _apply_presence (set presence flag)
    end
    Stage->>Stage: _log_metrics
    Stage->>Stage: _maybe_recycle_fetcher_after_success
    Stage-->>Caller: list[InterleavedBatch]

Loading

Comments Outside Diff (1)

1. nemo_curator/stages/interleaved/lance/materialize.py, line 978-988 (link)

   Fetcher recycled even on no-I/O batches

   _maybe_recycle_fetcher_after_success is called unconditionally at the end of every _process_tasks call, including batches where requested_addresses is empty (all rows were already populated or had null addresses). Each such batch still increments _fetcher_batches and counts toward fetcher_max_batches. When many consecutive batches require no fetching — e.g., a dataset that has already been fully materialized — the fetcher is repeatedly torn down and recreated, incurring unnecessary Lance dataset() open overhead. Consider guarding the recycle call on whether any I/O actually occurred.

_{Reviews (1): Last reviewed commit: "Refactor Lance Image Fetching and Remove..." | Re-trigger Greptile}

[greptile-apps]

LanceTableConfig is only imported under TYPE_CHECKING, so it is absent from the module's namespace at runtime. The @dataclass decorator itself works fine with from __future__ import annotations because annotations are stored as strings. However, anything that calls typing.get_type_hints(LanceRowIdImageMaterializationStage) — common in serialization libraries, dependency-injection frameworks, and documentation generators — will raise NameError: name 'LanceTableConfig' is not defined because it tries to resolve the string annotation in the module globals. Moving the import outside the TYPE_CHECKING guard avoids this without any functional cost.

Suggested change

	if TYPE_CHECKING:
	from .config import LanceTableConfig
	from .config import LanceTableConfig
	if TYPE_CHECKING:
	pass

[greptile-apps]

The raise RuntimeError("unreachable Lance fetch retry state") is dead code. The for loop spans exactly max_attempts iterations; the last iteration always either returns successfully or hits attempt >= max_attempts and raises RuntimeError. The loop body can never exhaust without one of those two outcomes, so execution never reaches the line after the loop.

Suggested change

	msg = "unreachable Lance fetch retry state"
	raise RuntimeError(msg)
	# Unreachable: the loop always returns or raises on the final attempt.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!