diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 44d16211c..c3f521509 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): @@ -687,7 +689,7 @@ class RemoteTable(Table): vector_column_name: Optional[str] = None, query_type="auto", fts_columns: Optional[Union[str, List[str]]] = None, - fast_search: bool = False, + fast_search: Optional[bool] = None, ) -> LanceVectorQueryBuilder: """Create a search query to find the nearest neighbors of the given query vector. We currently support @@ -738,7 +740,7 @@ class RemoteTable(Table): Skip a flat search of unindexed data. This may improve search performance but search results will not include unindexed data. - - *default False*. + - *default None*, which leaves fast search disabled. Returns ------- @@ -753,12 +755,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" + 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, query_type, vector_column_name=vector_column_name, - fts_columns=fts_columns, + fts_columns=fts_columns or [], fast_search=fast_search, ) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 8f534f678..d846a0120 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -17,10 +17,11 @@ from packaging.version import Version import lancedb from lancedb.conftest import MockTextEmbeddingFunction -from lancedb.query import ColumnOrdering +from lancedb.query import ColumnOrdering, MatchQuery from lancedb.remote import ClientConfig from lancedb.remote.errors import HttpError, RetryError import pytest +import numpy as np import pyarrow as pa @@ -1430,14 +1431,39 @@ 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=(("vector", 3),), +): 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": []}}') + 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")) body = request.rfile.read(content_len) @@ -1462,6 +1488,65 @@ def query_test_table(query_handler, *, server_version=Version("0.1.0")): yield table +@pytest.mark.parametrize( + "query,search_kwargs,expected_column", + [ + pytest.param([0.5] * 8, {}, "vector", id="vector-8"), + pytest.param([0.5] * 4, {}, "vec2", id="vector-4"), + pytest.param(np.array([0.5] * 4), {}, "vec2", id="numpy-vector"), + pytest.param([[0.5] * 4, [0.2] * 4], {}, "vec2", id="batch-vectors"), + pytest.param( + [0.5] * 4, {"vector_column_name": "vec2"}, "vec2", id="explicit-column" + ), + pytest.param(None, {}, None, id="scan"), + pytest.param("hello", {"query_type": "fts"}, None, id="fts"), + pytest.param( + "hello", {"query_type": "fts", "fts_columns": "text"}, None, id="fts-column" + ), + pytest.param(MatchQuery("hello", "text"), {}, None, id="structured-fts"), + ], +) +def test_query_object_sync_matches_local(mem_db, query, search_kwargs, expected_column): + vector_columns = (("vector", 8), ("vec2", 4)) + schema = pa.schema( + [pa.field(name, pa.list_(pa.float32(), dim)) for name, dim in vector_columns] + ) + local = mem_db.create_table("test", schema=schema) + local_query = local.search(query, **search_kwargs).limit(2).to_query_object() + assert local_query.vector_column == expected_column + assert local_query.fast_search is None + + with query_test_table(None, vector_columns=vector_columns) as remote: + remote_query = remote.search(query, **search_kwargs).limit(2).to_query_object() + assert remote_query == local_query + + +@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 () + with query_test_table(None, vector_columns=vector_columns) as table: + query_obj = table.search( + query, query_type=query_type, fast_search=fast_search + ).to_query_object() + assert query_obj.fast_search is fast_search + + +def test_query_sync_rejects_ambiguous_vector_column(mem_db): + vector_columns = (("vector", 4), ("vec2", 4)) + schema = pa.schema( + [pa.field(name, pa.list_(pa.float32(), dim)) for name, dim in vector_columns] + ) + local = mem_db.create_table("test", schema=schema) + with pytest.raises(ValueError, match="Candidates:.*vector.*vec2") as local_error: + local.search([0.5] * 4) + + with query_test_table(None, vector_columns=vector_columns) as remote: + with pytest.raises(ValueError) as remote_error: + remote.search([0.5] * 4) + assert str(remote_error.value) == str(local_error.value) + + def test_head(): def handler(body): assert body == { @@ -1488,6 +1573,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, @@ -1782,6 +1868,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, @@ -1793,7 +1880,7 @@ 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: + with query_test_table(handler, vector_columns=(("vector", 10),)) as table: embedding_func = MockTextEmbeddingFunction() embedding_config = MagicMock() embedding_config.function = embedding_func