fix(python): scope Instructor compatibility shim

python/python/lancedb/embeddings/instructor.py

@@ -3,6 +3,7 @@
 
 
 from typing import List
+from urllib.parse import unquote, urlparse
 
 import numpy as np
 
@@ -125,9 +126,20 @@ class InstructorEmbeddingFunction(TextEmbeddingFunction):
 
     @weak_lru(maxsize=1)
     def get_model(self):
-        instructor_embedding = attempt_import_or_raise(
-            "InstructorEmbedding", "InstructorEmbedding"
-        )
+        huggingface_hub = attempt_import_or_raise("huggingface_hub", "huggingface-hub")
+        missing = object()
+        original_cached_download = getattr(huggingface_hub, "cached_download", missing)
+        if original_cached_download is missing:
+            huggingface_hub.cached_download = _cached_download(huggingface_hub)
+
+        try:
+            instructor_embedding = attempt_import_or_raise(
+                "InstructorEmbedding", "InstructorEmbedding"
+            )
+        finally:
+            if original_cached_download is missing:
+                del huggingface_hub.cached_download
+
         torch = attempt_import_or_raise("torch", "torch")
 
         model = instructor_embedding.INSTRUCTOR(self.name)
@@ -140,3 +152,44 @@ class InstructorEmbeddingFunction(TextEmbeddingFunction):
                 model, {torch.nn.Linear}, dtype=torch.qint8
             )
         return model
+
+
+def _cached_download(huggingface_hub):
+    """Provide the legacy download API used by sentence-transformers 2.2.x."""
+
+    def cached_download(
+        *,
+        url,
+        cache_dir=None,
+        force_filename=None,
+        library_name=None,
+        library_version=None,
+        user_agent=None,
+        use_auth_token=None,
+        **_,
+    ):
+        path = urlparse(url).path.lstrip("/")
+        try:
+            repo_id, resolved_path = path.split("/resolve/", maxsplit=1)
+            revision, filename = resolved_path.split("/", maxsplit=1)
+        except ValueError as err:
+            raise ValueError(f"Unsupported Hugging Face Hub URL: {url}") from err
+
+        repo_id = unquote(repo_id)
+        revision = unquote(revision)
+        filename = unquote(filename)
+        # sentence-transformers derives force_filename from this Hub path with
+        # os.path.join. Using the URL path beneath local_dir produces the same
+        # local destination without sending Windows separators to the Hub.
+        return huggingface_hub.hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            revision=revision,
+            local_dir=cache_dir,
+            library_name=library_name,
+            library_version=library_version,
+            user_agent=user_agent,
+            token=use_auth_token,
+        )
+
+    return cached_download

python/python/lancedb/table.py

