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

nemo_retriever/src/nemo_retriever/common/params/__init__.py

@@ -33,6 +33,7 @@
 from nemo_retriever.common.params.models import TextChunkParams
 from nemo_retriever.common.params.models import TextGenerationParams
 from nemo_retriever.common.params.models import MetaJoinKey
+from nemo_retriever.common.params.models import VdbSinkParams
 from nemo_retriever.common.params.models import VdbUploadParams
 from nemo_retriever.common.params.models import VideoFrameParams
 from nemo_retriever.common.params.models import VideoFrameTextDedupParams
@@ -74,6 +75,7 @@
     "TextChunkParams",
     "TextGenerationParams",
     "MetaJoinKey",
+    "VdbSinkParams",
     "VdbUploadParams",
     "VideoFrameParams",
     "VideoFrameTextDedupParams",

nemo_retriever/src/nemo_retriever/common/params/models.py

@@ -692,6 +692,39 @@ def _warn_page_granularity_overrides(self) -> "EmbedParams":
 MetaJoinKey = Literal["auto", "source_id", "source_name"]
 
 
+class VdbSinkParams(_ParamsModel):
+    """Bounded terminal-sink policy for Ray Data VDB ingestion."""
+
+    max_batch_bytes: int = Field(
+        default=256 << 20,
+        gt=0,
+        description="Maximum retained Arrow batch size in bytes for the coordinated VDB sink.",
+    )
+    prefetch_batches: int = Field(
+        default=1,
+        ge=0,
+        description="Number of upstream batches buffered while the VDB sink writes.",
+    )
+    optimize: bool = Field(
+        default=False,
+        description="Run the backend optimization lifecycle after the bounded write completes.",
+    )
+    operation_id: Optional[str] = Field(
+        default=None,
+        description="Stable retry identity used to recover or reject repeated sink operations.",
+    )
+
+    @field_validator("operation_id")
+    @classmethod
+    def _validate_operation_id(cls, value: Optional[str]) -> Optional[str]:
+        if value is None:
+            return None
+        normalized = value.strip()
+        if not normalized:
+            raise ValueError("operation_id must be non-empty when provided")
+        return normalized
+
+
 class VdbUploadParams(_ParamsModel):
     """Post-graph vector DB upload configuration.
 
@@ -701,6 +734,7 @@ class VdbUploadParams(_ParamsModel):
 
     vdb_op: str = "lancedb"
     vdb_kwargs: dict[str, Any] = Field(default_factory=dict)
+    sink: VdbSinkParams = Field(default_factory=VdbSinkParams)
     meta_dataframe: Optional[Any] = None
     """Path to csv/json/parquet or an in-memory :class:`pandas.DataFrame`."""
     meta_source_field: Optional[str] = None

nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py

@@ -4,10 +4,10 @@
 
 import json
 import logging
+import math
 import os
 import threading
 import time
-
 from collections.abc import Iterable, Sequence
 from types import SimpleNamespace
 from typing import Any, Final, FrozenSet
@@ -27,9 +27,9 @@
     DocumentPage,
 )
 from nemo_retriever.common.vdb.adt_vdb import (
+    VDB,
     CollectionWriteContext,
     CollectionWriteResult,
-    VDB,
 )
 from nemo_retriever.common.vdb.lancedb_capabilities import (
     inspect_lancedb_table_object,
@@ -76,6 +76,43 @@ def _normalize_on_bad_vectors(value: str) -> str:
     return normalized
 
 
+def _stabilize_fill_vectors(
+    rows: list[dict[str, Any]],
+    *,
+    vector_dim: int,
+    fill_value: float,
+) -> list[dict[str, Any]]:
+    """Make ``fill`` independent of the installed LanceDB release.
+
+    LanceDB 0.34 replaces the complete vector when its width is wrong or any
+    element is NaN. Newer releases preserve valid elements and fill only the
+    invalid positions. NeMo Retriever supports both installation paths, so
+    normalize the historical public behavior before handing rows to LanceDB.
+    Values that LanceDB cannot coerce remain untouched so its normal validation
+    and error reporting still apply.
+    """
+
+    replacement = [float(fill_value)] * int(vector_dim)
+    stabilized: list[dict[str, Any]] = []
+    for row in rows:
+        vector = row.get("vector")
+        try:
+            wrong_dim = len(vector) != vector_dim
+        except TypeError:
+            wrong_dim = True
+
+        has_nan = False
+        if not wrong_dim:
+            try:
+                has_nan = any(value is not None and math.isnan(float(value)) for value in vector)
+            except (TypeError, ValueError):
+                stabilized.append(row)
+                continue
+
+        stabilized.append({**row, "vector": list(replacement)} if wrong_dim or has_nan else row)
+    return stabilized
+
+
 def _json_str(value) -> str:
     """
     Convert Python objects (dict/list/etc.) to a compact JSON string.
