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

fix(python): infer vector columns for sync remote search (#4409)

Browse Source 
Sync remote string searches ignored the table's embedding function when
`vector_column_name` was omitted. `auto` silently fell back to full-text
search, while `vector` and `hybrid` failed with `No embedding function
for None`.

Resolve vector columns from the schema before creating the query
builder, using the same inference helper as local search. Automatic
string queries without embedding functions select full-text search
first, so text-only tables and tables with multiple unconfigured vector
columns remain searchable. Vector and hybrid searches still reject
ambiguous columns, infer numeric queries by dimension, and honor
explicit column selection. Structured full-text queries and empty
queries retain their existing behavior.

Regression coverage uses embedding metadata returned by a mocked remote
server rather than a mocked embedding-function lookup. It also checks
automatic and explicit full-text queries on schemas with zero, one, or
multiple vector columns and no embedding functions.

Validation after merging `main` at `87ecfeac`:

- Confirmed the original embedding and ambiguity failures before the
fix, and reproduced the automatic FTS inference failures before
restoring the fallback.
- Remote, query, hybrid, and embedding test files through the uv-managed
environment: 240 passed, 4 skipped, 9 deselected (`not slow and not
s3_test`).
- Python extension built with `uv run --extra tests --extra dev maturin
develop --extras tests,dev`.
- `ruff check .` and `ruff format .` from the repository root, using
CI's Ruff version 0.9.9.
- `cargo fmt --all`.

Fixes #4377

<!-- lance-gatekeeper-fix:v1 agent=c167c6c89d857342e5d7415b08142589
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
This commit is contained in:
lancedb-gatefixer[bot]
2026-10-06 19:59:21 -07:00
committed by GitHub
parent a8baecb306
commit bbc4aebe55
2 changed files with 164 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/remote/table.py
+23
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):
@@ -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,
python/python/tests/test_remote_db.py
+141 -20
View File
@@ -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())