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)