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():