fix(python): reject NaNs in multivector columns

<!-- lance-gatekeeper-fix:v1 agent=5c80c44c083b3b8ad0da595419d468fc
generation=1 -->

## Root cause

The legacy synchronous Python table called `delete` on a shared, mutable
`lance.Dataset`. Concurrent table operations could hold a PyO3 borrow
while delete requested an exclusive borrow, producing `RuntimeError:
Already borrowed`. The current async-backed binding fixes this by
cloning its thread-safe Rust table handle before awaiting, but that
concurrency contract had no regression coverage.

## Fix

- Document why delete must clone the Rust table handle before entering
its async future.
- Add a barrier-synchronized regression test that deletes distinct rows
through one shared table from eight Python threads.
- Verify every delete commits exactly one row, every commit gets a
distinct version, and no rows remain.

## Validation

- `cargo check --quiet --features remote --tests --examples`
- `cargo fmt --all -- --check`
- `uv run --extra tests --extra dev ruff format --check
python/tests/test_table.py`
- `uv run --extra tests --extra dev ruff check
python/tests/test_table.py`
- `uv run --extra tests --extra dev pytest
python/tests/test_table.py::test_concurrent_deletes_are_thread_safe
python/tests/test_table.py::test_delete
python/tests/test_table.py::test_delete_expr
python/tests/test_table.py::test_delete_expr_async -q` (4 passed)
- Manual stress reproduction: 100 concurrent deletes on one table
completed at versions 2–101 with zero rows remaining.

Fixes #530

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>

python/python/lancedb/db.py

@@ -707,6 +707,9 @@ class LanceDBConnection(DBConnection):
         self._namespace_client_properties = namespace_client_properties
         if _inner is not None:
             self._conn = _inner
+            # Native-derived wrappers resolve this in their async reconstruction
+            # path so construction never synchronously re-enters LOOP.
+            self._read_consistency_interval = read_consistency_interval
             self._cached_namespace_client = None
             return
 
@@ -756,11 +759,14 @@ class LanceDBConnection(DBConnection):
         # storage_options.  Also, this class really shouldn't be holding any state
         # beyond _conn.
         self._conn = AsyncConnection(LOOP.run(do_connect()))
+        # Keep property access synchronous so debugger introspection cannot wait on
+        # the background loop while that thread is suspended at a breakpoint.
+        self._read_consistency_interval = read_consistency_interval
         self._cached_namespace_client: Optional[LanceNamespace] = None
 
     @property
     def read_consistency_interval(self) -> Optional[timedelta]:
-        return LOOP.run(self._conn.get_read_consistency_interval())
+        return self._read_consistency_interval
 
     @property
     def session(self) -> Optional[Session]:
@@ -771,8 +777,16 @@ class LanceDBConnection(DBConnection):
         return self._conn.uri
 
     @classmethod
-    def from_inner(cls, inner: LanceDbConnection):
-        return cls(None, _inner=inner)
+    def from_inner(
+        cls,
+        inner: LanceDbConnection,
+        read_consistency_interval: Optional[timedelta],
+    ):
+        return cls(
+            None,
+            read_consistency_interval=read_consistency_interval,
+            _inner=inner,
+        )
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(uri={self._conn.uri!r})"

python/python/lancedb/permutation.py

@@ -226,7 +226,7 @@ class PermutationBuilder:
 
         async def do_execute():
             inner_tbl = await self._async.execute()
-            return LanceTable.from_inner(inner_tbl)
+            return await LanceTable.from_inner(inner_tbl)
 
         return LOOP.run(do_execute())
 

python/python/lancedb/table.py

@@ -2182,11 +2182,15 @@ class LanceTable(Table):
         return self.name
 
     @classmethod
-    def from_inner(cls, tbl: LanceDBTable):
-        from .db import LanceDBConnection
+    async def from_inner(cls, tbl: LanceDBTable):
+        from .db import AsyncConnection, LanceDBConnection
 
         async_tbl = AsyncTable(tbl)
