mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 12:12:27 +00:00
Merge main into gatekeeper/fix-4272-1 and preserve query validation
This commit is contained in:
@@ -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;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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])
|
||||
|
||||
@@ -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
@@ -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) {
|
||||
|
||||
@@ -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"))?
|
||||
|
||||
Reference in New Issue
Block a user