Merge main into gatekeeper/fix-4272-1 and preserve query validation

This commit is contained in:
Gatefixer
2026-10-04 04:54:45 +00:00
30 changed files with 1173 additions and 69 deletions
+55 -6
View File
@@ -122,7 +122,48 @@ def _view_description(
class DBConnection(EnforceOverrides):
"""An active LanceDB connection interface."""
"""An active LanceDB connection interface.
Use [close][lancedb.db.DBConnection.close] to release the connection's
underlying resources, or use the connection as a context manager to close it
automatically when leaving the block, including when an exception is raised.
Examples
--------
>>> import lancedb
>>> with lancedb.connect("memory://") as db:
... assert db.is_open()
>>> db.is_open()
False
"""
def __enter__(self) -> DBConnection:
return self
def __exit__(self, *_) -> None:
self.close()
@abstractmethod
def is_open(self) -> bool:
"""Return True if the connection is open."""
pass
@abstractmethod
def close(self) -> None:
"""Close the connection, releasing any underlying resources.
It is safe to call this method multiple times. Database operations on a
closed connection raise ``RuntimeError: Connection is closed``.
Examples
--------
>>> import lancedb
>>> db = lancedb.connect("memory://")
>>> db.close()
>>> db.is_open()
False
"""
pass
def list_namespaces(
self,
@@ -368,9 +409,10 @@ class DBConnection(EnforceOverrides):
- [LanceModel][lancedb.pydantic.LanceModel]
mode: str; default "create"
The mode to use when creating the table.
Can be either "create" or "overwrite".
Can be "create", "overwrite", or "exist_ok".
By default, if the table already exists, an exception is raised.
If you want to overwrite the table, use mode="overwrite".
To open an existing table without adding data, use mode="exist_ok".
exist_ok: bool, default False
If a table by the same name already exists, then raise an exception
if exist_ok=False. If exist_ok=True, then open the existing table;
@@ -1275,6 +1317,14 @@ class LanceDBConnection(DBConnection):
def __repr__(self) -> str:
return f"{self.__class__.__name__}(uri={self._conn.uri!r})"
@override
def is_open(self) -> bool:
return self._conn.is_open()
@override
def close(self) -> None:
self._conn.close()
@override
def serialize(self) -> str:
import json
@@ -1497,8 +1547,6 @@ class LanceDBConnection(DBConnection):
"""
if namespace_path is None:
namespace_path = []
if mode.lower() not in ["create", "overwrite"]:
raise ValueError("mode must be either 'create' or 'overwrite'")
validate_table_name(name)
tbl = LanceTable.create(
@@ -2301,11 +2349,12 @@ class AsyncConnection(object):
- pyarrow.Schema
- [LanceModel][lancedb.pydantic.LanceModel]
mode: Literal["create", "overwrite"]; default "create"
mode: Literal["create", "overwrite", "exist_ok"]; default "create"
The mode to use when creating the table.
Can be either "create" or "overwrite".
Can be "create", "overwrite", or "exist_ok".
By default, if the table already exists, an exception is raised.
If you want to overwrite the table, use mode="overwrite".
To open an existing table without adding data, use mode="exist_ok".
exist_ok: bool, default False
If a table by the same name already exists, then raise an exception
if exist_ok=False. If exist_ok=True, then open the existing table;
+8
View File
@@ -489,6 +489,14 @@ class LanceNamespaceDBConnection(DBConnection):
)
self._uri = self._inner.uri
@override
def is_open(self) -> bool:
return self._inner.is_open()
@override
def close(self) -> None:
self._inner.close()
@override
def serialize(self) -> str:
import json
+13 -3
View File
@@ -40,7 +40,7 @@ from .rerankers.base import Reranker
from .rerankers.rrf import RRFReranker
from .rerankers.util import check_reranker_result
from .schema import blob_column_paths, is_blob_like_field, schema_has_blob_field
from .util import flatten_columns
from .util import _validate_query_vector, flatten_columns
from . import _wal_hybrid # WAL-PK-FUSION: delete.
from ._blob import (
BLOB_MODE_TO_HANDLING,
@@ -908,6 +908,9 @@ class LanceQueryBuilder(ABC):
fast_search: bool
Skip flat search of unindexed data.
"""
if query_type != "fts":
_validate_query_vector(query)
if ordering_field_name is not None:
import warnings
@@ -1620,6 +1623,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
str_query: Optional[str] = None,
fast_search: bool = None,
):
_validate_query_vector(query)
super().__init__(table)
self._query = query
self._distance_type = None
@@ -1694,13 +1698,15 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
Parameters
----------
nprobes: int
The number of probes to use.
The number of probes to use. Must be greater than 0.
Returns
-------
LanceVectorQueryBuilder
The LanceQueryBuilder object.
"""
if nprobes <= 0:
raise ValueError("Invalid input, nprobes must be greater than 0")
self._minimum_nprobes = nprobes
self._maximum_nprobes = nprobes
return self
@@ -2496,13 +2502,15 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
Parameters
----------
nprobes: int
The number of probes to use.
The number of probes to use. Must be greater than 0.
Returns
-------
LanceHybridQueryBuilder
The LanceHybridQueryBuilder object.
"""
if nprobes <= 0:
raise ValueError("Invalid input, nprobes must be greater than 0")
self._minimum_nprobes = nprobes
self._maximum_nprobes = nprobes
return self
@@ -3686,6 +3694,8 @@ class AsyncVectorQueryBase:
"""
Set the number of partitions to search (probe)
The number of probes must be greater than 0.
This argument is only used when the vector column has an IVF-based index.
If there is no index then this value is ignored.
+6 -2
View File
@@ -1054,6 +1054,10 @@ class RemoteDBConnection(DBConnection):
"""
return LOOP.run(self._conn.namespace_client())
async def close(self):
"""Close the connection to the database."""
@override
def is_open(self) -> bool:
return self._conn.is_open()
@override
def close(self) -> None:
self._conn.close()
+34 -14
View File
@@ -2617,6 +2617,10 @@ class Table(ABC):
to nullable. Currently, you cannot change a nullable column to
non-nullable.
Each alteration must specify at least one of "rename", "data_type",
or "nullable". The legacy key "name" is also accepted for renaming.
Unknown keys raise ValueError before any alterations are applied.
Returns
-------
AlterColumnsResult
@@ -4332,9 +4336,9 @@ class LanceTable(Table):
)
if storage_options is None:
storage_options = {}
storage_options["new_table_enable_v2_manifest_paths"] = (
storage_options["new_table_enable_v2_manifest_paths"] = str(
enable_v2_manifest_paths
)
).lower()
self._table = LOOP.run(
self._conn._conn.create_table(
@@ -4766,7 +4770,7 @@ class LanceTable(Table):
[LanceTable.uses_v2_manifest_paths][lancedb.table.LanceTable.uses_v2_manifest_paths]
to check if the table is already using the new path style.
"""
LOOP.run(self._table.migrate_v2_manifest_paths())
LOOP.run(self._table.migrate_manifest_paths_v2())
@deprecation.deprecated(
deprecated_in="0.33.1",
@@ -6122,7 +6126,17 @@ class AsyncTable:
@overload
async def search(
self,
query: Optional[str] = None,
query: None = None,
vector_column_name: Optional[str] = None,
query_type: QueryType = "auto",
ordering_field_name: Optional[str] = None,
fts_columns: Optional[Union[str, List[str]]] = None,
) -> AsyncQuery: ...
@overload
async def search(
self,
query: str,
vector_column_name: Optional[str] = None,
query_type: Literal["auto"] = ...,
ordering_field_name: Optional[str] = None,
@@ -6132,7 +6146,7 @@ class AsyncTable:
@overload
async def search(
self,
query: Optional[str] = None,
query: str,
vector_column_name: Optional[str] = None,
query_type: Literal["hybrid"] = ...,
ordering_field_name: Optional[str] = None,
@@ -6142,7 +6156,7 @@ class AsyncTable:
@overload
async def search(
self,
query: Optional[Union[VEC, "PIL.Image.Image", Tuple]] = None,
query: Union[VEC, "PIL.Image.Image", Tuple],
vector_column_name: Optional[str] = None,
query_type: Literal["auto"] = ...,
ordering_field_name: Optional[str] = None,
@@ -6152,7 +6166,7 @@ class AsyncTable:
@overload
async def search(
self,
query: Optional[str] = None,
query: str,
vector_column_name: Optional[str] = None,
query_type: Literal["fts"] = ...,
ordering_field_name: Optional[str] = None,
@@ -6162,9 +6176,7 @@ class AsyncTable:
@overload
async def search(
self,
query: Optional[
Union[VEC, str, "PIL.Image.Image", Tuple, FullTextQuery]
] = None,
query: Union[VEC, str, "PIL.Image.Image", Tuple, FullTextQuery],
vector_column_name: Optional[str] = None,
query_type: Literal["vector"] = ...,
ordering_field_name: Optional[str] = None,
@@ -6180,7 +6192,7 @@ class AsyncTable:
query_type: QueryType = "auto",
ordering_field_name: Optional[str] = None,
fts_columns: Optional[Union[str, List[str]]] = None,
) -> Union[AsyncHybridQuery, AsyncFTSQuery, AsyncVectorQuery]:
) -> Union[AsyncQuery, AsyncHybridQuery, AsyncFTSQuery, AsyncVectorQuery]:
"""Create a search query to find the nearest neighbors
of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/)
and [full-text search](https://lancedb.com/docs/search/full-text-search/).
@@ -6195,8 +6207,9 @@ class AsyncTable:
- *default None*.
Acceptable types are: list, np.ndarray, PIL.Image.Image
- If None then the select/where/limit clauses are applied to filter
the table
- If None then a plain [AsyncQuery][lancedb.query.AsyncQuery] is
returned, equivalent to calling [query][lancedb.table.AsyncTable.query].
The select/where/limit clauses are applied to filter the table.
vector_column_name: str, optional
The name of the vector column to search.
@@ -6224,10 +6237,13 @@ class AsyncTable:
Returns
-------
LanceQueryBuilder
AsyncQuery, AsyncHybridQuery, AsyncFTSQuery, or AsyncVectorQuery
A query builder object representing the query.
"""
if query is None:
return self.query()
def is_embedding(query):
return isinstance(query, (list, np.ndarray, pa.Array, pa.ChunkedArray))
@@ -6877,6 +6893,10 @@ class AsyncTable:
to nullable. Currently, you cannot change a nullable column to
non-nullable.
Each alteration must specify at least one of "rename", "data_type",
or "nullable". The legacy key "name" is also accepted for renaming.
Unknown keys raise ValueError before any alterations are applied.
Returns
-------
AlterColumnsResult
+11
View File
@@ -280,6 +280,14 @@ def infer_vector_column_dim(data_type: pa.DataType) -> Optional[int]:
return None
def _validate_query_vector(query: Any) -> None:
"""Reject empty vector inputs before vector-column inference or execution."""
if (isinstance(query, list) and not query) or (
isinstance(query, np.ndarray) and query.size == 0
):
raise ValueError("Query vector must not be empty")
def _query_vector_dim(query: Optional[Any]) -> Optional[int]:
if query is None:
return None
@@ -301,6 +309,9 @@ def infer_vector_column_name(
query: Optional[Any], # inferred later in query builder
vector_column_name: Optional[str],
):
if query_type != "fts":
_validate_query_vector(query)
if vector_column_name is not None:
return vector_column_name
+99
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import contextlib
import inspect
import re
import sys
@@ -350,6 +351,40 @@ async def test_create_table_from_iterator_async(mem_db_async: lancedb.AsyncConne
assert await table.count_rows() == 10
def test_create_exist_ok_mode(tmp_db: lancedb.DBConnection):
db = tmp_db
data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}]
table = db.create_table("t1", data, mode="exist_ok")
assert table.count_rows() == 2
original = table.to_arrow()
reopened = db.create_table("t1", [{"vector": [9.0, 8.0]}], mode="exist_ok")
assert reopened.to_arrow().equals(original)
def test_create_empty_exist_ok_mode(tmp_db: lancedb.DBConnection):
schema = pa.schema([pa.field("id", pa.int64())])
table = tmp_db.create_table("empty", schema=schema, mode="exist_ok")
assert table.schema == schema
assert table.count_rows() == 0
table.add([{"id": 7}])
reopened = tmp_db.create_table("empty", schema=schema, mode="exist_ok")
assert reopened.to_arrow().to_pylist() == [{"id": 7}]
bad_schema = pa.schema([pa.field("other", pa.int64())])
with pytest.raises(ValueError):
tmp_db.create_table("empty", schema=bad_schema, mode="exist_ok")
def test_create_table_invalid_mode(tmp_db: lancedb.DBConnection):
db = tmp_db
data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}]
with pytest.raises(ValueError, match="Invalid mode bogus"):
db.create_table("tb", data, mode="bogus")
assert "tb" not in db
def test_create_exist_ok(tmp_db: lancedb.DBConnection):
from conftest import pandas_string_type
@@ -405,6 +440,49 @@ async def test_connect(tmp_path):
assert str(db) == f"ListingDatabase(uri={tmp_path}, read_consistency_interval=5s)"
@pytest.fixture(params=["memory", "local", "namespace", "from_inner"])
def sync_db(request, tmp_path):
if request.param == "memory":
return lancedb.connect("memory://")
if request.param == "namespace":
return lancedb.connect_namespace("dir", {"root": str(tmp_path)})
db = lancedb.connect(tmp_path)
if request.param == "from_inner":
return lancedb.db.LanceDBConnection.from_inner(db._inner, None)
return db
def test_sync_close(sync_db):
assert sync_db.is_open()
assert sync_db.close() is None
assert not sync_db.is_open()
assert sync_db.close() is None
with pytest.warns(DeprecationWarning, match="table_names"):
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.table_names()
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.list_tables()
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.open_table("test")
@pytest.mark.parametrize("raise_error", [False, True])
def test_sync_context_manager(sync_db, raise_error):
with contextlib.ExitStack() as stack:
if raise_error:
stack.enter_context(pytest.raises(ValueError, match="test error"))
with sync_db as db:
assert db is sync_db
assert db.is_open()
if raise_error:
raise ValueError("test error")
assert not sync_db.is_open()
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.list_tables()
@pytest.mark.asyncio
async def test_close(mem_db_async: lancedb.AsyncConnection):
assert mem_db_async.is_open()
@@ -495,6 +573,27 @@ async def test_create_exist_ok_async(tmp_db_async: lancedb.AsyncConnection):
# await db.create_table("test", schema=bad_schema, exist_ok=True)
@pytest.mark.parametrize("enable_v2_manifest_paths", [False, True])
@pytest.mark.parametrize("empty", [False, True])
def test_create_table_deprecated_v2_manifest_paths(
tmp_path, enable_v2_manifest_paths, empty
):
db = lancedb.connect(tmp_path)
data = None if empty else [{"id": 1}]
with pytest.warns(DeprecationWarning, match="enable_v2_manifest_paths"):
table = db.create_table(
"test",
data=data,
schema=pa.schema([("id", pa.int64())]),
enable_v2_manifest_paths=enable_v2_manifest_paths,
)
assert table.uses_v2_manifest_paths() == enable_v2_manifest_paths
assert table.to_arrow().to_pylist() == ([] if empty else data)
assert db.open_table("test").uses_v2_manifest_paths() == enable_v2_manifest_paths
@pytest.mark.asyncio
async def test_create_table_v2_manifest_paths_async(tmp_path):
db_with_v2_paths = await lancedb.connect_async(
+53
View File
@@ -24,6 +24,7 @@ from lancedb.pydantic import LanceModel, Vector
from lancedb.query import (
AsyncFTSQuery,
AsyncHybridQuery,
AsyncQuery,
AsyncQueryBase,
AsyncVectorQuery,
ColumnOrdering,
@@ -974,6 +975,19 @@ def test_query_builder_with_filter(table):
assert all(np.array(rs[0]["vector"]) == [3, 4])
@pytest.mark.parametrize("query_type", ["vector", "hybrid"])
@pytest.mark.parametrize("nprobes", [0, -1])
def test_nprobes_nonpositive_sync(table, query_type, nprobes):
if query_type == "hybrid":
query = table.search(query_type="hybrid").vector([0, 0]).text("a")
else:
query = table.search([0, 0])
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(nprobes)
def test_invalid_nprobes_sync(table):
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
LanceVectorQueryBuilder(table, [0, 0], "vector").minimum_nprobes(0).to_list()
@@ -1005,6 +1019,18 @@ def test_multiple_nprobes_calls_works_sync(table):
).minimum_nprobes(20).to_list()
@pytest.mark.asyncio
@pytest.mark.parametrize("hybrid", [False, True])
async def test_nprobes_zero_async(table_async: AsyncTable, hybrid):
query = table_async.query().nearest_to([0, 0])
if hybrid:
query = query.nearest_to_text("dog")
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(0)
@pytest.mark.asyncio
async def test_invalid_nprobes_async(table_async: AsyncTable):
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
@@ -1279,6 +1305,27 @@ async def test_query_to_polars_async(table_async: AsyncTable):
assert df.shape == (0, num_columns)
@pytest.mark.asyncio
@pytest.mark.parametrize("search_kwargs", [{}, {"query": None}])
@pytest.mark.parametrize("with_vector", [False, True])
async def test_async_search_without_query(mem_db_async, search_kwargs, with_vector):
data = [{"id": i} for i in range(10)]
if with_vector:
for row in data:
row["vector"] = [float(row["id"]), 1.0]
table = await mem_db_async.create_table("test", data)
query = await table.search(**search_kwargs)
assert isinstance(query, AsyncQuery)
assert await query.limit(3).to_arrow() == await table.query().limit(3).to_arrow()
query = await table.search(**search_kwargs)
assert await query.where("id >= 8").select(["id"]).limit(3).to_list() == [
{"id": 8},
{"id": 9},
]
@pytest.mark.asyncio
async def test_none_query(table_async: AsyncTable):
with pytest.raises(ValueError):
@@ -2091,6 +2138,12 @@ def test_ensure_vector_query_nested_empty_list():
ensure_vector_query([[]])
@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)])
def test_vector_query_builder_empty_vector(table, query):
with pytest.raises(ValueError, match="^Query vector must not be empty$"):
LanceVectorQueryBuilder(table, query, "vector")
def test_fast_search(tmp_path):
db = lancedb.connect(tmp_path)
+209 -5
View File
@@ -16,8 +16,9 @@ import uuid
from packaging.version import Version
import lancedb
import numpy as np
from lancedb.conftest import MockTextEmbeddingFunction
from lancedb.query import ColumnOrdering
from lancedb.query import AsyncQuery, ColumnOrdering
from lancedb.remote import ClientConfig
from lancedb.remote.errors import HttpError, RetryError
import pytest
@@ -119,6 +120,77 @@ async def test_async_remote_db():
assert table_names == []
@pytest.mark.parametrize(
"alteration, match",
[
({"path": "id"}, "One of rename, nullable or data_type"),
(
{"path": "id", "nulable": False}, # spellchecker:disable-line
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
(
{
"path": "id",
"rename": "new_id",
"nulable": False, # spellchecker:disable-line
},
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
],
)
def test_remote_alter_columns_rejects_invalid_before_request(alteration, match):
requests = []
def handler(request):
requests.append(request.path)
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
with mock_lancedb_connection(handler) as db:
table = db.open_table("test")
requests.clear()
with pytest.raises(ValueError, match=match):
table.alter_columns(alteration)
assert requests == []
@pytest.mark.asyncio
@pytest.mark.parametrize(
"alteration, match",
[
({"path": "id"}, "One of rename, nullable or data_type"),
(
{"path": "id", "nulable": False}, # spellchecker:disable-line
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
],
)
async def test_async_remote_alter_columns_rejects_invalid_before_request(
alteration, match
):
requests = []
def handler(request):
requests.append(request.path)
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
async with mock_lancedb_connection_async(handler) as db:
table = await db.open_table("test")
requests.clear()
with pytest.raises(ValueError, match=match):
await table.alter_columns(alteration)
assert requests == []
@pytest.mark.asyncio
async def test_async_checkout():
def handler(request):
@@ -1521,6 +1593,60 @@ def test_query_sync_empty_query():
assert data == expected
@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)])
@pytest.mark.parametrize("vector_column_name", [None, "vector"])
@pytest.mark.parametrize("query_type", ["auto", "vector"])
def test_query_sync_empty_vector(query, vector_column_name, query_type):
def handler(body):
pytest.fail("An empty query vector must be rejected before sending a query")
with query_test_table(handler) as table:
with pytest.raises(ValueError, match="^Query vector must not be empty$"):
table.search(
query, vector_column_name=vector_column_name, query_type=query_type
).limit(3).to_arrow()
@pytest.mark.asyncio
@pytest.mark.parametrize("search_kwargs", [{}, {"query": None}])
async def test_async_search_without_query(search_kwargs):
expected = pa.table({"id": [7, 8, 9]})
def handler(request):
if request.path == "/v1/table/test/describe/":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
elif request.path == "/v1/table/test/query/":
body = json.loads(
request.rfile.read(int(request.headers["Content-Length"]))
)
assert body == {
"k": 3,
"filter": "id >= 7",
"vector": [],
"columns": ["id"],
"prefilter": True,
"version": None,
}
request.send_response(200)
request.send_header("Content-Type", "application/vnd.apache.arrow.file")
request.end_headers()
with pa.ipc.new_file(request.wfile, schema=expected.schema) as writer:
writer.write_table(expected)
else:
request.send_response(404)
request.end_headers()
async with mock_lancedb_connection_async(handler) as db:
table = await db.open_table("test")
query = await table.search(**search_kwargs)
assert isinstance(query, AsyncQuery)
result = await query.where("id >= 7").select(["id"]).limit(3).to_arrow()
assert result == expected
def test_query_sync_maximal():
def handler(body):
assert body == {
@@ -1581,6 +1707,53 @@ def test_query_sync_maximal():
)
@pytest.mark.parametrize("hybrid", [False, True])
def test_query_sync_nprobes_zero(hybrid):
query_requests = []
def handler(body):
query_requests.append(body)
return pa.table({"id": []})
with query_test_table(handler) as table:
if hybrid:
query = table.search(query_type="hybrid").vector([1, 2, 3]).text("dog")
else:
query = table.search([1, 2, 3])
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(0)
assert query_requests == []
@pytest.mark.asyncio
@pytest.mark.parametrize("hybrid", [False, True])
async def test_query_async_nprobes_zero(hybrid):
requests = []
def handler(request):
requests.append(request.path)
if request.path == "/v1/table/test/describe/":
send_json(request, {"version": 1, "schema": {"fields": []}})
else:
request.send_response(404)
request.end_headers()
async with mock_lancedb_connection_async(handler) as db:
table = await db.open_table("test")
query = table.query().nearest_to([1, 2, 3])
if hybrid:
query = query.nearest_to_text("dog")
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(0)
assert requests == ["/v1/table/test/describe/"]
def test_query_sync_nprobes():
def handler(body):
assert body == {
@@ -2145,15 +2318,46 @@ async def test_header_provider_overrides_static_headers():
def test_close():
"""Test that close() works without AttributeError."""
import asyncio
def handler(req):
req.send_response(200)
req.end_headers()
with mock_lancedb_connection(handler) as db:
asyncio.run(db.close())
assert db.close() is None
assert not db.is_open()
assert db.close() is None
with pytest.warns(DeprecationWarning, match="table_names"):
with pytest.raises(RuntimeError, match="Connection is closed"):
db.table_names()
with pytest.raises(RuntimeError, match="Connection is closed"):
db.list_tables()
with pytest.raises(RuntimeError, match="Connection is closed"):
db.open_table("test")
@pytest.mark.parametrize("raise_error", [False, True])
def test_sync_context_manager(raise_error):
def handler(req):
req.send_response(200)
req.send_header("Content-Type", "application/json")
req.end_headers()
req.wfile.write(b'{"tables": []}')
with mock_lancedb_connection(handler) as db:
with contextlib.ExitStack() as stack:
if raise_error:
stack.enter_context(pytest.raises(ValueError, match="test error"))
with db as conn:
assert conn is db
assert conn.is_open()
assert conn.list_tables().tables == []
if raise_error:
raise ValueError("test error")
assert not db.is_open()
with pytest.raises(RuntimeError, match="Connection is closed"):
db.list_tables()
@pytest.mark.parametrize("exception", [KeyboardInterrupt, SystemExit, GeneratorExit])
+109 -4
View File
@@ -113,6 +113,32 @@ def test_basic(mem_db: DBConnection):
assert table.to_arrow() == expected_data
@pytest.mark.parametrize("enable_v2", [False, True])
def test_migrate_v2_manifest_paths(tmp_path, enable_v2):
db = lancedb.connect(
tmp_path,
storage_options={"new_table_enable_v2_manifest_paths": str(enable_v2).lower()},
)
table = db.create_table("calls", [{"id": 1, "vector": [1.0, 1.0]}])
table.add([{"id": 2, "vector": [2.0, 2.0]}])
expected_data = table.to_arrow()
expected_versions = table.list_versions()
assert table.uses_v2_manifest_paths() == enable_v2
# Migration is also safe to repeat on a table already using v2 paths.
for _ in range(2):
table.migrate_v2_manifest_paths()
assert table.uses_v2_manifest_paths()
reopened = db.open_table("calls")
assert reopened.uses_v2_manifest_paths()
assert reopened.to_arrow() == expected_data
assert reopened.list_versions() == expected_versions
manifests = list((tmp_path / "calls.lance" / "_versions").glob("*.manifest"))
assert len(manifests) == len(expected_versions)
assert all(len(path.stem) == 20 and path.stem.isdigit() for path in manifests)
def test_search_preserves_nulls_from_sliced_arrow_table(mem_db: DBConnection):
data = pa.table(
{
@@ -3751,6 +3777,24 @@ def test_empty_query(mem_db: DBConnection):
assert df.num_rows == 42
@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)])
@pytest.mark.parametrize("vector_column_name", [None, "vector"])
@pytest.mark.parametrize("query_type", ["auto", "vector"])
@pytest.mark.parametrize("multiple_vector_columns", [False, True])
def test_search_empty_vector(
mem_db, query, vector_column_name, query_type, multiple_vector_columns
):
fields = [pa.field("vector", pa.list_(pa.float32(), 8))]
if multiple_vector_columns:
fields.append(pa.field("vec2", pa.list_(pa.float32(), 4)))
table = mem_db.create_table("empty_vector_query", schema=pa.schema(fields))
with pytest.raises(ValueError, match="^Query vector must not be empty$"):
table.search(
query, vector_column_name=vector_column_name, query_type=query_type
).limit(3).to_arrow()
def test_search_with_schema_inf_single_vector(mem_db: DBConnection):
class MyTable(LanceModel):
text: str
@@ -4144,14 +4188,74 @@ async def test_add_columns_with_schema(mem_db_async: AsyncConnection):
)
def test_alter_columns(mem_db: DBConnection):
@pytest.mark.parametrize("rename_key", ["rename", "name"])
def test_alter_columns(mem_db: DBConnection, rename_key):
data = pa.table({"id": [0, 1]})
table = mem_db.create_table("my_table", data=data)
alter_columns_res = table.alter_columns({"path": "id", "rename": "new_id"})
alter_columns_res = table.alter_columns({"path": "id", rename_key: "new_id"})
assert alter_columns_res.version == 2
assert table.to_arrow().column_names == ["new_id"]
INVALID_COLUMN_ALTERATIONS = [
(({"path": "id"},), "One of rename, nullable or data_type"),
(
({"path": "id", "nulable": False},), # spellchecker:disable-line
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
(
(
{
"path": "id",
"rename": "new_id",
"nulable": False, # spellchecker:disable-line
},
),
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
(
({"path": "id", "rename": "new_id"}, {"path": "id"}),
"One of rename, nullable or data_type",
),
]
def test_alter_columns_nullable_false(mem_db: DBConnection):
table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]}))
result = table.alter_columns({"path": "id", "nullable": False})
assert result.version == 2
assert not table.schema.field("id").nullable
@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS)
def test_alter_columns_rejects_invalid(mem_db: DBConnection, alterations, match):
table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]}))
initial_version = table.version
initial_schema = table.schema
with pytest.raises(ValueError, match=match):
table.alter_columns(*alterations)
assert table.version == initial_version
assert table.schema == initial_schema
@pytest.mark.asyncio
@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS)
async def test_alter_columns_rejects_invalid_async(
mem_db_async: AsyncConnection, alterations, match
):
table = await mem_db_async.create_table("my_table", data=pa.table({"id": [0, 1]}))
initial_version = await table.version()
initial_schema = await table.schema()
with pytest.raises(ValueError, match=match):
await table.alter_columns(*alterations)
assert await table.version() == initial_version
assert await table.schema() == initial_schema
def test_update_field_metadata(mem_db: DBConnection):
data = pa.table({"id": [0, 1], "category": ["a", "b"]})
table = mem_db.create_table("my_table", data=data)
@@ -4177,10 +4281,11 @@ def test_update_field_metadata(mem_db: DBConnection):
@pytest.mark.asyncio
async def test_alter_columns_async(mem_db_async: AsyncConnection):
@pytest.mark.parametrize("rename_key", ["rename", "name"])
async def test_alter_columns_async(mem_db_async: AsyncConnection, rename_key):
data = pa.table({"id": [0, 1]})
table = await mem_db_async.create_table("my_table", data=data)
alter_columns_res = await table.alter_columns({"path": "id", "rename": "new_id"})
alter_columns_res = await table.alter_columns({"path": "id", rename_key: "new_id"})
assert alter_columns_res.version == 2
assert (await table.to_arrow()).column_names == ["new_id"]
alter_columns_res = await table.alter_columns(
+5 -4
View File
@@ -995,8 +995,9 @@ impl VectorQuery {
self.inner = self.inner.clone().refine_factor(refine_factor);
}
pub fn nprobes(&mut self, nprobe: u32) {
self.inner = self.inner.clone().nprobes(nprobe as usize);
pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> {
self.inner = self.inner.clone().nprobes(nprobe as usize).infer_error()?;
Ok(())
}
pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> {
@@ -1188,8 +1189,8 @@ impl HybridQuery {
self.inner_vec.refine_factor(refine_factor);
}
pub fn nprobes(&mut self, nprobe: u32) {
self.inner_vec.nprobes(nprobe);
pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> {
self.inner_vec.nprobes(nprobe)
}
pub fn ef(&mut self, ef: u32) {
+11
View File
@@ -1902,6 +1902,17 @@ impl Table {
let alterations = alterations
.iter()
.map(|alteration| {
for key in alteration.keys().iter() {
let key = key.extract::<String>()?;
if !matches!(
key.as_str(),
"path" | "rename" | "name" | "nullable" | "data_type"
) {
return Err(PyValueError::new_err(format!(
"Unknown column alteration key '{key}'"
)));
}
}
let path = alteration
.get_item("path")?
.ok_or_else(|| PyValueError::new_err("Missing path"))?