Merge remote-tracking branch 'origin/main' into gatekeeper/fix-3153-1

docs/src/js/namespaces/rerankers/classes/RRFReranker.md

@@ -10,24 +10,6 @@ Reranks the results using the Reciprocal Rank Fusion (RRF) algorithm.
 
 ## Methods
 
-### outputSchema()
-
-```ts
-outputSchema(inputSchema): Promise<Schema<any>>
-```
-
-Declare the RRF output schema for vector-only query execution.
-
-#### Parameters
-
-* **inputSchema**: `Schema`&lt;`any`&gt;
-
-#### Returns
-
-`Promise`&lt;`Schema`&lt;`any`&gt;&gt;
-
-***
-
 ### rerankHybrid()
 
 ```ts

docs/src/js/namespaces/rerankers/interfaces/Reranker.md

@@ -8,27 +8,6 @@
 
 ## Methods
 
-### outputSchema()?
-
-```ts
-optional outputSchema(inputSchema): Promise<Schema<any>>
-```
-
-Declare the schema returned when reranking a vector-only query.
-
-This is required for vector-only reranking so query schema introspection
-and execution agree. Hybrid-only rerankers may omit it.
-
-#### Parameters
-
-* **inputSchema**: `Schema`&lt;`any`&gt;
-
-#### Returns
-
-`Promise`&lt;`Schema`&lt;`any`&gt;&gt;
-
-***
-
 ### rerankHybrid()
 
 ```ts

nodejs/__test__/rerankers.test.ts

@@ -79,22 +79,6 @@ describe("rerankers", function () {
     expect(result).toHaveLength(2);
   });
 
