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

Merge main into gatekeeper/fix-4391-1

Browse Source
This commit is contained in:
Gatefixer
2026-10-07 03:27:35 +00:00
parent db0612dd76 bbc4aebe55
commit 502ba17021
2 changed files with 164 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/remote/table.py
+6
View File
@@ -736,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.
python/python/tests/test_remote_db.py
+158 -53
View File
@@ -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, MatchQuery
from lancedb.query import (
AsyncQuery,
ColumnOrdering,
LanceVectorQueryBuilder,
MatchQuery,
)
from lancedb.remote import ClientConfig
from lancedb.remote.errors import HttpError, RetryError
import pytest
@@ -1607,34 +1612,46 @@ def query_test_table(
query_handler,
*,
server_version=Version("0.1.0"),
vector_columns=(("vector", 3),),
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()
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"))
@@ -1679,9 +1696,12 @@ def query_test_table(
],
)
def test_query_object_sync_matches_local(mem_db, query, search_kwargs, expected_column):
vector_columns = (("vector", 8), ("vec2", 4))
vector_columns = {"vector": 8, "vec2": 4}
schema = pa.schema(
[pa.field(name, pa.list_(pa.float32(), dim)) for name, dim in vector_columns]
[
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()
@@ -1696,7 +1716,7 @@ def test_query_object_sync_matches_local(mem_db, query, search_kwargs, expected_
@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 ()
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
@@ -1705,9 +1725,12 @@ def test_query_object_sync_preserves_fast_search(fast_search, query, query_type)
def test_query_sync_rejects_ambiguous_vector_column(mem_db):
vector_columns = (("vector", 4), ("vec2", 4))
vector_columns = {"vector": 4, "vec2": 4}
schema = pa.schema(
[pa.field(name, pa.list_(pa.float32(), dim)) for name, dim in vector_columns]
[
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:
@@ -1719,7 +1742,7 @@ def test_query_sync_rejects_ambiguous_vector_column(mem_db):
assert str(remote_error.value) == str(local_error.value)
@pytest.mark.parametrize("vector_columns", [(), (("v1", 4), ("v2", 8))])
@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 = []
@@ -1743,14 +1766,14 @@ def test_query_sync_auto_with_embedding_function():
assert "full_text_query" not in body
return pa.table({"id": [1]})
with query_test_table(handler, vector_columns=(("vector", 10),)) as table:
table.embedding_functions = {
"vector": EmbeddingFunctionConfig(
source_column="text",
vector_column="vector",
function=MockTextEmbeddingFunction(),
)
}
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}]
@@ -1797,7 +1820,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,
@@ -1810,7 +1910,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
@@ -2060,7 +2160,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": {
@@ -2075,8 +2177,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 == {
@@ -2103,9 +2205,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()
@@ -2143,7 +2247,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",
@@ -2190,15 +2296,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, vector_columns=(("vector", 10),)) 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())