@@ -4035,10 +4035,7 @@ 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_column_dim(
-                        batch[vector_column["name"]],
-                        vector_column["is_multivector"],
-                    )
+                    dim = _infer_vector_dim(batch[vector_column["name"]])
                     pending_dims.append(vector_column)
                 batch = _handle_bad_vector_column(
                     batch,
@@ -4047,13 +4044,11 @@ 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_column_dim(
-                        batch[vector_column["name"]],
-                        vector_column["is_multivector"],
+                    vector_column["expected_dim"] = _infer_vector_dim(
+                        batch[vector_column["name"]]
                     )
             if batch.schema.equals(output_schema, check_metadata=True):
                 yield batch
@@ -4079,30 +4074,22 @@ def _find_vector_columns(
     if target_schema is None:
         vector_columns = []
         for field in reader_schema:
-            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
+                _is_list_like(field.type)
+                and pa.types.is_floating(field.type.value_type)
+                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 or is_fixed_multivector:
-                vector_type = field.type.value_type if is_multivector else field.type
+            if named_vector_col or likely_vector_col:
                 vector_columns.append(
                     {
                         "name": field.name,
-                        "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,
+                        "expected_dim": None,
+                        "expected_value_type": None,
                     }
                 )
         return vector_columns
@@ -4116,8 +4103,9 @@ def _find_vector_columns(
     for field in target_schema:
         if field.name not in reader_column_names:
             continue
-        is_multivector = _is_multivector_type(field.type)
-        if not _is_float_vector_type(field.type) and not is_multivector:
+        if not _is_list_like(field.type) or not pa.types.is_floating(
+            field.type.value_type
+        ):
             continue
 
         reader_field = reader_schema.field(field.name)
@@ -4132,18 +4120,16 @@ def _find_vector_columns(
             and reader_field.type.list_size >= 10
         )
 
-        if named_vector_col or typed_fixed_vector_col or is_multivector:
-            vector_type = field.type.value_type if is_multivector else field.type
+        if named_vector_col or typed_fixed_vector_col:
             vector_columns.append(
                 {
                     "name": field.name,
                     "expected_dim": (
-                        vector_type.list_size
-                        if pa.types.is_fixed_size_list(vector_type)
+                        field.type.list_size
+                        if pa.types.is_fixed_size_list(field.type)
                         else None
                     ),
-                    "expected_value_type": vector_type.value_type,
-                    "is_multivector": is_multivector,
+                    "expected_value_type": field.type.value_type,
                 }
             )
 
@@ -4194,7 +4180,6 @@ 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)
@@ -4215,8 +4200,6 @@ 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
@@ -4233,21 +4216,12 @@ 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 is_multivector or pa.types.is_floating(vec_arr.type.value_type):
+    if 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 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:
+    if expected_dim is not None:
         dim = expected_dim
     elif pa.types.is_fixed_size_list(vec_arr.type):
         dim = vec_arr.type.list_size
@@ -4256,16 +4230,15 @@ def _handle_bad_vector_column(
         if dim is None:
             return data
 
-    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),
-        )
+    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()
 
@@ -4300,10 +4273,7 @@ def _handle_bad_vector_column(
                 raise ValueError(
                     "`fill_value` must not be None if `on_bad_vectors` is 'fill'"
                 )
-            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)
+            vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value)
         else:
             raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}")
 
@@ -4355,60 +4325,17 @@ 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):
-        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):
+        values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks])
+    else:
+        values = arr.flatten()
+    if 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()))
-    elif pa.types.is_floating(values.type):
-        values_has_nan = pc.is_nan(values)
     else:
-        return pa.array([False] * len(arr))
+        values_has_nan = pc.is_nan(values)
     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())
@@ -4423,16 +4350,6 @@ 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:
@@ -4524,16 +4441,6 @@ 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_embeddings.py

@@ -1,8 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
+import ntpath
 import os
 import pickle
+import sys
+from types import ModuleType
 from typing import List, Optional, Union
 from unittest.mock import MagicMock, patch
 
@@ -522,6 +525,59 @@ def test_embedding_function_safe_model_dump(embedding_type):
             )
 
 
+def test_instructor_embedding_supports_huggingface_hub_without_cached_download(
+    tmp_path, monkeypatch
+):
+    from lancedb.embeddings.instructor import InstructorEmbeddingFunction
+
+    hub_download = MagicMock(return_value="/cache/1_Pooling/config.json")
+    huggingface_hub = ModuleType("huggingface_hub")
+    huggingface_hub.hf_hub_download = hub_download
+    torch = ModuleType("torch")
+    monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub)
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    monkeypatch.delitem(sys.modules, "InstructorEmbedding", raising=False)
+    monkeypatch.syspath_prepend(str(tmp_path))
+
+    (tmp_path / "InstructorEmbedding.py").write_text(
+        "from huggingface_hub import cached_download\n\n"
+        "class INSTRUCTOR:\n"
+        "    def __init__(self, name):\n"
+        "        self.name = name\n"
+    )
+
+    embedding = InstructorEmbeddingFunction.create(show_progress_bar=False)
+    instructor_model = embedding.get_model()
+
+    assert instructor_model.name == "hkunlp/instructor-base"
+    assert not hasattr(huggingface_hub, "cached_download")
+
+    instructor_embedding = sys.modules["InstructorEmbedding"]
+    path = instructor_embedding.cached_download(
+        url=(
+            "https://huggingface.co/hkunlp/instructor-base/resolve/abc123/"
+            "1_Pooling/config.json"
+        ),
+        cache_dir="/cache",
+        force_filename=ntpath.join("1_Pooling", "config.json"),
+        library_name="sentence-transformers",
+        library_version="2.2.2",
+        use_auth_token="token",
+    )
+
+    assert path == "/cache/1_Pooling/config.json"
+    hub_download.assert_called_once_with(
+        repo_id="hkunlp/instructor-base",
+        filename="1_Pooling/config.json",
+        revision="abc123",
+        local_dir="/cache",
+        library_name="sentence-transformers",
+        library_version="2.2.2",
+        user_agent=None,
+        token="token",
+    )
+
+
 @patch("time.sleep")
 def test_retry(mock_sleep):
     test_function = MagicMock(side_effect=[Exception] * 9 + ["result"])

