diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 19a443950..2522abe6c 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -736,6 +736,12 @@ class RemoteTable(Table): - If the table has multiple vector columns then the *vector_column_name* needs to be specified. Otherwise, an error is raised. + query_type: str, default "auto" + The type of search: "auto", "vector", "fts", or "hybrid". + With "auto", a string query uses the embedding function for the + inferred or specified vector column if available, otherwise it uses + full-text search. + fast_search: bool, optional Skip a flat search of unindexed data. This may improve search performance but search results will not include unindexed data. diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 020f971d8..32d1e6e9c 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -19,7 +19,12 @@ import lancedb import numpy as np from lancedb.conftest import MockNonNormTextEmbeddingFunction, MockTextEmbeddingFunction from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry -from lancedb.query import AsyncQuery, ColumnOrdering, MatchQuery +from lancedb.query import ( + AsyncQuery, + ColumnOrdering, + LanceVectorQueryBuilder, + MatchQuery, +) from lancedb.remote import ClientConfig from lancedb.remote.errors import HttpError, RetryError import pytest @@ -1607,34 +1612,46 @@ def query_test_table( query_handler, *, server_version=Version("0.1.0"), - vector_columns=(("vector", 3),), + vector_columns=None, + embedding_functions=None, ): + if vector_columns is None: + vector_columns = {"vector": 3} + schema = { + "fields": [ + {"name": "text", "type": {"type": "string"}, "nullable": False}, + *[ + { + "name": name, + "type": { + "type": "fixed_size_list", + "fields": [ + { + "name": "item", + "type": {"type": "float"}, + "nullable": True, + } + ], + "length": dim, + }, + "nullable": False, + } + for name, dim in vector_columns.items() + ], + ] + } + if embedding_functions: + metadata = EmbeddingFunctionRegistry.get_instance().get_table_metadata( + embedding_functions + ) + schema["metadata"] = {key: value.decode() for key, value in metadata.items()} + def handler(request): if request.path == "/v1/table/test/describe/": request.send_response(200) request.send_header("Content-Type", "application/json") request.send_header("phalanx-version", str(server_version)) request.end_headers() - schema = { - "fields": [ - { - "name": name, - "type": { - "type": "fixed_size_list", - "fields": [ - { - "name": "item", - "type": {"type": "float"}, - "nullable": True, - } - ], - "length": dim, - }, - "nullable": True, - } - for name, dim in vector_columns - ] - } request.wfile.write(json.dumps({"version": 1, "schema": schema}).encode()) elif request.path == "/v1/table/test/query/": content_len = int(request.headers.get("Content-Length")) @@ -1679,9 +1696,12 @@ def query_test_table( ], ) def test_query_object_sync_matches_local(mem_db, query, search_kwargs, expected_column): - vector_columns = (("vector", 8), ("vec2", 4)) + vector_columns = {"vector": 8, "vec2": 4} schema = pa.schema( - [pa.field(name, pa.list_(pa.float32(), dim)) for name, dim in vector_columns] + [ + pa.field(name, pa.list_(pa.float32(), dim)) + for name, dim in vector_columns.items() + ] ) local = mem_db.create_table("test", schema=schema) local_query = local.search(query, **search_kwargs).limit(2).to_query_object() @@ -1696,7 +1716,7 @@ def test_query_object_sync_matches_local(mem_db, query, search_kwargs, expected_ @pytest.mark.parametrize("fast_search", [None, False, True]) @pytest.mark.parametrize("query,query_type", [([1, 2, 3], "vector"), ("hello", "fts")]) def test_query_object_sync_preserves_fast_search(fast_search, query, query_type): - vector_columns = (("vector", 3),) if query_type == "vector" else () + vector_columns = {"vector": 3} if query_type == "vector" else {} with query_test_table(None, vector_columns=vector_columns) as table: query_obj = table.search( query, query_type=query_type, fast_search=fast_search @@ -1705,9 +1725,12 @@ def test_query_object_sync_preserves_fast_search(fast_search, query, query_type) def test_query_sync_rejects_ambiguous_vector_column(mem_db): - vector_columns = (("vector", 4), ("vec2", 4)) + vector_columns = {"vector": 4, "vec2": 4} schema = pa.schema( - [pa.field(name, pa.list_(pa.float32(), dim)) for name, dim in vector_columns] + [ + pa.field(name, pa.list_(pa.float32(), dim)) + for name, dim in vector_columns.items() + ] ) local = mem_db.create_table("test", schema=schema) with pytest.raises(ValueError, match="Candidates:.*vector.*vec2") as local_error: @@ -1719,7 +1742,7 @@ def test_query_sync_rejects_ambiguous_vector_column(mem_db): assert str(remote_error.value) == str(local_error.value) -@pytest.mark.parametrize("vector_columns", [(), (("v1", 4), ("v2", 8))]) +@pytest.mark.parametrize("vector_columns", [{}, {"v1": 4, "v2": 8}]) @pytest.mark.parametrize("query_type", ["auto", "fts"]) def test_query_sync_fts_without_embeddings(vector_columns, query_type): seen = [] @@ -1743,14 +1766,14 @@ def test_query_sync_auto_with_embedding_function(): assert "full_text_query" not in body return pa.table({"id": [1]}) - with query_test_table(handler, vector_columns=(("vector", 10),)) as table: - table.embedding_functions = { - "vector": EmbeddingFunctionConfig( - source_column="text", - vector_column="vector", - function=MockTextEmbeddingFunction(), - ) - } + config = EmbeddingFunctionConfig( + source_column="text", + vector_column="vector", + function=MockTextEmbeddingFunction.create(), + ) + with query_test_table( + handler, vector_columns={"vector": 10}, embedding_functions=[config] + ) as table: query = table.search("hello") assert query.to_query_object().vector_column == "vector" assert query.to_list() == [{"id": 1}] @@ -1797,7 +1820,84 @@ def test_query_sync_minimal(): assert data == expected -def test_query_sync_empty_query(): +@pytest.mark.parametrize("query_type", ["auto", "vector", "hybrid"]) +@pytest.mark.parametrize("vector_column_name", [None, "embedding"]) +def test_query_sync_text_embedding(query_type, vector_column_name): + embedding_func = MockTextEmbeddingFunction.create() + config = EmbeddingFunctionConfig( + source_column="text", vector_column="embedding", function=embedding_func + ) + query = "quick brown fox" + expected_vector = embedding_func.compute_query_embeddings(query)[0] + queries = [] + + def handler(body): + queries.append(body) + if "full_text_query" in body: + assert query_type == "hybrid" + assert body["full_text_query"]["query"] == query + return pa.table({"id": [1], "_rowid": [1], "_score": [1.0]}) + assert body["vector_column"] == "embedding" + assert body["vector"] == pytest.approx(expected_vector) + return pa.table({"id": [1], "_rowid": [1], "_distance": [0.0]}) + + with query_test_table( + handler, + vector_columns={ + "embedding": embedding_func.ndims(), + **({"other": embedding_func.ndims()} if vector_column_name else {}), + }, + embedding_functions=[config], + ) as table: + assert list(table.embedding_functions) == ["embedding"] + builder = table.search( + query, query_type=query_type, vector_column_name=vector_column_name + ) + if query_type != "hybrid": + assert isinstance(builder, LanceVectorQueryBuilder) + results = builder.to_list() + assert results[0]["id"] == 1 + if query_type != "hybrid": + assert results[0]["_distance"] == 0.0 + assert len(queries) == (2 if query_type == "hybrid" else 1) + + +@pytest.mark.parametrize("query_type", ["auto", "vector", "hybrid"]) +def test_query_sync_ambiguous_vector_columns(query_type): + config = EmbeddingFunctionConfig( + source_column="text", + vector_column="v1", + function=MockTextEmbeddingFunction.create(), + ) + queries = [] + + def handler(body): + queries.append(body) + return pa.table({"id": [1]}) + + with query_test_table( + handler, + vector_columns={"v1": 10, "v2": 10}, + embedding_functions=[config], + ) as table: + with pytest.raises(ValueError, match=r"Candidates: \['v1', 'v2'\]"): + table.search("puppy", query_type=query_type).to_list() + assert queries == [] + + +@pytest.mark.parametrize("vector_column, dim", [("v1", 3), ("v2", 4)]) +def test_query_sync_infers_vector_column_by_dimension(vector_column, dim): + def handler(body): + assert body["vector_column"] == vector_column + assert body["vector"] == [1.0] * dim + return pa.table({"id": [1]}) + + with query_test_table(handler, vector_columns={"v1": 3, "v2": 4}) as table: + assert table.search([1.0] * dim).to_list() == [{"id": 1}] + + +@pytest.mark.parametrize("vector_columns", [{}, {"v1": 3, "v2": 3}]) +def test_query_sync_empty_query(vector_columns): def handler(body): assert body == { "k": 10, @@ -1810,7 +1910,7 @@ def test_query_sync_empty_query(): return pa.table({"id": [1, 2, 3]}) - with query_test_table(handler) as table: + with query_test_table(handler, vector_columns=vector_columns) as table: data = table.search(None).where("true").select(["id"]).limit(10).to_list() expected = [{"id": 1}, {"id": 2}, {"id": 3}] assert data == expected @@ -2060,7 +2160,9 @@ def test_query_sync_batch_queries(server_version): assert results == [{"id": 1, "query_index": 0}, {"id": 1, "query_index": 1}] -def test_query_sync_fts(): +@pytest.mark.parametrize("vector_columns", [{}, {"vector": 3}, {"v1": 3, "v2": 3}]) +@pytest.mark.parametrize("query_type", ["auto", "fts"]) +def test_query_sync_fts(vector_columns, query_type): def handler(body): assert body == { "full_text_query": { @@ -2075,8 +2177,8 @@ def test_query_sync_fts(): return pa.table({"id": [1, 2, 3]}) - with query_test_table(handler) as table: - (table.search("puppy", query_type="fts").to_list()) + with query_test_table(handler, vector_columns=vector_columns) as table: + (table.search("puppy", query_type=query_type).to_list()) def handler(body): assert body == { @@ -2103,9 +2205,11 @@ def test_query_sync_fts(): return pa.table({"id": [1, 2, 3]}) - with query_test_table(handler) as table: + with query_test_table(handler, vector_columns=vector_columns) as table: ( - table.search("puppy", query_type="fts", fts_columns=["name", "description"]) + table.search( + "puppy", query_type=query_type, fts_columns=["name", "description"] + ) .with_row_id(True) .limit(42) .to_list() @@ -2143,7 +2247,9 @@ def test_query_sync_fts_document_granularity(): } ) - with query_test_table(handler, server_version=Version("0.6.0")) as table: + with query_test_table( + handler, server_version=Version("0.6.0"), vector_columns={} + ) as table: result = table.search( MatchQuery( "alpha", @@ -2190,15 +2296,14 @@ def test_query_sync_hybrid(): } return pa.table({"_rowid": [1, 2, 3], "_distance": [0.1, 0.2, 0.3]}) - with query_test_table(handler, vector_columns=(("vector", 10),)) as table: - embedding_func = MockTextEmbeddingFunction() - embedding_config = MagicMock() - embedding_config.function = embedding_func - - embedding_funcs = MagicMock() - embedding_funcs.get = MagicMock(return_value=embedding_config) - table.embedding_functions = embedding_funcs - + config = EmbeddingFunctionConfig( + source_column="text", + vector_column="vector", + function=MockTextEmbeddingFunction.create(), + ) + with query_test_table( + handler, vector_columns={"vector": 10}, embedding_functions=[config] + ) as table: (table.search("puppy", query_type="hybrid").limit(42).to_list())