feat(vdb): bound batch LanceDB ingestion with Arrow streaming - #2567
feat(vdb): bound batch LanceDB ingestion with Arrow streaming#2567jioffe502 wants to merge 10 commits into
Conversation
Signed-off-by: jioffe502 <jioffe@nvidia.com>
Signed-off-by: jioffe502 <jioffe@nvidia.com>
|
Refreshed this branch onto current Post-merge validation:
The 12-dataset measurements in the PR body remain the source-frozen Architectural scope is now explicit: this PR is LanceDB-first. The Arrow schema, |
Signed-off-by: jioffe502 <jioffe@nvidia.com>
Greptile SummaryThe PR adds bounded Arrow streaming for LanceDB ingestion while preserving the terminal pandas result and introducing durable operation markers for interrupted writes.
|
| 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
Reviews (5): Last reviewed commit: "fix(vdb): stabilize fill policy across L..." | Re-trigger Greptile
Signed-off-by: jioffe502 <jioffe@nvidia.com>
3a7a33e to
89ba429
Compare
89ba429 to
29f9c6c
Compare
Signed-off-by: Jacob Ioffe <jioffe@nvidia.com>
Signed-off-by: Jacob Ioffe <jioffe@nvidia.com>
|
Follow-up on the The sink now owns and normalizes the historical/backward-compatible full-vector replacement policy before both create and append/put operations; the default |
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:
repartition(num_blocks=1);RecordBatchReaderto one LanceDB data mutation;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
PutVdbOperatorkeep their existing behavior.Measured result
We ran the existing implementation and this PR through the real
RayDataExecutoron the same source-frozen, CPU-only set of 12 nightly datasets: 197,796 stored rows per arm and 24 completed/validated runs.Correctness checks showed:
Validation
Local validation after merging current
main: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.