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

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

Browse Source
This commit is contained in:
Gatefixer
2026-10-03 01:56:05 +00:00
parent b0c0d8ee92
commit 35c59b2c6e
2 changed files with 154 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/remote/table.py
+17
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,15 @@ 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,
python/python/tests/test_remote_db.py
+137 -18
View File
@@ -17,7 +17,8 @@ from packaging.version import Version
import lancedb
from lancedb.conftest import MockTextEmbeddingFunction
from lancedb.query import ColumnOrdering
from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry
from lancedb.query import ColumnOrdering, LanceVectorQueryBuilder
from lancedb.remote import ClientConfig
from lancedb.remote.errors import HttpError, RetryError
import pytest
@@ -1430,14 +1431,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)
@@ -1488,6 +1526,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,
@@ -1502,7 +1541,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,
@@ -1515,7 +1631,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
@@ -1664,7 +1780,8 @@ 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", [{}, {"v1": 3, "v2": 3}])
def test_query_sync_fts(vector_columns):
def handler(body):
assert body == {
"full_text_query": {
@@ -1679,7 +1796,7 @@ 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").to_list())
def handler(body):
@@ -1707,7 +1824,7 @@ 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"])
.with_row_id(True)
@@ -1747,7 +1864,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",
@@ -1782,6 +1901,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,15 +1913,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())