Skip to content

feat(vdb): bound batch LanceDB ingestion with Arrow streaming - #2567

Open
jioffe502 wants to merge 10 commits into
NVIDIA:mainfrom
jioffe502:jioffe502/bounded-vdb-sink
Open

feat(vdb): bound batch LanceDB ingestion with Arrow streaming#2567
jioffe502 wants to merge 10 commits into
NVIDIA:mainfrom
jioffe502:jioffe502/bounded-vdb-sink

Conversation

@jioffe502

@jioffe502 jioffe502 commented Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Why

Batch VDB ingestion currently repartitions the complete embedded Ray dataset into one block before LanceDB starts writing. This creates a corpus-sized handoff and leaves the VDB write and index work until the end of the pipeline.

What this changes

This PR replaces that handoff with one coordinated LanceDB stream:

  • reads native Ray batches without repartition(num_blocks=1);
  • projects and packs canonical Arrow batches under a byte budget;
  • sends one lazy RecordBatchReader to one LanceDB data mutation;
  • validates the committed rows, schema, and index coverage before reporting success;
  • preserves the existing pandas ingest result without running upstream work twice;
  • records incomplete operations so retries fail safely instead of silently duplicating append data.

This is not a per-Ray-block append design. Ray workers do not independently commit, overwrite, or build indices.

The implementation is LanceDB-specific, but the Ray executor is not: it asks the ingest operator whether bounded terminal streaming is supported. Custom VDB implementations and PutVdbOperator keep their existing behavior.

Measured result

We ran the existing implementation and this PR through the real RayDataExecutor on the same source-frozen, CPU-only set of 12 nightly datasets: 197,796 stored rows per arm and 24 completed/validated runs.

Measured sink-slice result Existing global handoff This PR Change
Aggregate graph wall time 130.10 s 99.13 s -23.8%
Maximum sampled process-tree RSS 17.92 GiB 7.74 GiB -56.8%
Maximum sampled Ray object-store use 1.72 GiB 0.93 GiB -45.5%

Correctness checks showed:

  • exact stored row-count, multiset-digest, schema-fingerprint, and dense-index-coverage parity;
  • one table data mutation, one fragment, and one data file per run;
  • no candidate-side repartition or materialization.

Validation

Local validation after merging current main:

  • 53 real-Lance sink, recovery, concurrent-write, and index tests passed;
  • 164 operator and Ray executor tests passed;
  • the real Ray 2.56.1 multi-block integration passed on CPU;
  • Black and Flake8 passed.

Scope of the measurements

These are single-repetition sink-slice measurements, not end-to-end extraction/embedding results or statistically qualified performance claims. The retained inputs contain canonical VDB rows rather than the original wide graph rows.

The public pandas result remains corpus-sized. The configured byte limit bounds canonical Arrow batches owned by the sink; it does not claim a corpus-independent total process RSS bound. The performance matrix covers dense overwrite and indexing.

Detailed design and evidence: bounded Ray Data LanceDB sink research.

Relates to #2476.

Signed-off-by: jioffe502 <jioffe@nvidia.com>
Signed-off-by: jioffe502 <jioffe@nvidia.com>
@jioffe502

Copy link
Copy Markdown
Collaborator Author

