From 3575f972ac4ca599ba92fe5ea8081a1f3f4f8953 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=99=88=E5=BF=97=E8=B0=A6?= <89645338+simpleqt@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:49:31 +0800 Subject: [PATCH] fix(python): accept tuples and pyarrow arrays as hybrid search vectors (#4662) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The synchronous hybrid query builder rejected vector containers that the rest of the Python API accepts: ```python table.search(query_type="hybrid").vector((1.0, 2.0)).text("hello").to_pandas() # ValueError table.search(query_type="hybrid").vector(pa.array([1.0, 2.0], pa.float32()))... # ValueError: Vector query must be either a string or a vector ``` `lancedb.common.VEC` documents `list[float] | tuple[float, ...] | np.ndarray | pa.Array`, and the same two containers work for `table.search(vector)` and for the async hybrid API (`AsyncQuery.nearest_to(...).nearest_to_text(...)`). `LanceHybridQueryBuilder._validate_query` only listed `str | list | np.ndarray`, so a tuple or a pyarrow array was refused before reaching the query planner. This normalises the vector exactly the way `LanceQueryBuilder.create` already does — list/tuple to a `np.float32` array, `np.ndarray` through `astype(float32)`, `pa.Array`/`pa.ChunkedArray` through `np.array(..., dtype=float32)`, strings passed through — and keeps the existing `ValueError` for anything else. `python/python/tests/test_hybrid_query.py::test_hybrid_query_accepts_every_vector_container` is parametrised over tuple, pyarrow array and pyarrow chunked array, and asserts the returned `_rowid`s match the `list` form: 3 failed on `main` (`ValueError` from `query.py:2452`), 3 passed with the change. Neighbouring suites (`test_hybrid_query.py`, `test_query.py`, `test_rerankers.py`) are 220 passed / 9 skipped, and `test_table.py`, `test_remote_db.py`, `test_fts.py`, `test_index.py`, `test_permutation.py` are 724 passed / 5 skipped. `ruff` at the version pinned in the repository's pre-commit config (0.9.9) reports no findings and `ruff format --check` is clean. One source file (`python/python/lancedb/query.py`), no behaviour change for containers that already worked. Prepared with AI assistance; every line was verified by running the cases above. --------- Co-authored-by: Xuanwo --- python/python/lancedb/query.py | 28 ++++++++++++--- python/python/tests/test_hybrid_query.py | 44 ++++++++++++++++++++++-- 2 files changed, 66 insertions(+), 6 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 644b6571f..b2093712e 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -1655,12 +1655,12 @@ class LanceQueryBuilder(ABC): self.to_query_object(), distributed_metrics=distributed_metrics ) - def vector(self, vector: Union[np.ndarray, list]) -> Self: + def vector(self, vector: Union[VEC, Tuple]) -> Self: """Set the vector to search for. Parameters ---------- - vector: np.ndarray or list + vector: list, tuple, np.ndarray, pa.Array, or pa.ChunkedArray The vector to search for. Returns @@ -2450,7 +2450,15 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): ) vector_query = vector if vector is not None else query - if not isinstance(vector_query, (str, list, np.ndarray)): + # Accept the same containers as the vector search path, which + # normalizes them in `LanceQueryBuilder.create`. + if isinstance(vector_query, (list, tuple)): + vector_query = np.array(vector_query, dtype=np.float32) + elif isinstance(vector_query, np.ndarray): + vector_query = vector_query.astype(np.float32) + elif isinstance(vector_query, (pa.Array, pa.ChunkedArray)): + vector_query = np.array(vector_query.to_pylist(), dtype=np.float32) + elif not isinstance(vector_query, str): raise ValueError("Vector query must be either a string or a vector") text_query = text or query @@ -2923,7 +2931,19 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): self._refine_factor = refine_factor return self - def vector(self, vector: Union[np.ndarray, list]) -> LanceHybridQueryBuilder: + def vector(self, vector: Union[VEC, Tuple]) -> LanceHybridQueryBuilder: + """Set the vector to search for. + + Parameters + ---------- + vector: list, tuple, np.ndarray, pa.Array, or pa.ChunkedArray + The vector to search for. + + Returns + ------- + LanceHybridQueryBuilder + The LanceHybridQueryBuilder object. + """ self._vector = vector return self diff --git a/python/python/tests/test_hybrid_query.py b/python/python/tests/test_hybrid_query.py index 19050093c..c432bac59 100644 --- a/python/python/tests/test_hybrid_query.py +++ b/python/python/tests/test_hybrid_query.py @@ -315,13 +315,53 @@ def test_hybrid_query_offset(sync_table: Table): assert offset_result["_rowid"].to_pylist() == full["_rowid"].to_pylist()[2:] +@pytest.mark.parametrize( + "query_vector", + [ + pytest.param((0.0, 0.4), id="tuple"), + pytest.param(pa.array([0.0, 0.4], type=pa.float32()), id="pa-array"), + pytest.param( + pa.chunked_array([[0.0], [0.4]], type=pa.float32()), id="pa-chunked-array" + ), + ], +) +def test_hybrid_query_accepts_every_vector_container(sync_table: Table, query_vector): + # Vector search and the async hybrid API both accept tuples and pyarrow + # arrays for the query vector (``VEC`` in ``lancedb/common.py``), so the + # sync hybrid builder must accept them too. + expected = ( + sync_table.search(query_type="hybrid") + .vector([0.0, 0.4]) + .text("dog") + .limit(2) + .with_row_id(True) + .to_arrow() + ) + + result = ( + sync_table.search(query_type="hybrid") + .vector(query_vector) + .text("dog") + .limit(2) + .with_row_id(True) + .to_arrow() + ) + + assert result["_rowid"].to_pylist() == expected["_rowid"].to_pylist() + + def test_hybrid_query_default_limit(sync_table: Table): # The vector and FTS legs match disjoint rows, so fusing their results # yields more rows than the default limit. new_rows = [] + # Distinct distances and document lengths avoid ties in either leg, + # so comparing offset windows does not depend on tie ordering. for i in range(20): - new_rows.append({"text": "close_vec", "vector": [0.1, 0.1]}) - new_rows.append({"text": "dog", "vector": [50.0 + i, 50.0 + i]}) + coordinate = 0.11 + i * 0.001 + new_rows.append({"text": "close_vec", "vector": [coordinate, coordinate]}) + new_rows.append( + {"text": "dog " + "padding " * (i + 1), "vector": [50.0 + i, 50.0 + i]} + ) sync_table.add(new_rows) def query():