python/python/tests/test_table.py

@@ -1710,19 +1710,6 @@ 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(
         [

python/python/tests/test_util.py

@@ -400,83 +400,6 @@ 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(

rust/lancedb/src/table/add_data.rs

@@ -258,7 +258,6 @@ 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};
@@ -886,63 +885,6 @@ 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, ListArray};
+use arrow_array::{Array, FixedSizeListArray};
 use arrow_schema::{DataType, Field, FieldRef};
 use datafusion_common::config::ConfigOptions;
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
@@ -19,20 +19,16 @@ 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 or multivector column.
+/// Returns true if the field is a vector column: FixedSizeList<Float16/32/64>.
 fn is_vector_field(field: &Field) -> bool {
-    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,
-        }
+    if let DataType::FixedSizeList(child, _) = field.data_type() {
+        matches!(
+            child.data_type(),
+            DataType::Float16 | DataType::Float32 | DataType::Float64
+        )
+    } else {
+        false
     }
-
-    is_vector_data_type(field.data_type())
 }
 
 /// Wraps the input plan with a projection that checks vector columns for NaN values.
@@ -73,8 +69,8 @@ pub fn reject_nan_vectors(input: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn Execu
     Ok(Arc::new(projection))
 }
 
-/// A scalar UDF that passes through vector arrays unchanged, but errors if any
-/// float values in the vector or multivector are NaN.
+/// A scalar UDF that passes through FixedSizeList arrays unchanged, but errors
+/// if any float values in the list are NaN.
 #[derive(Debug, Hash, PartialEq, Eq)]
 struct RejectNanUdf {
     signature: Signature,
@@ -117,54 +113,44 @@ impl ScalarUDFImpl for RejectNanUdf {
 }
 
 fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> {
-    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 => Ok(array
-                .as_any()
-                .downcast_ref::<arrow_array::Float32Array>()
-                .unwrap()
-                .iter()
-                .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()))),
-            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}"
-            ))),
-        }
-    }
+    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.
-    if contains_nan(array)? {
+    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
+                .as_any()
+                .downcast_ref::<arrow_array::Float16Array>()
+                .unwrap()
+                .iter()
+                .any(|v| v.is_some_and(|v| v.is_nan())),
+            DataType::Float32 => row
+                .as_any()
+                .downcast_ref::<arrow_array::Float32Array>()
+                .unwrap()
+                .iter()
+                .any(|v| v.is_some_and(|v| v.is_nan())),
+            DataType::Float64 => row
+                .as_any()
+                .downcast_ref::<arrow_array::Float64Array>()
+                .unwrap()
+                .iter()
+                .any(|v| v.is_some_and(|v| v.is_nan())),
+            _ => false,
+        }
+    });
+
+    if has_nan {
         return Err(datafusion_common::DataFusionError::ArrowError(
             Box::new(arrow_schema::ArrowError::ComputeError(
                 "Vector column contains NaN values".to_string(),
@@ -180,7 +166,6 @@ 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() {
@@ -206,46 +191,6 @@ 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.
@@ -309,15 +254,6 @@ 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",