@@ -897,6 +934,13 @@ def create_index(self, records=None, table_name: str = "nv-ingest", **kwargs):
                     record_batches, expected_dim=vector_dim if enforce_dim else None
                 )
 
+            if self.on_bad_vectors == "fill":
+                results = _stabilize_fill_vectors(
+                    results,
+                    vector_dim=vector_dim,
+                    fill_value=self.fill_value,
+                )
+
             if self._service_table_schema:
                 results = _to_service_lancedb_rows(results)
                 schema = _with_retrieval_mode_metadata(
@@ -1003,13 +1047,20 @@ def write_to_index(
         hybrid = hybrid if hybrid is not None else self.hybrid
         sparse = sparse if sparse is not None else self.sparse
         fts_language = fts_language or self.fts_language
+        phase_timings = kwargs.pop("_phase_timings", None)
+        if isinstance(phase_timings, dict):
+            phase_timings.setdefault("vector_index", 0.0)
+            phase_timings.setdefault("fts_index", 0.0)
 
         if sparse:
             fts_index_start = time.perf_counter()
             sparse_rows = int(table.count_rows())
             table.create_fts_index("text", language=fts_language, replace=True)
             wait_for_column_index(table, "text", covered_rows=sparse_rows)
-            _record_timing("lancedb.fts_index_ready", time.perf_counter() - fts_index_start)
+            fts_duration = time.perf_counter() - fts_index_start
+            _record_timing("lancedb.fts_index_ready", fts_duration)
+            if isinstance(phase_timings, dict):
+                phase_timings["fts_index"] = fts_duration
             return
 
         num_rows = int(table.count_rows())
@@ -1054,13 +1105,19 @@ def write_to_index(
                 replace=True,
             )
             wait_for_column_index(table, "vector", covered_rows=num_rows)
-            _record_timing("lancedb.vector_index_ready", time.perf_counter() - vector_index_start)
+        vector_duration = time.perf_counter() - vector_index_start
+        _record_timing("lancedb.vector_index_ready", vector_duration)
+        if isinstance(phase_timings, dict):
+            phase_timings["vector_index"] = vector_duration
 
         if hybrid:
             fts_index_start = time.perf_counter()
             table.create_fts_index("text", language=fts_language, replace=True)
             wait_for_column_index(table, "text", covered_rows=num_rows)
-            _record_timing("lancedb.fts_index_ready", time.perf_counter() - fts_index_start)
+            fts_duration = time.perf_counter() - fts_index_start
+            _record_timing("lancedb.fts_index_ready", fts_duration)
+            if isinstance(phase_timings, dict):
+                phase_timings["fts_index"] = fts_duration
 
     def run(self, records):
         """Commit rows, then bring the table indexes up to date.
@@ -1209,6 +1266,13 @@ def put(
                 "put() only updates existing rows and will not create tables."
             ) from exc
 
+        if self.on_bad_vectors == "fill":
+            rows = _stabilize_fill_vectors(
+                rows,
+                vector_dim=_schema_vector_dim(_table_schema(table)),
+                fill_value=self.fill_value,
+            )
+
         input_ids = [r[key] for r in rows]
         unique_input_ids = list(dict.fromkeys(input_ids))
 
@@ -1272,6 +1336,9 @@ def sparse_retrieval(self, query_texts: Iterable[str], **kwargs: Any) -> list[li
             where_clause = str(where_clause).strip() or None
 
         table = lancedb.connect(uri=table_path).open_table(table_name)
+        from nemo_retriever.common.vdb.sink import assert_lancedb_table_ready
+
+        assert_lancedb_table_ready(table)
 
         search_results = []
         for query_text in query_texts:
@@ -1350,6 +1417,9 @@ def retrieval(self, vectors: Iterable[Sequence[float]], **kwargs: Any) -> list[l
             where_clause = str(where_clause).strip() or None
 
         table = lancedb.connect(uri=table_path).open_table(table_name)
+        from nemo_retriever.common.vdb.sink import assert_lancedb_table_ready
+
+        assert_lancedb_table_ready(table)
 
         if hybrid:
             vectors_for_search = list(vectors)