diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 44d16211c..57d01e4c9 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -63,6 +63,7 @@ from lancedb.table import _normalize_progress from ..query import ( AnalyzePlanDistributedMetrics, DocumentGranularity, + FullTextQuery, LanceQueryBuilder, LanceTakeQueryBuilder, LanceVectorQueryBuilder, @@ -77,6 +78,7 @@ from ..table import ( Tags, ) from ..types import BaseTokenizerType +from ..util import infer_vector_column_name class RemoteTable(Table): @@ -734,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. @@ -753,6 +761,21 @@ class RemoteTable(Table): - and also the "_distance" column which is the distance between the query vector and the returned vector. """ + if isinstance(query, FullTextQuery): + query_type = "fts" + elif ( + query_type == "auto" + and isinstance(query, str) + and not self.embedding_functions + ): + query_type = "fts" + vector_column_name = infer_vector_column_name( + schema=self.schema, + query_type=query_type, + query=query, + vector_column_name=vector_column_name, + ) + return LanceQueryBuilder.create( self, query, diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 36c318591..787a39e1d 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -19,7 +19,7 @@ import lancedb import numpy as np from lancedb.conftest import MockNonNormTextEmbeddingFunction, MockTextEmbeddingFunction from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry -from lancedb.query import AsyncQuery, ColumnOrdering +from lancedb.query import AsyncQuery, ColumnOrdering, LanceVectorQueryBuilder from lancedb.remote import ClientConfig from lancedb.remote.errors import HttpError, RetryError import pytest @@ -1603,14 +1603,51 @@ def test_checkpoint_lsm_sync(): @contextlib.contextmanager -def query_test_table(query_handler, *, server_version=Version("0.1.0")): +def query_test_table( + query_handler, + *, + server_version=Version("0.1.0"), + 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() - request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') + 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")) body = request.rfile.read(content_len) @@ -1661,6 +1698,7 @@ def test_query_sync_minimal(): "upper_bound": None, "ef": None, "vector": [1.0, 2.0, 3.0], + "vector_column": "vector", "nprobes": 20, "minimum_nprobes": 20, "maximum_nprobes": 20, @@ -1675,7 +1713,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, @@ -1688,7 +1803,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 @@ -1938,7 +2053,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": { @@ -1953,8 +2070,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 == { @@ -1981,9 +2098,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() @@ -2021,7 +2140,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", @@ -2056,6 +2177,7 @@ def test_query_sync_hybrid(): "prefilter": True, "refine_factor": None, "vector": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "vector_column": "vector", "nprobes": 20, "minimum_nprobes": 20, "maximum_nprobes": 20, @@ -2067,15 +2189,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) 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())