mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-12 23:02:28 +00:00
fix(python): align remote and local search query objects (#4413)
Python sync remote searches default `fast_search` to `False` and omitted full-text search columns to `None`, while local searches leave fast search unset and use `[]`. This makes otherwise identical `to_query_object()` results depend on the backend. Align these omitted defaults in `RemoteTable.search` while preserving explicit fast-search settings. Keep the shared vector-column inference and automatic FTS fallback provided by the base branch. Regression tests compare complete local and remote query objects for both reported vector dimensions, NumPy and batch vectors, explicit columns, scans, and full-text queries. They also cover matching ambiguity errors, explicit fast-search values, automatic FTS without embeddings, and automatic vector search with embedding metadata, using the shared remote fixture. Validation after merging `main` at `bbc4aebe`: - `cd python && uv run --no-sync --extra tests pytest python/tests/test_remote_db.py -q`: 128 passed. - `cd python && uv run --no-sync --extra tests pytest python/tests/test_query.py python/tests/test_table.py -q -k 'query_serialization or async_search_without_query or none_query or nprobes_zero or empty_vector'`: 27 passed. - `ruff format .` and `ruff check .` from the repository root (Ruff 0.9.9): passed. - `git diff --check`: passed. Fixes #4391 <!-- lance-gatekeeper-fix:v1 agent=7cad560371d7aa64fb2511e510dd1bfd generation=1 --> --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
This commit is contained in:
@@ -686,7 +686,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
|
||||
@@ -743,7 +743,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
|
||||
-------
|
||||
@@ -778,7 +778,7 @@ class RemoteTable(Table):
|
||||
query,
|
||||
query_type,
|
||||
vector_column_name=vector_column_name,
|
||||
fts_columns=fts_columns,
|
||||
fts_columns=fts_columns or [],
|
||||
fast_search=fast_search,
|
||||
)
|
||||
|
||||
|
||||
@@ -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, LanceVectorQueryBuilder
|
||||
from lancedb.query import (
|
||||
AsyncQuery,
|
||||
ColumnOrdering,
|
||||
LanceVectorQueryBuilder,
|
||||
MatchQuery,
|
||||
)
|
||||
from lancedb.remote import ClientConfig
|
||||
from lancedb.remote.errors import HttpError, RetryError
|
||||
import pytest
|
||||
@@ -1831,6 +1836,108 @@ def query_test_table(
|
||||
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.items()
|
||||
]
|
||||
)
|
||||
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.items()
|
||||
]
|
||||
)
|
||||
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)
|
||||
|
||||
|
||||
@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 = []
|
||||
|
||||
def handler(body):
|
||||
seen.append(body)
|
||||
assert body["full_text_query"] == {"query": "hello", "columns": []}
|
||||
assert body["vector"] == []
|
||||
assert "vector_column" not in body
|
||||
return pa.table({"id": [1]})
|
||||
|
||||
with query_test_table(handler, vector_columns=vector_columns) as table:
|
||||
assert table.search("hello", query_type=query_type).to_list() == [{"id": 1}]
|
||||
assert len(seen) == 1
|
||||
|
||||
|
||||
def test_query_sync_auto_with_embedding_function():
|
||||
def handler(body):
|
||||
assert body["vector_column"] == "vector"
|
||||
assert body["vector"] == [0.0] * 10
|
||||
assert "full_text_query" not in body
|
||||
return pa.table({"id": [1]})
|
||||
|
||||
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}]
|
||||
|
||||
|
||||
def test_head():
|
||||
def handler(body):
|
||||
assert body == {
|
||||
|
||||
Reference in New Issue
Block a user