rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-20 18:52:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(python): align remote and local search query objects

Browse Source
This commit is contained in:
Gatefixer
2026-10-03 02:19:36 +00:00
parent b0c0d8ee92
commit b44c5257bc
2 changed files with 105 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/remote/table.py
+14 -3
View File
@@ -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,
)
python/python/tests/test_remote_db.py
+91 -4
View File
@@ -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