From 6d92307fc34b81adc0c040da367c6f9b64f0aeba Mon Sep 17 00:00:00 2001 From: Zhewen Tan <1090179959@qq.com> Date: Fri, 23 Oct 2026 01:10:15 +0800 Subject: [PATCH] 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 --- python/python/lancedb/query.py | 9 +++++--- python/python/lancedb/util.py | 11 +++++++++- python/python/tests/test_table.py | 36 +++++++++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 4 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index c427820e8..53f56a2c0 100644 --- a/python/python/lancedb/query.py +++ b/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: diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index d4533887b..3babc3872 100644 --- a/python/python/lancedb/util.py +++ b/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) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 095f0b911..7a4f5981c 100644 --- a/python/python/tests/test_table.py +++ b/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)