mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user