Refreshed this branch onto current upstream/main (bc9f2ce5, including #2566) and resolved the LanceDB index-wait overlap without dropping either side's behavior.

Post-merge validation:

  • 53 real-Lance sink, recovery, concurrent-write, and IVF tests passed;
  • 163 operator and Ray executor tests passed;
  • the CPU-only real Ray 2.56.1 multi-block integration passed;
  • merge commit hooks passed (Black and Flake8).

The 12-dataset measurements in the PR body remain the source-frozen 41ab2753 control/candidate run; I did not relabel or silently rerun that evidence after the merge.

Architectural scope is now explicit: this PR is LanceDB-first. The Arrow schema, RecordBatchReader write, retry markers, version/index validation, and optimization lifecycle remain Lance-owned. Decoupling terminal streaming capability detection from concrete LanceDB knowledge is tracked separately in #2570 and is not being folded into this PR.

Signed-off-by: jioffe502 <jioffe@nvidia.com>
@jioffe502
jioffe502 marked this pull request as ready for review August 20, 2026 22:53
@jioffe502
jioffe502 requested review from a team as code owners August 20, 2026 22:53
@jioffe502
jioffe502 requested a review from jdye64 August 20, 2026 22:53
@greptile-apps

greptile-apps Bot commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR adds bounded Arrow streaming for LanceDB ingestion while preserving the terminal pandas result and introducing durable operation markers for interrupted writes.

  • Adds configurable batch-byte, prefetch, optimization, and retry-identity policies.
  • Coordinates one streamed LanceDB mutation followed by validation and index finalization.
  • Updates Ray execution to consume the bounded sink without globally repartitioning the upstream dataset.
  • Adds focused sink lifecycle, failure, parameter, graph, and Ray integration tests.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains from the prior review thread.

The previously reported configuration-schema issue is resolved by descriptive metadata on every new VDB sink policy field, and no other eligible blocking failure remains.

Important Files Changed

Filename Overview
nemo_retriever/src/nemo_retriever/common/params/models.py Adds a validated, descriptively documented Pydantic policy for bounded VDB ingestion; the previously reported missing field descriptions are fixed.
nemo_retriever/src/nemo_retriever/common/vdb/sink.py Implements canonical Arrow batching, one-mutation LanceDB ingestion, validation, index finalization, and write reporting.
nemo_retriever/src/nemo_retriever/common/vdb/sink_operation.py Adds durable pending, data, and success markers for retry classification and incomplete-operation visibility.
nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py Integrates vector normalization, index-phase timing, and readiness checks with the coordinated sink lifecycle.
nemo_retriever/src/nemo_retriever/graph/executor.py Adds the Ray executor path that streams upstream batches into a bounded sink and preserves downstream and terminal results.
nemo_retriever/src/nemo_retriever/operators/vdb.py Exposes bounded-sink capability and delegates streamed batches to the LanceDB coordinator.
nemo_retriever/tests/test_bounded_vdb_sink.py Covers bounded batching and successful LanceDB sink behavior.
nemo_retriever/tests/test_bounded_vdb_sink_failures.py Exercises interrupted writes, retry classification, and finalization failures.
nemo_retriever/tests/test_ray_vdb_sink_integration.py Verifies real multi-block Ray streaming, one upstream execution, and post-sink ordering.

Sequence Diagram

sequenceDiagram
    participant Ray as Ray Data executor
    participant Sink as Bounded VDB sink
    participant Lance as LanceDB
    participant Index as Index finalization
    Ray->>Sink: Stream native batches
    Sink->>Sink: Project canonical rows
    Sink->>Sink: Pack byte-bounded Arrow batches
    Sink->>Lance: One RecordBatchReader mutation
    Lance-->>Sink: Committed data version
    Sink->>Sink: Validate rows and schema
    Sink->>Index: Build and validate requested indexes
    Index-->>Sink: Coverage confirmed
    Sink-->>Ray: Write report and retained result
Loading

Reviews (5): Last reviewed commit: "fix(vdb): stabilize fill policy across L..." | Re-trigger Greptile

Comment thread nemo_retriever/src/nemo_retriever/common/params/models.py Outdated
@jioffe502
jioffe502 force-pushed the jioffe502/bounded-vdb-sink branch from 3a7a33e to 89ba429 Compare August 24, 2026 13:51
@jioffe502
jioffe502 force-pushed the jioffe502/bounded-vdb-sink branch from 89ba429 to 29f9c6c Compare August 24, 2026 13:53
@jioffe502

Copy link
Copy Markdown
Collaborator Author

Follow-up on the fill compatibility regression: the repository lock (LanceDB 0.34.0) replaces the entire vector when any element is invalid, while the newer LanceDB exercised by current container CI preserves otherwise-valid elements. That made the public behavior dependency-version-sensitive.

The sink now owns and normalizes the historical/backward-compatible full-vector replacement policy before both create and append/put operations; the default drop policy is unchanged. The focused real-Lance 0.34 tests pass locally, and the complete current CI matrix—including Docker against its resolved dependency set—is green at fc8166ef. This keeps the application contract deterministic rather than branching on the installed LanceDB version.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant