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>
This commit is contained in:
Zhewen Tan
2026-10-23 01:10:15 +08:00
committed by GitHub
co-authored by Xuanwo
parent 548a190321
commit 6d92307fc3
3 changed files with 52 additions and 4 deletions
+36
View File
@@ -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)