-        conn = LanceDBConnection.from_inner(tbl.database())
+        inner_conn = tbl.database()
+        read_consistency_interval = await AsyncConnection(
+            inner_conn
+        ).get_read_consistency_interval()
+        conn = LanceDBConnection.from_inner(inner_conn, read_consistency_interval)
         return cls(
             conn,
             async_tbl.name,
@@ -4031,7 +4035,10 @@ def _handle_bad_vectors(
             for vector_column in vector_columns:
                 dim = vector_column["expected_dim"]
                 if target_schema is not None and dim is None:
-                    dim = _infer_vector_dim(batch[vector_column["name"]])
+                    dim = _infer_vector_column_dim(
+                        batch[vector_column["name"]],
+                        vector_column["is_multivector"],
+                    )
                     pending_dims.append(vector_column)
                 batch = _handle_bad_vector_column(
                     batch,
@@ -4040,11 +4047,13 @@ def _handle_bad_vectors(
                     fill_value=fill_value,
                     expected_dim=dim,
                     expected_value_type=vector_column["expected_value_type"],
+                    is_multivector=vector_column["is_multivector"],
                 )
             for vector_column in pending_dims:
                 if vector_column["expected_dim"] is None:
-                    vector_column["expected_dim"] = _infer_vector_dim(
-                        batch[vector_column["name"]]
+                    vector_column["expected_dim"] = _infer_vector_column_dim(
+                        batch[vector_column["name"]],
+                        vector_column["is_multivector"],
                     )
             if batch.schema.equals(output_schema, check_metadata=True):
                 yield batch
@@ -4070,22 +4079,30 @@ def _find_vector_columns(
     if target_schema is None:
         vector_columns = []
         for field in reader_schema:
-            named_vector_col = (
-                _is_list_like(field.type)
-                and pa.types.is_floating(field.type.value_type)
-                and field.name == VECTOR_COLUMN_NAME
+            is_multivector = _is_multivector_type(field.type)
+            is_fixed_multivector = is_multivector and pa.types.is_fixed_size_list(
+                field.type.value_type
             )
+            named_vector_col = (
+                _is_float_vector_type(field.type) or is_multivector
+            ) and field.name == VECTOR_COLUMN_NAME
             likely_vector_col = (
                 pa.types.is_fixed_size_list(field.type)
                 and pa.types.is_floating(field.type.value_type)
                 and (field.type.list_size >= 10)
             )
-            if named_vector_col or likely_vector_col:
+            if named_vector_col or likely_vector_col or is_fixed_multivector:
+                vector_type = field.type.value_type if is_multivector else field.type
                 vector_columns.append(
                     {
                         "name": field.name,
-                        "expected_dim": None,
-                        "expected_value_type": None,
+                        "expected_dim": (
+                            vector_type.list_size
+                            if pa.types.is_fixed_size_list(vector_type)
+                            else None
+                        ),
+                        "expected_value_type": vector_type.value_type,
+                        "is_multivector": is_multivector,
                     }
                 )
         return vector_columns
@@ -4099,9 +4116,8 @@ def _find_vector_columns(
     for field in target_schema:
         if field.name not in reader_column_names:
             continue
-        if not _is_list_like(field.type) or not pa.types.is_floating(
-            field.type.value_type
-        ):
+        is_multivector = _is_multivector_type(field.type)
+        if not _is_float_vector_type(field.type) and not is_multivector:
             continue
 
         reader_field = reader_schema.field(field.name)
@@ -4116,16 +4132,18 @@ def _find_vector_columns(
             and reader_field.type.list_size >= 10
         )
 
-        if named_vector_col or typed_fixed_vector_col:
+        if named_vector_col or typed_fixed_vector_col or is_multivector:
+            vector_type = field.type.value_type if is_multivector else field.type
             vector_columns.append(
                 {
                     "name": field.name,
                     "expected_dim": (
-                        field.type.list_size
-                        if pa.types.is_fixed_size_list(field.type)
+                        vector_type.list_size
+                        if pa.types.is_fixed_size_list(vector_type)
                         else None
                     ),
-                    "expected_value_type": field.type.value_type,
+                    "expected_value_type": vector_type.value_type,
+                    "is_multivector": is_multivector,
                 }
             )
 
@@ -4176,6 +4194,7 @@ def _handle_bad_vector_column(
     fill_value: float = 0.0,
     expected_dim: Optional[int] = None,
     expected_value_type: Optional[pa.DataType] = None,
+    is_multivector: bool = False,
 ) -> pa.RecordBatch:
     """
     Ensure that the vector column exists and has type fixed_size_list(float)
@@ -4196,6 +4215,8 @@ def _handle_bad_vector_column(
     vec_arr = data[vector_column_name]
     if not _is_list_like(vec_arr.type):
         return data
+    if is_multivector and not _is_multivector_type(vec_arr.type):
+        return data
 
     if (
         expected_dim is not None
@@ -4212,12 +4233,21 @@ def _handle_bad_vector_column(
         vec_arr = pa.array(vec_arr.to_pylist(), type=pa.list_(expected_value_type))
         data = data.set_column(position, vector_column_name, vec_arr)
 
-    if pa.types.is_floating(vec_arr.type.value_type):
+    if is_multivector or pa.types.is_floating(vec_arr.type.value_type):
         has_nan = has_nan_values(vec_arr)
     else:
         has_nan = pa.array([False] * len(vec_arr))
 
-    if expected_dim is not None:
+    if is_multivector:
+        dim = (
+            expected_dim
+            if expected_dim is not None
+            else _infer_vector_column_dim(vec_arr, True)
+        )
+        if dim is None:
+            return data
+        has_wrong_dim = _multivector_has_wrong_dim(vec_arr, dim)
+    elif expected_dim is not None:
         dim = expected_dim
     elif pa.types.is_fixed_size_list(vec_arr.type):
         dim = vec_arr.type.list_size
@@ -4226,15 +4256,16 @@ def _handle_bad_vector_column(
         if dim is None:
             return data
 
-    is_null = pc.is_null(vec_arr)
-    # pc.list_value_length returns null for null list entries, so
-    # pc.not_equal(null, dim) also returns null. Use or_kleene so that
-    # True OR null = True (Kleene three-valued logic), ensuring null vectors
-    # are counted as wrong-dim.
-    has_wrong_dim = pc.or_kleene(
-        is_null,
-        pc.not_equal(pc.list_value_length(vec_arr), dim),
-    )
+    if not is_multivector:
+        is_null = pc.is_null(vec_arr)
+        # pc.list_value_length returns null for null list entries, so
+        # pc.not_equal(null, dim) also returns null. Use or_kleene so that
+        # True OR null = True (Kleene three-valued logic), ensuring null vectors
+        # are counted as wrong-dim.
+        has_wrong_dim = pc.or_kleene(
+            is_null,
+            pc.not_equal(pc.list_value_length(vec_arr), dim),
+        )
 
     has_bad_vectors = pc.any(has_nan).as_py() or pc.any(has_wrong_dim).as_py()
 
@@ -4269,7 +4300,10 @@ def _handle_bad_vector_column(
                 raise ValueError(
                     "`fill_value` must not be None if `on_bad_vectors` is 'fill'"
                 )
-            vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value)
+            if is_multivector:
+                vec_arr = _fill_bad_multivector_values(vec_arr, dim, fill_value)
+            else:
+                vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value)
         else:
             raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}")
 
@@ -4321,17 +4355,60 @@ def _fill_bad_vector_values(
     return filled.cast(arr.type)
 
 
+def _fill_bad_multivector_values(
+    arr: Union[pa.Array, pa.ChunkedArray], dim: int, fill_value: float
+) -> pa.Array:
+    if not isinstance(arr, pa.ChunkedArray):
+        arr = pa.chunked_array([arr])
+    arr = arr.combine_chunks()
+
+    filled_vectors = _fill_bad_vector_values(arr.values, dim, fill_value)
+    parent_nulls = pc.is_null(arr)
+    if pa.types.is_large_list(arr.type):
+        filled = pa.LargeListArray.from_arrays(
+            arr.offsets, filled_vectors, mask=parent_nulls
+        )
+    else:
+        filled = pa.ListArray.from_arrays(
+            arr.offsets, filled_vectors, mask=parent_nulls
+        )
+    return filled.cast(arr.type)
+
+
+def _multivector_has_wrong_dim(
+    arr: Union[pa.Array, pa.ChunkedArray], dim: int
+) -> pa.BooleanArray:
+    if isinstance(arr, pa.ChunkedArray):
+        results = [_multivector_has_wrong_dim(chunk, dim) for chunk in arr.chunks]
+        return pa.concat_arrays(results) if results else pa.array([], type=pa.bool_())
+
+    vectors = arr.flatten()
+    vector_is_wrong = pc.or_kleene(
+        pc.is_null(vectors),
+        pc.not_equal(pc.list_value_length(vectors), dim),
+    )
+    parent_indices = pc.list_parent_indices(arr)
+    wrong_parent_indices = pc.unique(pc.filter(parent_indices, vector_is_wrong))
+    indices = pa.array(range(len(arr)), type=pa.uint32())
+    return pc.or_(pc.is_null(arr), pc.is_in(indices, wrong_parent_indices))
+
+
 def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray:
     if isinstance(arr, pa.ChunkedArray):
-        values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks])
-    else:
-        values = arr.flatten()
-    if pa.types.is_float16(values.type):
+        results = [has_nan_values(chunk) for chunk in arr.chunks]
+        return pa.concat_arrays(results) if results else pa.array([], type=pa.bool_())
+
+    values = arr.flatten()
+    if _is_list_like(values.type):
+        values_has_nan = has_nan_values(values)
+    elif pa.types.is_float16(values.type):
         # is_nan isn't yet implemented for f16, so we cast to f32
         # https://github.com/apache/arrow/issues/45083
         values_has_nan = pc.is_nan(values.cast(pa.float32()))
-    else:
+    elif pa.types.is_floating(values.type):
         values_has_nan = pc.is_nan(values)
+    else:
+        return pa.array([False] * len(arr))
     values_indices = pc.list_parent_indices(arr)
     has_nan_indices = pc.unique(pc.filter(values_indices, values_has_nan))
     indices = pa.array(range(len(arr)), type=pa.uint32())
@@ -4346,6 +4423,16 @@ def _is_list_like(data_type: pa.DataType) -> bool:
     )
 
 
+def _is_float_vector_type(data_type: pa.DataType) -> bool:
+    return _is_list_like(data_type) and pa.types.is_floating(data_type.value_type)
+
+
+def _is_multivector_type(data_type: pa.DataType) -> bool:
+    return (
+        pa.types.is_list(data_type) or pa.types.is_large_list(data_type)
+    ) and _is_float_vector_type(data_type.value_type)
+
+
 def _merge_metadata(*metadata_dicts: Optional[dict]) -> dict:
     merged = {}
     for metadata in metadata_dicts:
@@ -4437,6 +4524,16 @@ def _infer_vector_dim(arr: Union[pa.Array, pa.ChunkedArray]) -> Optional[int]:
     return pc.mode(lengths)[0].as_py()["mode"]
 
 
+def _infer_vector_column_dim(
+    arr: Union[pa.Array, pa.ChunkedArray], is_multivector: bool
+) -> Optional[int]:
+    if not is_multivector:
+        return _infer_vector_dim(arr)
+    if isinstance(arr, pa.ChunkedArray):
+        arr = arr.combine_chunks()
+    return _infer_vector_dim(arr.flatten())
+
+
 def _validate_schema(schema: pa.Schema):
     """
     Make sure the metadata is valid utf8

python/python/tests/test_db.py

@@ -77,6 +77,23 @@ def test_sync_repr_does_not_use_background_loop(tmp_path, monkeypatch):
     assert repr(table) == f"LanceTable(name='test', _conn={db!r})"
 
 
+def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkeypatch):
+    from lancedb.background_loop import LOOP
+    from lancedb.db import LanceDBConnection
+
+    consistency_interval = timedelta(seconds=5)
+    db = lancedb.connect(tmp_path, read_consistency_interval=consistency_interval)
+    db_from_inner = LanceDBConnection.from_inner(db._inner, consistency_interval)
+
+    def fail_run(*args, **kwargs):
+        raise AssertionError("properties should not use the Python background loop")
+
+    monkeypatch.setattr(LOOP, "run", fail_run)
+
+    assert db.read_consistency_interval == consistency_interval
+    assert db_from_inner.read_consistency_interval == consistency_interval
+
+
 def test_ingest_pd(tmp_path):
     db = lancedb.connect(tmp_path)
 

python/python/tests/test_permutation.py

@@ -6,6 +6,7 @@ import math
 import pytest
 
 from lancedb import DBConnection, Table, connect
+from lancedb.background_loop import LOOP
 from lancedb.permutation import Permutation, Permutations, permutation_builder
 
 
@@ -31,6 +32,25 @@ def test_split_random_ratios(mem_db):
     assert 65 <= split_1_count <= 75  # ~70% ± tolerance
 
 
+def test_execute_does_not_reenter_background_loop(tmp_path, monkeypatch):
+    import threading
+
+    db = connect(tmp_path)
+    tbl = db.create_table("test_table", pa.table({"x": range(10)}))
+    original_run = LOOP.run
+
+    def fail_on_reentry(future):
+        assert threading.current_thread() is not LOOP.thread
+        return original_run(future)
+
+    monkeypatch.setattr(LOOP, "run", fail_on_reentry)
+
+    permutation_tbl = permutation_builder(tbl).execute()
+
+    assert permutation_tbl.count_rows() == 10
+    assert permutation_tbl._conn.read_consistency_interval is None
+
+
 def test_split_random_counts(mem_db):
     """Test random splitting with absolute counts."""
     tbl = mem_db.create_table(

python/python/tests/test_table.py

@@ -6,6 +6,7 @@ import os
 import sys
 import threading
 import warnings
+from concurrent.futures import ThreadPoolExecutor
 from datetime import date, datetime, timedelta
 from time import sleep
 from typing import List
@@ -1709,6 +1710,19 @@ def test_create_with_nans(mem_db: DBConnection):
     assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
 
 
+def test_create_with_nans_in_multivectors(mem_db: DBConnection):
+    multivector_type = pa.list_(pa.list_(pa.float32(), 128))
+    schema = pa.schema(
+        [pa.field("filename", pa.string()), pa.field("vector", multivector_type)]
+    )
+    vector = [0.1] * 128
+    vector[-1] = np.nan
+    data = [{"filename": "img1.jpg", "vector": [vector]}]
+
+    with pytest.raises(RuntimeError, match="Vector column 'vector' has NaNs"):
+        mem_db.create_table("nan_multivector", data=data, schema=schema)
+
+
 def test_add_with_nans(mem_db: DBConnection):
     schema = pa.schema(
         [
@@ -2124,6 +2138,27 @@ def test_delete(mem_db: DBConnection):
     assert table.to_arrow()["id"].to_pylist() == [1]
 
 
+def test_concurrent_deletes_are_thread_safe(mem_db: DBConnection):
+    num_workers = 8
+    table = mem_db.create_table(
+        "my_table", data=[{"id": row_id} for row_id in range(num_workers)]
+    )
+    barrier = threading.Barrier(num_workers)
+
+    def delete(row_id: int):
+        barrier.wait()
+        return table.delete(f"id = {row_id}")
+
+    with ThreadPoolExecutor(max_workers=num_workers) as pool:
+        results = list(pool.map(delete, range(num_workers)))
+
+    assert all(result.num_deleted_rows == 1 for result in results)
+    assert sorted(result.version for result in results) == list(
+        range(2, num_workers + 2)
+    )
+    assert table.count_rows() == 0
+
+
 def test_delete_expr(mem_db: DBConnection):
     table = mem_db.create_table(
         "my_table",

python/python/tests/test_util.py

@@ -400,6 +400,83 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
     assert output["vector"].combine_chunks() == expected
 
 
+@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
+def test_handle_bad_multivectors_nan(on_bad_vectors):
+    multivector_type = pa.list_(pa.list_(pa.float32(), 2))
+    vectors = pa.array(
+        [
+            [[1.0, float("nan")], [2.0, 3.0]],
+            [[4.0, 5.0]],
+        ],
+        type=multivector_type,
+    )
+    data = pa.table({"vector": vectors})
+
+    if on_bad_vectors == "error":
+        with pytest.raises(ValueError, match="Vector column 'vector' has NaNs"):
+            _handle_bad_vectors(data.to_reader()).read_all()
+        return
+
+    output = _handle_bad_vectors(
+        data.to_reader(),
+        on_bad_vectors=on_bad_vectors,
+        fill_value=42.0,
+    ).read_all()
+
+    if on_bad_vectors == "drop":
+        expected = pa.array([[[4.0, 5.0]]], type=multivector_type)
+    elif on_bad_vectors == "fill":
+        expected = pa.array(
+            [[[1.0, 42.0], [2.0, 3.0]], [[4.0, 5.0]]],
+            type=multivector_type,
+        )
+    else:
+        expected = pa.array([None, [[4.0, 5.0]]], type=multivector_type)
+
+    assert output["vector"].combine_chunks() == expected
+
+
+@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
+def test_handle_bad_variable_multivectors(on_bad_vectors):
+    target_type = pa.list_(pa.list_(pa.float32(), 2))
+    vectors = pa.array(
+        [
+            [[1.0, float("nan")], [2.0, 3.0]],
+            [[4.0]],
+            [[5.0, 6.0]],
+        ]
+    )
+    data = pa.table({"vector": vectors})
+
+    if on_bad_vectors == "error":
+        with pytest.raises(ValueError, match="variable length vectors"):
+            _handle_bad_vectors(
+                data.to_reader(),
+                target_schema=pa.schema({"vector": target_type}),
+            ).read_all()
+        return
+
+    output = _handle_bad_vectors(
+        data.to_reader(),
+        on_bad_vectors=on_bad_vectors,
+        fill_value=42.0,
+        target_schema=pa.schema({"vector": target_type}),
+    ).read_all()
+
+    if on_bad_vectors == "drop":
+        expected = [[[5.0, 6.0]]]
+    elif on_bad_vectors == "fill":
+        expected = [
+            [[1.0, 42.0], [2.0, 3.0]],
+            [[4.0, 42.0]],
+            [[5.0, 6.0]],
+        ]
+    else:
+        expected = [None, None, [[5.0, 6.0]]]
+
+    assert output["vector"].combine_chunks().to_pylist() == expected
+
+
 def test_handle_bad_vectors_noop():
     # ChunkedArray should be preserved as-is
     vector = pa.chunked_array(

python/src/table.rs

@@ -745,6 +745,9 @@ impl Table {
 
     #[allow(private_interfaces)]
     pub fn delete(self_: PyRef<'_, Self>, condition: PredicateArg) -> PyResult<Bound<'_, PyAny>> {
+        // Do not hold the Python borrow across the await.  The cloned Rust table
+        // handle is thread-safe and allows deletes on the same Python table to
+        // run concurrently without PyO3 reporting "Already borrowed".
         let inner = self_.inner_ref()?.clone();
         future_into_py(self_.py(), async move {
             let result = match &condition {

rust/lancedb/src/table/add_data.rs

@@ -258,6 +258,7 @@ mod tests {
         FixedSizeListArray, Float32Array, Int32Array, LargeStringArray, ListArray, RecordBatch,
         RecordBatchIterator, record_batch,
     };
+    use arrow_buffer::OffsetBuffer;
     use arrow_schema::{ArrowError, DataType, Field, Schema};
     use futures::TryStreamExt;
     use lance::dataset::{WriteMode, WriteParams};
@@ -885,6 +886,63 @@ mod tests {
         assert_eq!(row_count, 1);
     }
 
+    #[tokio::test]
+    async fn test_add_rejects_nan_multivectors() {
+        let vector_type =
+            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4);
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "embedding",
+            DataType::List(Arc::new(Field::new("item", vector_type.clone(), true))),
+            false,
+        )]));
+
+        let db = connect("memory://").execute().await.unwrap();
+        let table = db
+            .create_empty_table("nan_multivector_test", schema.clone())
+            .execute()
+            .await
+            .unwrap();
+
+        let vectors = FixedSizeListArray::try_new(
+            Arc::new(Field::new("item", DataType::Float32, true)),
+            4,
+            Arc::new(Float32Array::from(vec![
+                0.1,
+                0.2,
+                0.3,
+                0.4,
+                0.5,
+                f32::NAN,
+                0.7,
+                0.8,
+            ])),
+            None,
+        )
+        .unwrap();
+        let multivectors = ListArray::try_new(
+            Arc::new(Field::new("item", vector_type, true)),
+            OffsetBuffer::from_lengths([2]),
+            Arc::new(vectors),
+            None,
+        )
+        .unwrap();
+        let batch = RecordBatch::try_new(schema, vec![Arc::new(multivectors)]).unwrap();
+
+        let err = table.add(batch.clone()).execute().await.unwrap_err();
+        assert!(
+            err.to_string().contains("NaN"),
+            "Expected error mentioning NaN values, but got: {err:?}"
+        );
+
+        table
+            .add(batch)
+            .on_nan_vectors(NaNVectorBehavior::Keep)
+            .execute()
+            .await
+            .unwrap();
+        assert_eq!(table.count_rows(None).await.unwrap(), 1);
+    }
+
     #[tokio::test]
     async fn test_add_subschema() {
         let data = record_batch!(("id", Int64, [4, 5]), ("text", Utf8, ["foo", "bar"])).unwrap();

rust/lancedb/src/table/datafusion/reject_nan.rs

@@ -5,7 +5,7 @@
 
 use std::sync::{Arc, LazyLock};
 
-use arrow_array::{Array, FixedSizeListArray};
+use arrow_array::{Array, FixedSizeListArray, ListArray};
 use arrow_schema::{DataType, Field, FieldRef};
 use datafusion_common::config::ConfigOptions;
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
@@ -19,16 +19,20 @@ use crate::{Error, Result};
 static REJECT_NAN_UDF: LazyLock<Arc<datafusion_expr::ScalarUDF>> =
     LazyLock::new(|| Arc::new(datafusion_expr::ScalarUDF::from(RejectNanUdf::new())));
 
-/// Returns true if the field is a vector column: FixedSizeList<Float16/32/64>.
+/// Returns true if the field is a vector or multivector column.
 fn is_vector_field(field: &Field) -> bool {
-    if let DataType::FixedSizeList(child, _) = field.data_type() {
-        matches!(
-            child.data_type(),
-            DataType::Float16 | DataType::Float32 | DataType::Float64
-        )
-    } else {
-        false
+    fn is_vector_data_type(data_type: &DataType) -> bool {
+        match data_type {
+            DataType::FixedSizeList(child, _) => matches!(
+                child.data_type(),
+                DataType::Float16 | DataType::Float32 | DataType::Float64
+            ),
+            DataType::List(child) => is_vector_data_type(child.data_type()),
+            _ => false,
+        }
     }
+
+    is_vector_data_type(field.data_type())
 }
 
 /// Wraps the input plan with a projection that checks vector columns for NaN values.
@@ -69,8 +73,8 @@ pub fn reject_nan_vectors(input: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn Execu
     Ok(Arc::new(projection))
 }
 
-/// A scalar UDF that passes through FixedSizeList arrays unchanged, but errors
-/// if any float values in the list are NaN.
+/// A scalar UDF that passes through vector arrays unchanged, but errors if any
+/// float values in the vector or multivector are NaN.
 #[derive(Debug, Hash, PartialEq, Eq)]
 struct RejectNanUdf {
     signature: Signature,
@@ -113,44 +117,54 @@ impl ScalarUDFImpl for RejectNanUdf {
 }
 
 fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> {
-    let fsl = array
-        .as_any()
-        .downcast_ref::<FixedSizeListArray>()
-        .ok_or_else(|| {
-            datafusion_common::DataFusionError::Internal(
-                "reject_nan expected FixedSizeList".to_string(),
-            )
-        })?;
-
-    // Only inspect elements that are both in a valid parent row and non-null
-    // themselves. Values backing null parent rows or null child elements may
-    // contain garbage (including NaN) per the Arrow spec.
-    let has_nan = (0..fsl.len()).filter(|i| fsl.is_valid(*i)).any(|i| {
-        let row = fsl.value(i);
-        match row.data_type() {
-            DataType::Float16 => row
+    fn contains_nan(array: &dyn Array) -> datafusion_common::Result<bool> {
+        match array.data_type() {
+            DataType::Float16 => Ok(array
                 .as_any()
                 .downcast_ref::<arrow_array::Float16Array>()
                 .unwrap()
                 .iter()
-                .any(|v| v.is_some_and(|v| v.is_nan())),
-            DataType::Float32 => row
+                .any(|v| v.is_some_and(|v| v.is_nan()))),
+            DataType::Float32 => Ok(array
                 .as_any()
                 .downcast_ref::<arrow_array::Float32Array>()
                 .unwrap()
                 .iter()
-                .any(|v| v.is_some_and(|v| v.is_nan())),
-            DataType::Float64 => row
+                .any(|v| v.is_some_and(|v| v.is_nan()))),
+            DataType::Float64 => Ok(array
                 .as_any()
                 .downcast_ref::<arrow_array::Float64Array>()
                 .unwrap()
                 .iter()
-                .any(|v| v.is_some_and(|v| v.is_nan())),
-            _ => false,
+                .any(|v| v.is_some_and(|v| v.is_nan()))),
+            DataType::FixedSizeList(_, _) => {
+                let lists = array.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
+                for i in (0..lists.len()).filter(|i| lists.is_valid(*i)) {
+                    if contains_nan(lists.value(i).as_ref())? {
+                        return Ok(true);
+                    }
+                }
+                Ok(false)
+            }
+            DataType::List(_) => {
+                let lists = array.as_any().downcast_ref::<ListArray>().unwrap();
+                for i in (0..lists.len()).filter(|i| lists.is_valid(*i)) {
+                    if contains_nan(lists.value(i).as_ref())? {
+                        return Ok(true);
+                    }
+                }
+                Ok(false)
+            }
+            data_type => Err(datafusion_common::DataFusionError::Internal(format!(
+                "reject_nan expected a vector or multivector, got {data_type}"
+            ))),
         }
-    });
+    }
 
-    if has_nan {
+    // Only inspect elements that are both in a valid parent row and non-null
+    // themselves. Values backing null parent rows or null child elements may
+    // contain garbage (including NaN) per the Arrow spec.
+    if contains_nan(array)? {
         return Err(datafusion_common::DataFusionError::ArrowError(
             Box::new(arrow_schema::ArrowError::ComputeError(
                 "Vector column contains NaN values".to_string(),
@@ -166,6 +180,7 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> {
 mod tests {
     use super::*;
     use arrow_array::Float32Array;
+    use arrow_buffer::OffsetBuffer;
 
     #[test]
     fn test_passes_clean_vectors() {
@@ -191,6 +206,46 @@ mod tests {
         assert!(check_no_nans(&fsl).is_err());
     }
 
+    #[test]
+    fn test_rejects_nan_multivectors() {
+        let vectors = FixedSizeListArray::try_new(
+            Arc::new(Field::new("item", DataType::Float32, true)),
+            2,
+            Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, f32::NAN])),
+            None,
+        )
+        .unwrap();
+        let multivectors = ListArray::try_new(
+            Arc::new(Field::new("item", vectors.data_type().clone(), true)),
+            OffsetBuffer::from_lengths([2]),
+            Arc::new(vectors),
+            None,
+        )
+        .unwrap();
+
+        assert!(check_no_nans(&multivectors).is_err());
+    }
+
+    #[test]
+    fn test_skips_null_multivector_rows() {
+        let vectors = FixedSizeListArray::try_new(
+            Arc::new(Field::new("item", DataType::Float32, true)),
+            2,
+            Arc::new(Float32Array::from(vec![f32::NAN, f32::NAN, 1.0, 2.0])),
+            None,
+        )
+        .unwrap();
+        let multivectors = ListArray::try_new(
+            Arc::new(Field::new("item", vectors.data_type().clone(), true)),
+            OffsetBuffer::from_lengths([1, 1]),
+            Arc::new(vectors),
+            Some(vec![false, true].into()),
+        )
+        .unwrap();
+
+        assert!(check_no_nans(&multivectors).is_ok());
+    }
+
     #[test]
     fn test_skips_null_rows() {
         // Values backing null rows may contain NaN per the Arrow spec.
@@ -254,6 +309,15 @@ mod tests {
             DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float64, true)), 4),
             false,
         )));
+        assert!(is_vector_field(&Field::new(
+            "v",
+            DataType::List(Arc::new(Field::new(
+                "item",
+                DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4,),
+                true,
+            ))),
+            false,
+        )));
         assert!(!is_vector_field(&Field::new("id", DataType::Int32, false)));
         assert!(!is_vector_field(&Field::new(
             "v",