-  it("returns relevance scores when reranking a vector search", async function () {
-    const query = table
-      .vectorSearch([0.1, 0.1])
-      .limit(2)
-      .rerank(await RRFReranker.create());
-    const schema = await query.outputSchema();
-    const result = await query.toArray();
-
-    expect(schema.fields.map((field) => field.name)).toContain(
-      "_relevance_score",
-    );
-    expect(result).toHaveLength(2);
-    expect(result[0]._relevance_score).toBeCloseTo(1 / 60);
-    expect(result[1]._relevance_score).toBeCloseTo(1 / 61);
-  });
-
   it("does not keep process alive after rerank query", async function () {
     const script = `
 import * as lancedb from "./dist/index.js";

nodejs/lancedb/query.ts

@@ -5,11 +5,9 @@ import {
   Table as ArrowTable,
   type IntoVector,
   RecordBatch,
-  createEmptyTable,
   extractVectorBuffer,
   fromBufferToRecordBatch,
   fromRecordBatchToBuffer,
-  fromTableToBuffer,
   tableFromIPC,
 } from "./arrow";
 import { type IvfPqOptions } from "./indices";
@@ -746,30 +744,20 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
   }
 
   rerank(reranker: Reranker): VectorQuery {
-    super.doCall((inner) => {
-      const outputSchema = reranker.outputSchema?.bind(reranker);
-      inner.rerank(
-        async (args) => {
-          const vecResults = await fromBufferToRecordBatch(args.vecResults);
-          const ftsResults = await fromBufferToRecordBatch(args.ftsResults);
-          const result = await reranker.rerankHybrid(
-            args.query,
-            vecResults as RecordBatch,
-            ftsResults as RecordBatch,
-          );
+    super.doCall((inner) =>
+      inner.rerank(async (args) => {
+        const vecResults = await fromBufferToRecordBatch(args.vecResults);
+        const ftsResults = await fromBufferToRecordBatch(args.ftsResults);
+        const result = await reranker.rerankHybrid(
+          args.query,
+          vecResults as RecordBatch,
+          ftsResults as RecordBatch,
+        );
 
-          const buffer = fromRecordBatchToBuffer(result);
-          return buffer;
-        },
-        outputSchema
-          ? async (args) => {
-              const inputSchema = tableFromIPC(args.inputSchema).schema;
-              const result = await outputSchema(inputSchema);
-              return fromTableToBuffer(createEmptyTable(result));
-            }
-          : undefined,
-      );
-    });
+        const buffer = fromRecordBatchToBuffer(result);
+        return buffer;
+      }),
+    );
 
     return this;
   }

nodejs/lancedb/rerankers/index.ts

@@ -1,22 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
-import { RecordBatch, Schema } from "apache-arrow";
+import { RecordBatch } from "apache-arrow";
 
 export * from "./rrf";
 
-// Interface for a reranker. A reranker is used to rerank vector and hybrid
-// search results. For vector-only searches, query is empty and ftsResults is an
-// empty batch with the same schema as vecResults.
+// Interface for a reranker. A reranker is used to rerank the results from a
+// vector and FTS search. This is useful for combining the results from both
+// search methods.
 export interface Reranker {
-  /**
-   * Declare the schema returned when reranking a vector-only query.
-   *
-   * This is required for vector-only reranking so query schema introspection
-   * and execution agree. Hybrid-only rerankers may omit it.
-   */
-  outputSchema?(inputSchema: Schema): Promise<Schema>;
-
   rerankHybrid(
     query: string,
     vecResults: RecordBatch,

nodejs/lancedb/rerankers/rrf.ts

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
-import { Field, Float32, RecordBatch, Schema } from "apache-arrow";
+import { RecordBatch } from "apache-arrow";
 import { fromBufferToRecordBatch, fromRecordBatchToBuffer } from "../arrow";
 import { RrfReranker as NativeRRFReranker } from "../native";
 
@@ -24,17 +24,6 @@ export class RRFReranker {
     );
   }
 
-  /** Declare the RRF output schema for vector-only query execution. */
-  async outputSchema(inputSchema: Schema): Promise<Schema> {
-    return new Schema(
-      [
-        ...inputSchema.fields,
-        new Field("_relevance_score", new Float32(), false),
-      ],
-      inputSchema.metadata,
-    );
-  }
-
   async rerankHybrid(
     query: string,
     vecResults: RecordBatch,

nodejs/src/query.rs

@@ -6,8 +6,8 @@ use std::sync::Arc;
 use crate::error::NapiErrorExt;
 use crate::error::convert_error;
 use crate::iterator::RecordBatchIterator;
+use crate::rerankers::RerankHybridCallbackArgs;
 use crate::rerankers::Reranker;
-use crate::rerankers::{RerankHybridCallbackArgs, RerankOutputSchemaCallbackArgs};
 use crate::util::{parse_distance_type, schema_to_buffer};
 use arrow_array::{
     Array, Float16Array as ArrowFloat16Array, Float32Array as ArrowFloat32Array,
@@ -388,9 +388,8 @@ impl VectorQuery {
     pub fn rerank(
         &mut self,
         rerank_hybrid: Function<RerankHybridCallbackArgs, Promise<Buffer>>,
-        output_schema: Option<Function<RerankOutputSchemaCallbackArgs, Promise<Buffer>>>,
     ) -> napi::Result<()> {
-        let reranker = Reranker::new(rerank_hybrid, output_schema)?;
+        let reranker = Reranker::new(rerank_hybrid)?;
         self.inner = self.inner.clone().rerank(Arc::new(reranker));
         Ok(())
     }

nodejs/src/rerankers.rs

@@ -6,7 +6,7 @@ use async_trait::async_trait;
 use napi::{bindgen_prelude::*, threadsafe_function::ThreadsafeFunction};
 use napi_derive::napi;
 
-use lancedb::ipc::{batches_to_ipc_file, ipc_file_to_schema, schema_to_ipc_file};
+use lancedb::ipc::batches_to_ipc_file;
 use lancedb::rerankers::Reranker as LanceDBReranker;
 use lancedb::{error::Error, ipc::ipc_file_to_batches};
 
@@ -21,72 +21,28 @@ type RerankHybridFn = ThreadsafeFunction<
     true,
 >;
 
-type RerankOutputSchemaFn = ThreadsafeFunction<
-    RerankOutputSchemaCallbackArgs,
-    Promise<Buffer>,
-    RerankOutputSchemaCallbackArgs,
-    Status,
-    false,
-    true,
->;
-
 /// Reranker implementation that "wraps" a NodeJS Reranker implementation.
 /// This contains references to the callbacks that can be used to invoke the
 /// reranking methods on the NodeJS implementation and handles serializing the
 /// record batches to Arrow IPC buffers.
 pub struct Reranker {
     rerank_hybrid: RerankHybridFn,
-    output_schema: Option<RerankOutputSchemaFn>,
 }
 
 impl Reranker {
     pub fn new(
         rerank_hybrid: Function<RerankHybridCallbackArgs, Promise<Buffer>>,
-        output_schema: Option<Function<RerankOutputSchemaCallbackArgs, Promise<Buffer>>>,
     ) -> napi::Result<Self> {
         let rerank_hybrid = rerank_hybrid
             .build_threadsafe_function()
             .weak::<true>()
             .build()?;
-        let output_schema = output_schema
-            .map(|output_schema| {
-                output_schema
-                    .build_threadsafe_function()
-                    .weak::<true>()
-                    .build()
-            })
-            .transpose()?;
-        Ok(Self {
-            rerank_hybrid,
-            output_schema,
-        })
+        Ok(Self { rerank_hybrid })
     }
 }
 
 #[async_trait]
 impl lancedb::rerankers::Reranker for Reranker {
-    async fn output_schema(
-        &self,
-        input: &arrow_schema::SchemaRef,
-    ) -> lancedb::error::Result<arrow_schema::SchemaRef> {
-        let output_schema = self.output_schema.as_ref().ok_or(Error::NotSupported {
-            message: "vector rerankers must declare their output schema".to_string(),
-        })?;
-        let callback_args = RerankOutputSchemaCallbackArgs {
-            input_schema: Buffer::from(schema_to_ipc_file(input.as_ref())?),
-        };
-        let promised_buffer: Promise<Buffer> = output_schema
-            .call_async(callback_args)
-            .await
-            .map_err(|e| Error::Runtime {
-                message: format!("napi error status={}, reason={}", e.status, e.reason),
-            })?;
-        let buffer = promised_buffer.await.map_err(|e| Error::Runtime {
-            message: format!("napi error status={}, reason={}", e.status, e.reason),
-        })?;
-        ipc_file_to_schema(buffer.to_vec())
-    }
-
     async fn rerank_hybrid(
         &self,
         query: &str,
@@ -130,11 +86,6 @@ pub struct RerankHybridCallbackArgs {
     pub fts_results: Buffer,
 }
 
-#[napi(object)]
-pub struct RerankOutputSchemaCallbackArgs {
-    pub input_schema: Buffer,
-}
-
 fn buffer_to_record_batch(buffer: Buffer) -> Result<RecordBatch> {
     let mut reader = ipc_file_to_batches(buffer.to_vec()).default_error()?;
     reader

python/python/lancedb/_lancedb.pyi

@@ -269,6 +269,7 @@ class Table:
         mode: Literal["append", "overwrite"],
         progress: Optional[Any] = None,
         write_parallelism: Optional[int] = None,
+        on_nan_vectors: Optional[Literal["error", "keep"]] = None,
     ) -> AddResult: ...
     async def update(
         self, updates: Dict[str, str], where: Optional[str]

python/python/lancedb/db.py

@@ -333,7 +333,9 @@ class DBConnection(EnforceOverrides):
             schema that's specified.
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill".
+            One of "error", "drop", "fill", "null", or "keep". With "keep",
+            vectors containing NaNs are preserved, but vectors with the wrong
+            dimension still raise an error.
         fill_value: float
             The value to use when filling vectors. Only used if on_bad_vectors="fill".
         storage_options: dict, optional
@@ -1595,7 +1597,9 @@ class AsyncConnection(object):
             schema that's specified.
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill".
+            One of "error", "drop", "fill", "null", or "keep". With "keep",
+            vectors containing NaNs are preserved, but vectors with the wrong
+            dimension still raise an error.
         fill_value: float
             The value to use when filling vectors. Only used if on_bad_vectors="fill".
         storage_options: dict, optional

python/python/lancedb/merge.py

@@ -175,7 +175,9 @@ class LanceMergeInsertBuilder(object):
             can be anything you use for [`add`][lancedb.table.Table.add]
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill".
+            One of "error", "drop", "fill", "null", or "keep". With "keep",
+            vectors containing NaNs are preserved, but vectors with the wrong
+            dimension still raise an error.
         fill_value: float, default 0.
             The value to use when filling vectors. Only used if on_bad_vectors="fill".
         timeout: Optional[timedelta], default None

python/python/lancedb/remote/db.py

@@ -543,7 +543,9 @@ class RemoteDBConnection(DBConnection):
             to "exist_ok".
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill".
+            One of "error", "drop", "fill", "null", or "keep". With "keep",
+            vectors containing NaNs are preserved, but vectors with the wrong
+            dimension still raise an error.
         fill_value: float
             The value to use when filling vectors. Only used if on_bad_vectors="fill".
 

python/python/lancedb/remote/table.py

@@ -629,7 +629,9 @@ class RemoteTable(Table):
             "append" and "overwrite".
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill".
+            One of "error", "drop", "fill", "null", or "keep". With "keep",
+            vectors containing NaNs are preserved, but vectors with the wrong
+            dimension still raise an error.
         fill_value: float, default 0.
             The value to use when filling vectors. Only used if on_bad_vectors="fill".
         progress: bool, callable, or tqdm-like, optional

python/python/lancedb/table.py

@@ -350,8 +350,10 @@ def _sanitize_data(
         in the input table before casting.
     metadata : Optional[dict], default None
         The embedding metadata to add to the schema.
-    on_bad_vectors : Literal["error", "drop", "fill", "null"], default "error"
+    on_bad_vectors : Literal["error", "drop", "fill", "null", "keep"], default "error"
         What to do if any of the vectors are not the same size or contains NaNs.
+        With "keep", vectors containing NaNs are preserved, but vectors with the
+        wrong dimension still raise an error.
     fill_value : float, default 0.0
         The value to use when filling vectors. Only used if on_bad_vectors="fill".
         All entries in the vector will be set to this value.
@@ -1247,7 +1249,9 @@ class Table(ABC):
             "append" and "overwrite".
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill".
+            One of "error", "drop", "fill", "null", or "keep". With "keep",
+            vectors containing NaNs are preserved but are not indexed for vector
+            search; vectors with the wrong dimension still raise an error.
         fill_value: float, default 0.
             The value to use when filling vectors. Only used if on_bad_vectors="fill".
         progress: bool, callable, or tqdm-like, optional
@@ -3269,7 +3273,9 @@ class LanceTable(Table):
             "append" and "overwrite".
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill", "null".
+            One of "error", "drop", "fill", "null", or "keep". With "keep",
+            vectors containing NaNs are preserved, but vectors with the wrong
+            dimension still raise an error.
         fill_value: float, default 0.
             The value to use when filling vectors. Only used if on_bad_vectors="fill".
         progress: bool, callable, or tqdm-like, optional
@@ -3582,7 +3588,9 @@ class LanceTable(Table):
             data but will validate against any schema that's specified.
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill", "null".
+            One of "error", "drop", "fill", "null", or "keep". With "keep",
+            vectors containing NaNs are preserved, but vectors with the wrong
+            dimension still raise an error.
         fill_value: float, default 0.
             The value to use when filling vectors. Only used if on_bad_vectors="fill".
         embedding_functions: list of EmbeddingFunctionModel, default None
@@ -4018,7 +4026,7 @@ class LanceTable(Table):
 
 def _handle_bad_vectors(
     reader: pa.RecordBatchReader,
-    on_bad_vectors: Literal["error", "drop", "fill", "null"] = "error",
+    on_bad_vectors: OnBadVectorsType = "error",
     fill_value: float = 0.0,
     target_schema: Optional[pa.Schema] = None,
     metadata: Optional[dict] = None,
@@ -4192,7 +4200,9 @@ def _handle_bad_vector_column(
         The name of the vector column.
     on_bad_vectors: str, default "error"
         What to do if any of the vectors are not the same size or contains NaNs.
-        One of "error", "drop", "fill", "null".
+        One of "error", "drop", "fill", "null", or "keep". With "keep",
+        vectors containing NaNs are preserved, but vectors with the wrong dimension
+        still raise an error.
     fill_value: float, default 0.0
         The value to use when filling vectors. Only used if on_bad_vectors="fill".
     """
@@ -4257,7 +4267,8 @@ def _handle_bad_vector_column(
                     f"Vector column '{vector_column_name}' has NaNs. "
                     "Set on_bad_vectors='drop' to remove them, "
                     "set on_bad_vectors='fill' and fill_value=<value> to replace them, "
-                    "or set on_bad_vectors='null' to replace them with null."
+                    "set on_bad_vectors='null' to replace them with null, "
+                    "or set on_bad_vectors='keep' to preserve them."
                 )
         elif on_bad_vectors == "null":
             vec_arr = pc.if_else(
@@ -4274,6 +4285,16 @@ def _handle_bad_vector_column(
                     "`fill_value` must not be None if `on_bad_vectors` is 'fill'"
                 )
             vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value)
+        elif on_bad_vectors == "keep":
+            if pc.any(has_wrong_dim).as_py():
+                raise ValueError(
+                    f"Vector column '{vector_column_name}' has variable length "
+                    "vectors. on_bad_vectors='keep' only preserves vectors "
+                    "containing NaNs. Set on_bad_vectors='drop' to remove "
+                    "wrong-size vectors, set on_bad_vectors='fill' and "
+                    "fill_value=<value> to replace them, or set "
+                    "on_bad_vectors='null' to replace them with null."
+                )
         else:
             raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}")
 
@@ -5118,7 +5139,9 @@ class AsyncTable:
             "append" and "overwrite".
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill", "null".
+            One of "error", "drop", "fill", "null", or "keep". With "keep",
+            vectors containing NaNs are preserved but are not indexed for vector
+            search; vectors with the wrong dimension still raise an error.
         fill_value: float, default 0.
             The value to use when filling vectors. Only used if on_bad_vectors="fill".
         progress: callable or tqdm-like, optional
@@ -5166,6 +5189,7 @@ class AsyncTable:
                 mode or "append",
                 progress=progress,
                 write_parallelism=write_parallelism,
+                on_nan_vectors="keep" if on_bad_vectors == "keep" else None,
             )
         except RuntimeError as e:
             if "Cast error" in str(e):

python/python/lancedb/types.py

@@ -24,7 +24,7 @@ DistanceType = Literal["l2", "cosine", "dot"]
 DistanceTypeWithHamming = Literal["l2", "cosine", "dot", "hamming"]
 
 # Vector handling literals
-OnBadVectorsType = Literal["error", "drop", "fill", "null"]
+OnBadVectorsType = Literal["error", "drop", "fill", "null", "keep"]
 
 # Mode literals
 AddMode = Literal["append", "overwrite"]

python/python/tests/test_table.py

@@ -1767,6 +1767,40 @@ def test_add_with_nans(mem_db: DBConnection):
     assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
 
 
+def test_add_with_non_finite_values_keep(mem_db: DBConnection):
+    schema = pa.schema([pa.field("data", pa.list_(pa.float32(), 4))])
+    table = mem_db.create_table("test", schema=schema)
+    batch = pa.table(
+        {
+            "data": pa.array(
+                [[np.nan, np.inf, -np.inf, -0.0]],
+                type=schema.field("data").type,
+            )
+        },
+        schema=schema,
+    )
+
+    with pytest.raises(ValueError, match="NaN"):
+        table.add(batch)
+
+    table.add(batch, on_bad_vectors="keep")
+
+    values = table.to_arrow()["data"][0].as_py()
+    assert np.isnan(values[0])
+    assert np.isposinf(values[1])
+    assert np.isneginf(values[2])
+    assert values[3] == 0.0
+    assert np.signbit(values[3])
+
+
+def test_add_keep_rejects_wrong_dimension(mem_db: DBConnection):
+    schema = pa.schema([pa.field("vector", pa.list_(pa.float32(), 2))])
+    table = mem_db.create_table("test", schema=schema)
+
+    with pytest.raises((ValueError, RuntimeError), match="variable length"):
+        table.add([{"vector": [1.0]}], on_bad_vectors="keep")
+
+
 def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection):
     class Schema(LanceModel):
         text: str

python/python/tests/test_util.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
 
+import math
 import os
 import pathlib
 from typing import Optional
@@ -364,7 +365,7 @@ def test_fill_bad_vector_values_arrow_types(vector_type, vectors, expected):
     assert actual.to_pylist() == expected
 
 
-@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
+@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null", "keep"])
 def test_handle_bad_vectors_nan(on_bad_vectors):
     vector = pa.array([[1.0, float("nan")], [3.0, 4.0]])
     data = pa.table({"vector": vector})
@@ -379,8 +380,9 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
         assert output == (
             "ValueError: Vector column 'vector' has NaNs. Set "
             "on_bad_vectors='drop' to remove them, set on_bad_vectors='fill' "
-            "and fill_value=<value> to replace them, or set on_bad_vectors='null' "
-            "to replace them with null."
+            "and fill_value=<value> to replace them, set on_bad_vectors='null' "
+            "to replace them with null, or set on_bad_vectors='keep' to preserve "
+            "them."
         )
         return
     else:
@@ -396,10 +398,26 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
         expected = pa.array([[1.0, 42.0], [3.0, 4.0]])
     elif on_bad_vectors == "null":
         expected = pa.array([None, [3.0, 4.0]])
+    elif on_bad_vectors == "keep":
+        actual = output["vector"].to_pylist()
+        assert actual[0][0] == 1.0
+        assert math.isnan(actual[0][1])
+        assert actual[1] == [3.0, 4.0]
+        return
 
     assert output["vector"].combine_chunks() == expected
 
 
+def test_handle_bad_vectors_keep_rejects_wrong_dimension():
+    data = pa.table({"vector": [[1.0, 2.0], [3.0]]})
+
+    with pytest.raises(ValueError, match="only preserves vectors containing NaNs"):
+        _handle_bad_vectors(
+            data.to_reader(),
+            on_bad_vectors="keep",
+        ).read_all()
+
+
 def test_handle_bad_vectors_noop():
     # ChunkedArray should be preserved as-is
     vector = pa.chunked_array(

python/src/table.rs

@@ -21,7 +21,8 @@ use lancedb::blob::{BlobFile, BlobRangeRequest};
 use lancedb::index::scalar::FtsIndexBuilder;
 use lancedb::table::{
     AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken,
-    NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
+    NaNVectorBehavior, NewColumnTransform, OptimizeAction, OptimizeOptions, Ref,
+    Table as LanceDbTable,
 };
 use lancedb::tokenize as lancedb_tokenize;
 use pyo3::{
@@ -642,13 +643,14 @@ impl Table {
         })
     }
 
-    #[pyo3(signature = (data, mode, progress=None, write_parallelism=None))]
+    #[pyo3(signature = (data, mode, progress=None, write_parallelism=None, on_nan_vectors=None))]
     pub fn add<'a>(
         self_: PyRef<'a, Self>,
         data: PyScannable,
         mode: String,
         progress: Option<Py<PyAny>>,
         write_parallelism: Option<usize>,
+        on_nan_vectors: Option<String>,
     ) -> PyResult<Bound<'a, PyAny>> {
         let mut op = self_.inner_ref()?.add(data);
         if mode == "append" {
@@ -658,6 +660,18 @@ impl Table {
         } else {
             return Err(PyValueError::new_err(format!("Invalid mode: {}", mode)));
         }
+        match on_nan_vectors.as_deref() {
+            None | Some("error") => {}
+            Some("keep") => {
+                op = op.on_nan_vectors(NaNVectorBehavior::Keep);
+            }
+            Some(other) => {
+                return Err(PyValueError::new_err(format!(
+                    "Invalid on_nan_vectors: {}",
+                    other
+                )));
+            }
+        }
         if let Some(write_parallelism) = write_parallelism {
             op = op.write_parallelism(write_parallelism);
         }

rust/lancedb/src/query.rs

@@ -511,9 +511,7 @@ pub trait QueryBase {
 
     /// Rerank the results using the specified reranker.
     ///
-    /// For vector-only searches, the reranker receives the vector results and an
-    /// empty full-text result set and must declare its output schema. Reranking
-    /// multiple query vectors in one query is not supported.
+    /// This is currently only supported for Hybrid Search.
     fn rerank(self, reranker: Arc<dyn Reranker>) -> Self;
 
     /// The method to normalize the scores. Can be "rank" or "Score". If "Rank",
@@ -1140,44 +1138,6 @@ pub struct VectorQuery {
 }
 
 impl VectorQuery {
-    fn check_vector_rerank_supported(&self) -> Result<()> {
-        if self.request.query_vector.len() > 1 {
-            return Err(Error::NotSupported {
-                message: "reranking multiple query vectors is not supported; execute one query per vector"
-                    .to_string(),
-            });
-        }
-        Ok(())
-    }
-
-    async fn vector_rerank_output_schema(&self) -> Result<SchemaRef> {
-        self.check_vector_rerank_supported()?;
-
-        // Rerankers receive row IDs internally. Apply their schema transform to
-        // that exact input and then hide the row ID from the declared public
-        // schema unless it was explicitly requested.
-        let vector_query = self.clone().with_row_id();
-        let plan = vector_query
-            .create_plan(QueryExecutionOptions::default())
-            .await?;
-        let reranker = self
-            .request
-            .base
-            .reranker
-            .as_ref()
-            .expect("vector_rerank_output_schema requires a reranker");
-        let input_schema = plan.schema();
-        let output_schema = reranker.output_schema(&input_schema).await?;
-
-        if self.request.base.with_row_id {
-            Ok(output_schema)
-        } else {
-            Ok(RecordBatch::new_empty(output_schema)
-                .drop_column(ROW_ID)?
-                .schema())
-        }
-    }
-
     fn new(base: Query) -> Self {
         Self {
             parent: base.parent,
@@ -1483,61 +1443,6 @@ impl VectorQuery {
         Ok(single_batch_stream(results, max_batch_length))
     }
 
-    async fn execute_vector_rerank(
-        &self,
-        options: QueryExecutionOptions,
-    ) -> Result<SendableRecordBatchStream> {
-        self.check_vector_rerank_supported()?;
-
-        let max_batch_length = options.max_batch_length as usize;
-        let internal_options = options.without_output_batch_length_limit();
-
-        // RRF needs row IDs to assign and preserve scores. Keep them internal unless
-        // the caller explicitly requested them.
-        let vector_query = self.clone().with_row_id();
-        let vector_results = vector_query
-            .inner_execute_with_options(internal_options)
-            .await?;
-        let schema = vector_results.schema();
-        let vector_results = vector_results.try_collect::<Vec<_>>().await?;
-        let vector_results = concat_batches(&schema, vector_results.iter())?;
-        let vector_schema = vector_results.schema();
-        let fts_results = RecordBatch::new_empty(vector_schema.clone());
-
-        let reranker = self
-            .request
-            .base
-            .reranker
-            .as_ref()
-            .expect("execute_vector_rerank requires a reranker");
-        let expected_schema = reranker.output_schema(&vector_schema).await?;
-        let mut results = reranker
-            .rerank_hybrid("", vector_results, fts_results)
-            .await?;
-
-        check_reranker_result(&results)?;
-        if results.schema() != expected_schema {
-            return Err(Error::Schema {
-                message: format!(
-                    "reranker returned schema {:?}, but declared {:?}",
-                    results.schema(),
-                    expected_schema
-                ),
-            });
-        }
-
-        let limit = self.request.base.limit.unwrap_or(DEFAULT_TOP_K);
-        if results.num_rows() > limit {
-            results = results.slice(0, limit);
-        }
-
-        if !self.request.base.with_row_id {
-            results = results.drop_column(ROW_ID)?;
-        }
-
-        Ok(single_batch_stream(results, max_batch_length))
-    }
-
     async fn inner_execute_with_options(
         &self,
         options: QueryExecutionOptions,
@@ -1590,23 +1495,6 @@ impl ExecutableQuery for VectorQuery {
             return Ok(hybrid_result);
         }
 
-        if self.request.base.reranker.is_some() {
-            let timeout = options.timeout;
-            let mut rerank_options = options;
-            // A single outer deadline covers planning, candidate collection,
-            // schema declaration, and the complete reranker callback.
-            rerank_options.timeout = None;
-            let execution = self.execute_vector_rerank(rerank_options);
-            return match timeout {
-                Some(timeout) => tokio::time::timeout(timeout, execution)
-                    .await
-                    .map_err(|_| Error::Timeout {
-                        message: format!("Query timeout after {} ms", timeout.as_millis()),
-                    })?,
-                None => execution.await,
-            };
-        }
-
         self.inner_execute_with_options(options).await
     }
 
@@ -1619,15 +1507,6 @@ impl ExecutableQuery for VectorQuery {
         let query = AnyQuery::VectorQuery(self.request.clone());
         self.parent.analyze_plan(&query, options).await
     }
-
-    async fn output_schema(&self) -> Result<SchemaRef> {
-        if self.request.base.full_text_search.is_none() && self.request.base.reranker.is_some() {
-            self.vector_rerank_output_schema().await
-        } else {
-            let plan = self.create_plan(QueryExecutionOptions::default()).await?;
-            Ok(plan.schema())
-        }
-    }
 }
 
 impl HasQuery for VectorQuery {
@@ -1764,13 +1643,7 @@ impl ExecutableQuery for TakeQuery {
 
 #[cfg(test)]
 mod tests {
-    use std::{
-        collections::HashSet,
-        sync::{
-            Arc,
-            atomic::{AtomicBool, Ordering},
-        },
-    };
+    use std::{collections::HashSet, sync::Arc};
 
     use super::*;
     use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type};
@@ -1786,31 +1659,6 @@ mod tests {
 
     use crate::{Table, connect, database::CreateTableMode, index::Index};
 
-    #[derive(Debug)]
-    struct SlowReranker {
-        invoked: Arc<AtomicBool>,
-    }
-
-    #[async_trait::async_trait]
-    impl Reranker for SlowReranker {
-        async fn output_schema(&self, input: &SchemaRef) -> Result<SchemaRef> {
-            RRFReranker::default().output_schema(input).await
-        }
-
-        async fn rerank_hybrid(
-            &self,
-            query: &str,
-            vector_results: RecordBatch,
-            fts_results: RecordBatch,
-        ) -> Result<RecordBatch> {
-            self.invoked.store(true, Ordering::SeqCst);
-            tokio::time::sleep(Duration::from_secs(2)).await;
-            RRFReranker::default()
-                .rerank_hybrid(query, vector_results, fts_results)
-                .await
-        }
-    }
-
     #[tokio::test]
     async fn test_setters_getters() {
         // TODO: Switch back to memory://foo after https://github.com/lancedb/lancedb/issues/1051
@@ -2522,71 +2370,6 @@ mod tests {
         // We don't guarantee order.
         assert!(query_index.values().contains(&0));
         assert!(query_index.values().contains(&1));
-
-        let reranked = query.rerank(Arc::new(RRFReranker::default()));
-        let Err(execute_error) = reranked.execute().await else {
-            panic!("multi-vector reranking should be rejected");
-        };
-        assert!(
-            execute_error
-                .to_string()
-                .contains("reranking multiple query vectors is not supported")
-        );
-        let schema_error = reranked.output_schema().await.unwrap_err();
-        assert!(
-            schema_error
-                .to_string()
-                .contains("reranking multiple query vectors is not supported")
-        );
-    }
-
-    #[tokio::test]
-    async fn test_vector_rerank_timeout_covers_reranker() {
-        let tmp_dir = tempdir().unwrap();
-        let table = make_test_table(&tmp_dir).await;
-        let invoked = Arc::new(AtomicBool::new(false));
-        let reranker = SlowReranker {
-            invoked: invoked.clone(),
-        };
-
-        let result = table
-            .vector_search(&[0.1, 0.2, 0.3, 0.4])
-            .unwrap()
-            .limit(1)
-            .rerank(Arc::new(reranker))
-            .execute_with_options(QueryExecutionOptions {
-                timeout: Some(Duration::from_secs(1)),
-                ..Default::default()
-            })
-            .await;
-
-        assert!(invoked.load(Ordering::SeqCst));
-        assert!(matches!(result, Err(Error::Timeout { .. })));
-    }
-
-    #[tokio::test]
-    async fn test_vector_rerank_output_schema_matches_execution() {
-        let tmp_dir = tempdir().unwrap();
-        let table = make_test_table(&tmp_dir).await;
-        let query = table
-            .vector_search(&[0.1, 0.2, 0.3, 0.4])
-            .unwrap()
-            .limit(1)
-            .rerank(Arc::new(RRFReranker::default()));
-
-        let promised = query.output_schema().await.unwrap();
-        let actual = query
-            .execute()
-            .await
-            .unwrap()
-            .next()
-            .await
-            .unwrap()
-            .unwrap()
-            .schema();
-
-        assert_eq!(promised, actual);
-        assert!(promised.column_with_name("_relevance_score").is_some());
     }
 
     #[tokio::test]

rust/lancedb/src/rerankers.rs

@@ -8,7 +8,6 @@ use arrow::{
     compute::{concat_batches, filter_record_batch},
 };
 use arrow_array::{BooleanArray, RecordBatch, UInt64Array};
-use arrow_schema::SchemaRef;
 use async_trait::async_trait;
 use lance::dataset::ROW_ID;
 
@@ -48,28 +47,16 @@ impl std::fmt::Display for NormalizeMethod {
     }
 }
 
-/// Interface for a reranker. A reranker is used to rerank vector and hybrid
-/// search results. This is useful for combining results from multiple search
-/// methods or assigning a relevance score to vector search results.
+/// Interface for a reranker. A reranker is used to rerank the results from a
+/// vector and FTS search. This is useful for combining the results from both
+/// search methods.
 #[async_trait]
 pub trait Reranker: std::fmt::Debug + Sync + Send {
-    /// Declare the schema returned by [`Self::rerank_hybrid`] for a vector-only
-    /// query.
-    ///
-    /// Vector reranking validates the returned batch against this schema so
-    /// [`crate::query::ExecutableQuery::output_schema`] and execution cannot
-    /// disagree. Rerankers that only support hybrid search do not need to
-    /// implement this method.
-    async fn output_schema(&self, _input: &SchemaRef) -> Result<SchemaRef> {
-        Err(Error::NotSupported {
-            message: "vector rerankers must declare their output schema".to_string(),
-        })
-    }
+    // TODO support vector reranking and FTS reranking. Currently only hybrid reranking is supported.
 
     /// Rerank function receives the individual results from the vector and FTS search
     /// results. You can choose to use any of the results to generate the final results,
-    /// allowing maximum flexibility. For a vector-only search, `query` is empty and
-    /// `fts_results` is an empty batch with the same schema as `vector_results`.
+    /// allowing maximum flexibility.
     async fn rerank_hybrid(
         &self,
         query: &str,

rust/lancedb/src/rerankers/rrf.rs

@@ -9,7 +9,7 @@ use arrow::{
     compute::{sort_to_indices, take},
 };
 use arrow_array::{Float32Array, RecordBatch, UInt64Array};
-use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions};
+use arrow_schema::{DataType, Field, Schema, SortOptions};
 use async_trait::async_trait;
 use lance::dataset::ROW_ID;
 
@@ -44,19 +44,6 @@ impl Default for RRFReranker {
 
 #[async_trait]
 impl Reranker for RRFReranker {
-    async fn output_schema(&self, input: &SchemaRef) -> Result<SchemaRef> {
-        let mut fields = input.fields().to_vec();
-        fields.push(Arc::new(Field::new(
-            RELEVANCE_SCORE,
-            DataType::Float32,
-            false,
-        )));
-        Ok(Arc::new(Schema::new_with_metadata(
-            fields,
-            input.metadata().clone(),
-        )))
-    }
-
     async fn rerank_hybrid(
         &self,
         _query: &str,
@@ -148,9 +135,15 @@ impl Reranker for RRFReranker {
             .collect();
 
         // add relevance score to schema
-        let schema = self.output_schema(&combined_results.schema()).await?;
+        let mut fields = combined_results.schema().fields().to_vec();
+        fields.push(Arc::new(Field::new(
+            RELEVANCE_SCORE,
+            DataType::Float32,
+            false,
+        )));
+        let schema = Schema::new(fields);
 
-        let combined_results = RecordBatch::try_new(schema, columns)?;
+        let combined_results = RecordBatch::try_new(Arc::new(schema), columns)?;
 
         Ok(combined_results)
     }