fix(python): accept tuple and Arrow vector queries (#3972)

## Summary

Sync `Table.search()` declares tuple, `pyarrow.Array`, and
`pyarrow.ChunkedArray` query vectors as supported, and the async query
path already accepts them. The sync query builder only classified lists
and NumPy arrays as vectors, however, so these inputs were routed
through embedding/FTS resolution and then rejected. Arrow and tuple
inputs also did not contribute a query dimension when selecting among
multiple vector columns.

This change:

- recognizes the supported tuple and Arrow containers as sync vector
queries for both `auto` and explicit `vector` query types;
- converts them to the existing internal `float32` NumPy representation
at the sync builder boundary;
- infers dimensions from tuple, primitive Arrow, and list-valued Arrow
queries so schemas with differently sized vector columns select the
matching column.

The regression test uses two vector columns with dimensions 2 and 3. It
verifies tuple, Arrow array, and chunked Arrow queries through both sync
query modes and confirms the dimension-2 column is selected.

## Test plan

- `pytest
python/tests/test_table.py::test_sync_search_accepts_supported_vector_containers
-q` (6 passed)
- `pytest python/tests/test_util.py python/tests/test_query.py
python/tests/test_table.py -q -m "not slow and not s3_test"` (308
passed, 1 skipped, 1 deselected)
- `pytest python/tests -q -m "not slow and not s3_test"
--ignore=python/tests/docs` (981 passed, 25 skipped, 65 deselected)
- `ruff format --check python`
- `ruff check python`
- `cargo fmt --all --check`

Co-authored-by: Xuanwo <github@xuanwo.io>

python/python/lancedb/query.py

@@ -82,6 +82,7 @@ if TYPE_CHECKING:
 
 T = TypeVar("T", bound="LanceModel")
 AnalyzePlanDistributedMetrics = Literal["aggregate", "per_worker", "full"]
+_SYNC_VECTOR_QUERY_TYPES = (list, tuple, np.ndarray, pa.Array, pa.ChunkedArray)
 
 # Number of rows a hybrid query returns when no limit was set on it. This
 # mirrors the default the Rust query builder applies to its sub-queries.
@@ -1024,10 +1025,12 @@ class LanceQueryBuilder(ABC):
                 fast_search=fast_search,
             )
 
-        if isinstance(query, list):
+        if isinstance(query, (list, tuple)):
             query = np.array(query, dtype=np.float32)
         elif isinstance(query, np.ndarray):
             query = query.astype(np.float32)
+        elif isinstance(query, (pa.Array, pa.ChunkedArray)):
+            query = np.array(query.to_pylist(), dtype=np.float32)
         else:
             raise TypeError(f"Unsupported query type: {type(query)}")
 
@@ -1049,7 +1052,7 @@ class LanceQueryBuilder(ABC):
             query = cls._query_to_vector(table, query, vector_column_name)
             return query, query_type
         elif query_type == "auto":
-            if isinstance(query, (list, np.ndarray)):
+            if isinstance(query, _SYNC_VECTOR_QUERY_TYPES):
                 return query, "vector"
             else:
                 conf = table.embedding_functions.get(vector_column_name)
@@ -1065,7 +1068,7 @@ class LanceQueryBuilder(ABC):
 
     @classmethod
     def _query_to_vector(cls, table, query, vector_column_name):
-        if isinstance(query, (list, np.ndarray)):
+        if isinstance(query, _SYNC_VECTOR_QUERY_TYPES):
             return query
         conf = table.embedding_functions.get(vector_column_name)
         if conf is not None:

python/python/lancedb/util.py

@@ -309,7 +309,16 @@ def _query_vector_dim(query: Optional[Any]) -> Optional[int]:
         if query.ndim == 0:
             return None
         return query.shape[-1]
-    if isinstance(query, list) and query:
+    if isinstance(query, (pa.Array, pa.ChunkedArray)):
+        if len(query) == 0:
+            return None
+        if pa.types.is_fixed_size_list(query.type):
+            return query.type.list_size
+        if pa.types.is_list(query.type) or pa.types.is_large_list(query.type):
+            first = query[0].as_py()
+            return len(first) if first is not None else None
+        return len(query)
+    if isinstance(query, (list, tuple)) and query:
         first = query[0]
         if isinstance(first, (list, tuple, np.ndarray)):
             return len(first)

python/python/tests/test_table.py

@@ -3697,6 +3697,42 @@ def test_multiple_vector_columns(mem_db: DBConnection):
     assert result1["text"][0] != result2["text"][0]
 
 
+@pytest.mark.parametrize(
+    "query",
+    [
+        (0.0, 0.0),
+        pa.array([0.0, 0.0]),
+        pa.chunked_array([[0.0], [0.0]]),
+    ],
+)
+@pytest.mark.parametrize("query_type", ["auto", "vector"])
+def test_sync_search_accepts_supported_vector_containers(
+    mem_db: DBConnection, query, query_type
+):
+    schema = pa.schema(
+        [
+            pa.field("id", pa.int64()),
+            pa.field("small", pa.list_(pa.float32(), 2)),
+            pa.field("large", pa.list_(pa.float32(), 3)),
+        ]
+    )
+    table = mem_db.create_table(
+        "supported_vector_containers",
+        data=pa.table(
+            {
+                "id": [1, 2],
+                "small": [[0.0, 0.0], [1.0, 1.0]],
+                "large": [[1.0, 1.0, 1.0], [0.0, 0.0, 0.0]],
+            },
+            schema=schema,
+        ),
+    )
+
+    result = table.search(query, query_type=query_type).limit(1).to_arrow()
+
+    assert result["id"].to_pylist() == [1]
+
+
 def test_create_scalar_index(mem_db: DBConnection):
     vec_array = pa.array(
         [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]], pa.list_(pa.float32(), 2)