Merge remote-tracking branch 'origin/main' into xuanwo/pr3851-merge

This commit is contained in:
Xuanwo
2026-10-26 23:35:36 +08:00
137 changed files with 6687 additions and 1967 deletions
+1 -2
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb-python"
version = "0.41.0-beta.13"
version = "0.41.0-beta.14"
publish = false
edition.workspace = true
description = "Python bindings for LanceDB"
@@ -22,7 +22,6 @@ bytes.workspace = true
lancedb.workspace = true
datafusion-common.workspace = true
lance-core.workspace = true
# The workspace path dependency includes the patched crate in Maturin sources.
lance-arrow.workspace = true
lance-namespace.workspace = true
lance-namespace-impls.workspace = true
+76
View File
@@ -40,6 +40,33 @@ Stable releases are created about every 2 weeks. For the latest features and bug
pip install --pre --extra-index-url https://pypi.fury.io/lancedb/ lancedb
```
If you also need the optional `lance` Python API, install the `pylance` extra
and include the Lance preview index:
```bash
pip install --pre \
--extra-index-url https://pypi.fury.io/lancedb/ \
--extra-index-url https://pypi.fury.io/lance-format/ \
"lancedb[pylance]"
```
With uv, also select the best matching version across the indexes:
```bash
uv pip install --prerelease allow \
--index https://pypi.fury.io/lancedb/ \
--index https://pypi.fury.io/lance-format/ \
--index-strategy unsafe-best-match \
"lancedb[pylance]"
```
The LanceDB preview index contains old `pylance` releases; current Lance previews
are published to the `lance-format` index. By default, uv only considers versions
from the first index containing a package, which can select an old `pylance` or
backtrack to an older LanceDB release. `unsafe-best-match` searches all configured
indexes, like pip; use it only with indexes you trust. Install the extra rather
than an unconstrained `pylance` so LanceDB's minimum supported version is enforced.
### Threading in CPU-limited containers
LanceDB uses separate pools for compute work and storage I/O. On a container with
@@ -74,3 +101,52 @@ print(results)
### Development
See [CONTRIBUTING.md](./CONTRIBUTING.md) for information on how to contribute to LanceDB.
### Opening SQL blob descriptors
A SQL `blob_descriptor(video)` result keeps the Blob V2 storage metadata and a
`source` struct containing its table URI, version, commit timestamp, branch,
column and row identity. This example uses a remote SQL connection whose SQL
service produces those source-bearing descriptors. Open the reference on the
source table to read incrementally:
```python
rows = db.execute_query("""
SELECT f.timestamp, blob_descriptor(v.video) AS video
FROM videos v JOIN frames f ON v.id = f.video_id
""").read_all()
videos = db.open_table("videos")
for row in rows.to_pylist():
if row["video"] is not None:
with videos.open_blob_file(row["video"]) as video:
frame = extract_frame(video, row["timestamp"])
```
`open_blob_file` accepts the dictionary or an Arrow struct scalar directly. It
keeps the query's version and branch without changing the table handle's checkout.
Payload bytes load on reads. An external blob with an unspecified size may need a
metadata-only HEAD request to open.
`videos.open_blob_files(rows["video"])` returns an iterator over the same handles,
preserving order, duplicates and nulls (`None`). It accepts Arrow arrays/chunked
arrays or another iterable and opens each handle as iteration advances. Close each
non-null handle after use. Async tables offer `await videos.open_blob_file(...)`
and `async for video in videos.open_blob_files(...)`.
```python
for video in videos.open_blob_files(rows["video"]):
if video is not None:
with video:
header = video.read(4096)
```
To extract a frame while retaining each row's timestamp, iterate over the
handles and rows together:
```python
for row, video in zip(rows.to_pylist(), videos.open_blob_files(rows["video"])):
if video is not None:
with video:
frame = extract_frame(video, row["timestamp"])
```
+3 -2
View File
@@ -45,7 +45,8 @@ repository = "https://github.com/lancedb/lancedb"
[project.optional-dependencies]
pylance = [
"pylance>=5.0.0b5",
# Keep the minimum in the bundled Lance major and compatible with tests below.
"pylance>=15.0.0b1",
]
# A library only needs the OpenTelemetry API; the application supplies and
# configures the SDK (the actual exporter/reader). See
@@ -63,7 +64,7 @@ tests = [
"polars>=0.19, <=1.32.3",
"pyarrow<25",
"pyarrow-stubs>=16.0",
"pylance==14.0.0b12",
"pylance==15.0.0b1",
"requests>=2.31.0",
"datafusion>=54,<55",
"opentelemetry-sdk>=1.30.0",
+16
View File
@@ -24,6 +24,7 @@ from .remote import ClientConfig
from .remote.db import RemoteDBConnection
from .expr import Expr, col, lit, func
from .schema import blob, vector
from ._blob_reader import BlobReader as BlobReader
from .job import AsyncJob, Job
from .sql import AsyncQuery as AsyncSqlQuery
from .sql import Query as SqlQuery
@@ -280,6 +281,21 @@ def connect(
"uri is required when not connecting through a namespace client"
)
if isinstance(uri, str) and uri.startswith("db://"):
# Preserve the deprecated remote timeout options, but validate all other
# unknown keywords just as the local connection does.
unknown_kwargs = {
key: value
for key, value in kwargs.items()
if key not in ("connection_timeout", "read_timeout")
}
if unknown_kwargs:
raise ValueError(f"Unknown keyword arguments: {unknown_kwargs}")
if session is not None:
raise NotImplementedError("session is not supported for remote connections")
if manifest_enabled:
raise NotImplementedError(
"manifest_enabled is not supported for remote connections"
)
if api_key is None:
api_key = os.environ.get("LANCEDB_API_KEY")
if api_key is None:
+13 -1
View File
@@ -7,7 +7,7 @@ from __future__ import annotations
import io
import operator
from collections.abc import Awaitable, Callable, Iterable
from collections.abc import Awaitable, Callable, Iterable, Mapping
from typing import TYPE_CHECKING, Optional, Union
import pyarrow as pa
@@ -31,6 +31,18 @@ FetchBlobsSync = Callable[[str, list[int]], pa.Array | pa.ChunkedArray]
FetchBlobsAsync = Callable[[str, list[int]], Awaitable[pa.Array | pa.ChunkedArray]]
def _normalize_blob_descriptor(value):
if isinstance(value, pa.Scalar):
value = value.as_py()
if value is None:
return None
if not isinstance(value, Mapping):
raise TypeError(
"blob descriptor must be a mapping, an Arrow struct scalar, or None"
)
return dict(value)
class BlobFile(io.RawIOBase):
"""Seekable lazy handle from :meth:`~lancedb.table.Table.fetch_blob_files`.
+120
View File
@@ -0,0 +1,120 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Seekable binary input for Functions that consume Blob V2 values lazily."""
import io
import operator
import urllib.request
from collections.abc import Callable
from http import HTTPStatus
_READ_CHUNK_BYTES = 1024 * 1024
_READ_TIMEOUT_SECONDS = 300
class BlobReader(io.RawIOBase):
"""Read a blob incrementally inside a Function.
Annotating a Function input with ``BlobReader`` requests a lazy, seekable
reader. Annotating it with ``bytes`` requests materialized bytes instead.
``read_range`` does not change the sequential cursor. Readers are scoped
to the invocation and must not be retained for a later call.
For local calls, construct a reader from a range-reading callable and the
blob's size. The callable accepts an offset and length and returns bytes.
"""
def __init__(self, read_range: Callable[[int, int], bytes], size: int):
super().__init__()
self._read_range = read_range
self._size = operator.index(size)
if self._size < 0:
raise ValueError("blob size must be nonnegative")
self._position = 0
@classmethod
def _from_uri(cls, uri: str, size: int):
opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
def read_range(offset, length):
request = urllib.request.Request(
uri, headers={"Range": f"bytes={offset}-{offset + length - 1}"}
)
with opener.open(request, timeout=_READ_TIMEOUT_SECONDS) as response:
if response.status != HTTPStatus.PARTIAL_CONTENT:
raise OSError(
"blob range service did not return a partial response"
)
data = response.read(length + 1)
if len(data) != length:
raise OSError(
"blob range service returned an invalid payload length"
)
return data
return cls(read_range, size)
def readable(self) -> bool:
return not self.closed
def seekable(self) -> bool:
return not self.closed
def size(self) -> int:
"""Return the blob size without reading its payload."""
self._checkClosed()
return self._size
def tell(self) -> int:
self._checkClosed()
return self._position
def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
self._checkClosed()
offset = operator.index(offset)
if whence == io.SEEK_SET:
position = offset
elif whence == io.SEEK_CUR:
position = self._position + offset
elif whence == io.SEEK_END:
position = self._size + offset
else:
raise ValueError(f"invalid whence: {whence}")
if position < 0:
raise ValueError("negative seek position")
self._position = position
return position
def read_range(self, offset: int, length: int) -> bytes:
"""Read a byte range without changing the cursor; truncate at EOF."""
self._checkClosed()
offset, length = operator.index(offset), operator.index(length)
if offset < 0 or length < 0:
raise ValueError("offset and length must be nonnegative")
remaining = min(length, max(0, self._size - offset))
chunks = []
while remaining:
count = min(remaining, _READ_CHUNK_BYTES)
chunk = self._read_range(offset, count)
if len(chunk) != count:
raise OSError("blob range reader returned an invalid payload length")
chunks.append(chunk)
offset += count
remaining -= count
return b"".join(chunks)
def read(self, size: int = -1) -> bytes:
self._checkClosed()
size = -1 if size is None else operator.index(size)
if size < 0:
size = max(0, self._size - self._position)
data = self.read_range(self._position, size)
self._position += len(data)
return data
def readinto(self, buffer) -> int:
view = memoryview(buffer).cast("B")
data = self.read(len(view))
view[: len(data)] = data
return len(data)
+8 -6
View File
@@ -19,11 +19,9 @@ from .index import (
FTS,
)
from lance_namespace import (
ListNamespacesResponse,
CreateNamespaceResponse,
DropNamespaceResponse,
DescribeNamespaceResponse,
ListTablesResponse,
)
from .remote import ClientConfig
from .types import BaseTokenizerType
@@ -114,12 +112,12 @@ class Connection(object):
uri: str
async def is_open(self): ...
async def close(self): ...
async def list_namespaces(
def list_namespaces(
self,
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListNamespacesResponse: ...
) -> NameListing: ...
async def create_namespace(
self,
namespace_path: List[str],
@@ -136,12 +134,12 @@ class Connection(object):
self,
namespace_path: List[str],
) -> DescribeNamespaceResponse: ...
async def list_tables(
def list_tables(
self,
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListTablesResponse: ...
) -> NameListing: ...
async def table_names(
self,
namespace_path: Optional[List[str]],
@@ -322,6 +320,8 @@ class Job:
@property
def id(self) -> Optional[str]: ...
@property
def job_type(self) -> Optional[str]: ...
@property
def _state(self) -> Optional[str]: ...
@property
def _description(self) -> Optional[JobDescription]: ...
@@ -580,6 +580,7 @@ class Table:
async def fetch_blob_files(
self, column: str, row_ids: list[int]
) -> list[Optional[BlobFile]]: ...
async def open_blob_file(self, descriptor: str) -> BlobFile: ...
def vector_search(self) -> VectorQuery: ...
class Tags:
@@ -1010,6 +1011,7 @@ async def connect_catalog(
class NameListing:
def num_page_results(self) -> int: ...
def page_token(self) -> Optional[str]: ...
async def next_page(self) -> Optional[dict[str, Any]]: ...
def __aiter__(self) -> NameListing: ...
async def __anext__(self) -> str: ...
+116 -70
View File
@@ -32,18 +32,24 @@ if sys.version_info >= (3, 12):
else:
from overrides import EnforceOverrides, override # type: ignore
from lancedb.listing import AsyncListing, Listing
from lancedb.listing import (
AsyncListing,
AsyncNamespaceListing,
AsyncTableListing,
Listing,
NamespaceListing,
TableListing,
validate_page_limit,
)
from lancedb.embeddings.registry import EmbeddingFunctionRegistry
from lancedb.common import data_to_reader, sanitize_uri, validate_schema
from lancedb.background_loop import LOOP
from lance_namespace import (
LanceNamespace,
ListNamespacesResponse,
CreateNamespaceResponse,
DropNamespaceResponse,
DescribeNamespaceResponse,
ListTablesResponse,
connect as namespace_connect,
)
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
@@ -171,7 +177,7 @@ class DBConnection(EnforceOverrides):
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListNamespacesResponse:
) -> NamespaceListing:
"""List immediate child namespace names in the given namespace.
Parameters
@@ -187,12 +193,12 @@ class DBConnection(EnforceOverrides):
Returns
-------
ListNamespacesResponse
NamespaceListing
Response containing namespace names and optional page_token for pagination.
"""
if namespace_path is None:
namespace_path = []
return ListNamespacesResponse(namespaces=[], page_token=None)
return NamespaceListing(AsyncNamespaceListing(None, empty=True))
def create_namespace(
self,
@@ -317,8 +323,8 @@ class DBConnection(EnforceOverrides):
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListTablesResponse:
"""List a page of tables in this database.
) -> TableListing:
"""List tables in this database, one page at a time.
Parameters
----------
@@ -334,7 +340,7 @@ class DBConnection(EnforceOverrides):
Returns
-------
ListTablesResponse
TableListing
Response containing table names and optional page_token for pagination.
"""
raise NotImplementedError(
@@ -422,10 +428,13 @@ class DBConnection(EnforceOverrides):
it will not add the provided data but will validate against any
schema that's specified.
on_bad_vectors: str, default "error"
What to do if any of the vectors are not the same size or contains NaNs.
One of "error", "drop", "fill".
How to handle vectors with incorrect sizes, NaNs, or null elements.
One of "error", "drop", "fill", "null".
fill_value: float
The value to use when filling vectors. Only used if on_bad_vectors="fill".
NaN or null elements and missing trailing elements are replaced with this
value. Valid elements are preserved, and vectors longer than the expected
dimension are truncated.
storage_options: dict, optional
Additional options for the storage backend. Options already set on the
connection will be inherited by the table, but can be overridden here.
@@ -542,7 +551,9 @@ class DBConnection(EnforceOverrides):
LanceTable(name='table4', ...)
"""
raise NotImplementedError
raise NotImplementedError(
"create_table is not supported for this connection type"
)
def __getitem__(self, name: str) -> LanceTable:
return self.open_table(name)
@@ -597,7 +608,9 @@ class DBConnection(EnforceOverrides):
-------
A LanceTable object representing the table.
"""
raise NotImplementedError
raise NotImplementedError(
"open_table is not supported for this connection type"
)
def create_materialized_view(
self,
@@ -664,10 +677,11 @@ class DBConnection(EnforceOverrides):
) -> Job[None]:
"""Submit materialized-view creation and return its job.
The job may already be complete for a local database. On LanceDB
Cloud and Enterprise, its ``id`` is the server job identifier from
the ``202 Accepted`` create response. Wait for the job before opening
or querying the view.
The job may already be complete for a local database. Remote creation
submits CREATE MATERIALIZED VIEW through the SQL server. Its job has
``job_type == "query"`` and ``id`` identifies the connection-scoped
query. Registry event history is unavailable.
Wait for the job before opening or querying the view.
"""
raise NotImplementedError(
"materialized views are not supported on this connection type"
@@ -737,13 +751,17 @@ class DBConnection(EnforceOverrides):
"""
if namespace_path is None:
namespace_path = []
raise NotImplementedError
raise NotImplementedError(
"drop_table is not supported for this connection type"
)
def drop_table_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Job:
"""Start dropping a table and return its cleanup job."""
raise NotImplementedError
raise NotImplementedError(
"drop_table_async is not supported for this connection type"
)
def rename_table(
self,
@@ -771,19 +789,26 @@ class DBConnection(EnforceOverrides):
cur_namespace_path = []
if new_namespace_path is None:
new_namespace_path = []
raise NotImplementedError
raise NotImplementedError(
"rename_table is not supported for this connection type"
)
def drop_database(self):
"""
Drop database
This is the same thing as dropping all the tables
"""
raise NotImplementedError
raise NotImplementedError(
"drop_database is not supported for this connection type"
)
def drop_all_tables(self, namespace_path: Optional[List[str]] = None):
"""
Drop all tables from the database
Remote connections do not currently support this operation. Use
``list_tables()`` and ``drop_table()`` to drop tables individually.
Parameters
----------
namespace_path: List[str], optional
@@ -792,12 +817,32 @@ class DBConnection(EnforceOverrides):
"""
if namespace_path is None:
namespace_path = []
raise NotImplementedError
raise NotImplementedError(
"drop_all_tables is not supported for this connection type"
)
@property
def uri(self) -> str:
"""The URI of this database."""
return self._uri
def _all_table_names(self) -> Generator[str, None, None]:
page_token = None
while True:
response = self.list_tables(page_token=page_token)
yield from response.tables
page_token = response.page_token
if not page_token:
return
def __len__(self) -> int:
"""Return the number of tables in the root namespace across all pages."""
return sum(1 for _ in self._all_table_names())
def __contains__(self, name: str) -> bool:
"""Return whether the root namespace contains the named table."""
return name in self._all_table_names()
def namespace_client(self) -> LanceNamespace:
"""Get the equivalent namespace client for this connection.
@@ -1230,18 +1275,19 @@ class DBConnection(EnforceOverrides):
parameters=[np.array([0.1, 0.2, 0.3], dtype=np.float32), 10],
)
Parameters are supported in queries, not in DDL or DML. A
parameterized query runs on the call that returns its rows:
cancelling it, or dropping its handle before reading, stops it on the
server. Its status follows that stream as well: ``describe()``
reports ``running`` until the reader has received every row, so open
``reader()`` directly rather than waiting for ``finished``.
Parameters are supported in queries, not in DDL or DML. With or
without parameters, the server chooses whether to stream rows on the
initial call or return a polling continuation. The reader follows
either path transparently. Dropping an unread handle or its reader
stops a directly streamed query; deferred execution requires explicit
cancellation. Open ``reader()`` directly: a streamed query needs its
reader to make progress and reports ``running`` until the stream ends.
Local connections do not support SQL.
"""
raise NotImplementedError("SQL is not supported for this connection type")
def describe_query(self, query_id: UUID) -> QueryDescription:
def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription:
"""Describe a submitted SQL query by its connection-scoped id."""
raise NotImplementedError("SQL is not supported for this connection type")
@@ -1436,7 +1482,7 @@ class LanceDBConnection(DBConnection):
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListNamespacesResponse:
) -> NamespaceListing:
"""List immediate child namespace names in the given namespace.
Parameters
@@ -1452,16 +1498,14 @@ class LanceDBConnection(DBConnection):
Returns
-------
ListNamespacesResponse
NamespaceListing
Response containing namespace names and optional page_token for pagination.
"""
if namespace_path is None:
namespace_path = []
return LOOP.run(
return NamespaceListing(
self._conn.list_namespaces(
namespace_path=namespace_path,
page_token=page_token,
limit=limit,
namespace_path=namespace_path, page_token=page_token, limit=limit
)
)
@@ -1512,8 +1556,8 @@ class LanceDBConnection(DBConnection):
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListTablesResponse:
"""List a page of tables in this database.
) -> TableListing:
"""List tables in this database, one page at a time.
Parameters
----------
@@ -1529,12 +1573,12 @@ class LanceDBConnection(DBConnection):
Returns
-------
ListTablesResponse
TableListing
Response containing table names and optional page_token for pagination.
"""
if namespace_path is None:
namespace_path = []
return LOOP.run(
return TableListing(
self._conn.list_tables(
namespace_path=namespace_path, page_token=page_token, limit=limit
)
@@ -1583,21 +1627,6 @@ class LanceDBConnection(DBConnection):
)
)
def _all_table_names(self) -> Generator[str, None, None]:
page_token = None
while True:
response = self.list_tables(page_token=page_token)
yield from response.tables
page_token = response.page_token
if not page_token:
return
def __len__(self) -> int:
return sum(1 for _ in self._all_table_names())
def __contains__(self, name: str) -> bool:
return name in self._all_table_names()
@override
def create_table(
self,
@@ -2286,12 +2315,12 @@ class AsyncConnection(object):
else:
return None
async def list_namespaces(
def list_namespaces(
self,
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListNamespacesResponse:
) -> AsyncNamespaceListing:
"""List immediate child namespace names in the given namespace.
Parameters
@@ -2306,15 +2335,18 @@ class AsyncConnection(object):
Returns
-------
ListNamespacesResponse
AsyncNamespaceListing
Response containing namespace names and optional pagination token
"""
if namespace_path is None:
namespace_path = []
result = await self._inner.list_namespaces(
validate_page_limit(limit)
inner = self._inner.list_namespaces(
namespace_path=namespace_path, page_token=page_token, limit=limit
)
return ListNamespacesResponse(**result)
if limit == 0:
return AsyncNamespaceListing(None, empty=True)
return AsyncNamespaceListing(inner, page_token)
async def create_namespace(
self,
@@ -2393,13 +2425,13 @@ class AsyncConnection(object):
result = await self._inner.describe_namespace(namespace_path)
return DescribeNamespaceResponse(**result)
async def list_tables(
def list_tables(
self,
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListTablesResponse:
"""List a page of tables in this database.
) -> AsyncTableListing:
"""List tables in this database, one page at a time.
Parameters
----------
@@ -2415,15 +2447,18 @@ class AsyncConnection(object):
Returns
-------
ListTablesResponse
AsyncTableListing
Response containing table names and optional page_token for pagination.
"""
if namespace_path is None:
namespace_path = []
result = await self._inner.list_tables(
validate_page_limit(limit)
inner = self._inner.list_tables(
namespace_path=namespace_path, page_token=page_token, limit=limit
)
return ListTablesResponse(**result)
if limit == 0:
return AsyncTableListing(None, empty=True)
return AsyncTableListing(inner, page_token)
async def table_names(
self,
@@ -2521,10 +2556,13 @@ class AsyncConnection(object):
it will not add the provided data but will validate against any
schema that's specified.
on_bad_vectors: str, default "error"
What to do if any of the vectors are not the same size or contains NaNs.
One of "error", "drop", "fill".
How to handle vectors with incorrect sizes, NaNs, or null elements.
One of "error", "drop", "fill", "null".
fill_value: float
The value to use when filling vectors. Only used if on_bad_vectors="fill".
NaN or null elements and missing trailing elements are replaced with this
value. Valid elements are preserved, and vectors longer than the expected
dimension are truncated.
storage_options: dict, optional
Additional options for the storage backend. Options already set on the
connection will be inherited by the table, but can be overridden here.
@@ -2841,7 +2879,9 @@ class AsyncConnection(object):
) -> AsyncJob[None]:
"""Submit materialized-view creation and return its job.
Wait for the returned job before opening or querying the view.
Wait for the returned job before opening or querying the view. Remote
creation uses SQL; ``job_type`` is ``"query"`` and ``id``
identifies that query on this connection.
"""
source, native_options = await self._materialized_view_source(
source, select, where, limit
@@ -3381,8 +3421,14 @@ class AsyncConnection(object):
)
)
async def describe_query(self, query_id: UUID) -> QueryDescription:
"""Describe a submitted SQL query by its connection-scoped id."""
async def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription:
"""Describe a submitted SQL query by its UUID or string ID.
For a handle with ``job_type == "query"``, pass its ``id`` on the
connection that submitted it. Registry job IDs are not query IDs.
"""
if isinstance(query_id, str):
query_id = UUID(query_id)
return await self._inner.describe_query(query_id)
async def namespace_client(self) -> LanceNamespace:
@@ -49,6 +49,9 @@ class OpenClipEmbeddings(EmbeddingFunction):
self.name, pretrained=self.pretrained
)
model.to(self.device)
# open_clip returns the model in training mode, which keeps dropout,
# stochastic depth and batch norm in their training behaviour.
model.eval()
self._model, self._preprocess = model, preprocess
self._tokenizer = open_clip.get_tokenizer(self.name)
self._ndims = None
+59 -21
View File
@@ -33,6 +33,7 @@ from typing import (
Annotated,
Any,
Callable,
Literal,
Optional,
Union,
get_args,
@@ -52,7 +53,9 @@ from pydantic import (
model_validator,
)
from .schema import blob as _blob_field
from .schema import is_blob_v2_field as _is_blob_v2_field
from ._blob_reader import BlobReader
from .secrets import EnvVarSecret
_Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1)
@@ -213,6 +216,7 @@ class FunctionParameter(_RemoteValue):
name: str
arrow_type: str
nullable: bool
blob_access: Optional[Literal["bytes", "reader"]] = None
class FunctionResultField(_OpenRemoteValue):
@@ -1096,6 +1100,43 @@ def _infer_signature(
if (input_schema is None) != (output_schema is None):
raise ValueError("input_schema and output_schema must be provided together")
try:
annotations = get_type_hints(function, include_extras=True)
except Exception as error:
raise TypeError(f"failed to resolve Function annotations: {error}") from error
def input_parameter(field):
annotation = annotations.get(field.name)
if get_origin(annotation) in (Union, types.UnionType):
choices = [
value for value in get_args(annotation) if value is not type(None)
]
if len(choices) == 1:
annotation = choices[0]
arrow_type = _canonical_arrow_field(field)
if annotation is BlobReader:
if arrow_type != _FUNCTION_BLOB_V2_TYPE:
raise TypeError(
f"BlobReader input '{field.name}' requires a Blob V2 field"
)
access = "reader"
elif annotation is bytes:
if arrow_type != _FUNCTION_BLOB_V2_TYPE and not (
pa.types.is_binary(field.type) or pa.types.is_large_binary(field.type)
):
raise TypeError(
f"bytes input '{field.name}' requires binary or Blob V2"
)
access = "bytes"
else:
access = "bytes" if arrow_type == _FUNCTION_BLOB_V2_TYPE else None
return FunctionParameter(
name=field.name,
arrow_type=arrow_type,
nullable=field.nullable,
blob_access=access,
)
if input_schema is not None:
if not isinstance(input_schema, pa.Schema):
raise TypeError("input_schema must be a PyArrow Schema")
@@ -1110,20 +1151,9 @@ def _infer_signature(
"input_schema fields must exactly match callable parameters in order: "
f"expected {expected!r}, got {actual!r}"
)
inputs = tuple(
FunctionParameter(
name=field.name,
arrow_type=_canonical_arrow_field(field),
nullable=field.nullable,
)
for field in input_schema
)
inputs = tuple(input_parameter(field) for field in input_schema)
return FunctionSignature(inputs=inputs, output=_function_output(output_schema))
try:
annotations = get_type_hints(function, include_extras=True)
except Exception as error:
raise TypeError(f"failed to resolve Function annotations: {error}") from error
missing = [
parameter.name for parameter in parameters if parameter.name not in annotations
]
@@ -1132,15 +1162,21 @@ def _infer_signature(
raise TypeError(f"missing Function annotations: {names!r}")
inputs = []
for parameter in parameters:
data_type, nullable = _annotation_type(annotations[parameter.name])
inputs.append(
FunctionParameter(
name=parameter.name,
arrow_type=_canonical_arrow_field(
pa.field(parameter.name, data_type, nullable=nullable)
),
nullable=nullable,
annotation = annotations[parameter.name]
base = annotation
nullable = False
if get_origin(base) in (Union, types.UnionType):
choices = [value for value in get_args(base) if value is not type(None)]
if len(choices) == 1:
base, nullable = choices[0], True
if base is BlobReader:
inputs.append(
input_parameter(_blob_field(parameter.name, nullable=nullable))
)
continue
data_type, nullable = _annotation_type(annotation)
inputs.append(
input_parameter(pa.field(parameter.name, data_type, nullable=nullable))
)
output_type, output_nullable = _annotation_type(annotations["return"])
if output_nullable:
@@ -1239,7 +1275,7 @@ _PYTHON_CALLABLE_ARTIFACT = "python_callable"
_PYTHON_BUNDLE_ARTIFACT = "python_bundle"
# Modules the Function image places beside user code.
_RESERVED_SOURCE_MODULES = frozenset(
{"user_function", "lance_source", "lance_callable"}
{"user_function", "lance_source", "lance_callable", "lance_blob"}
)
# The service accepts bundle paths made of these components only.
_SOURCE_PATH_COMPONENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
@@ -1499,6 +1535,8 @@ class _SourcePackager:
"""One module-level line that rebinds `name` to `value` in the artifact:
an import for modules and importable classes/functions, a literal
otherwise."""
if value is BlobReader:
return f"from lance_blob import BlobReader as {name}"
if isinstance(value, types.ModuleType):
if value.__name__.split(".")[0] in _NAMESPACE_MODULES:
raise ValueError(
+6 -2
View File
@@ -882,8 +882,12 @@ class IvfRq:
num_partitions: int, default sqrt(num_rows)
Number of IVF partitions to create.
num_bits: int, default 1
num_bits: int, optional
Number of bits to encode each dimension in the RabitQ codebook.
If None, uses Lance's default (currently 5), matching Rust and Node.
Set to 1 to retain the previous Python default and maximum compression.
Five-bit vector codes use roughly five times the storage of one-bit codes,
in addition to other index overhead.
max_iterations: int, default 50
Max iterations to train kmeans when computing IVF partitions.
@@ -897,7 +901,7 @@ class IvfRq:
distance_type: Literal["l2", "cosine", "dot"] = "l2"
num_partitions: Optional[int] = None
num_bits: int = 1
num_bits: Optional[int] = None
max_iterations: int = 50
sample_rate: int = 256
target_partition_size: Optional[int] = None
+17 -13
View File
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Handles to operations a server may run asynchronously."""
"""Handles to asynchronous registry jobs, SQL queries, and local operations."""
import asyncio
import json
@@ -44,12 +44,12 @@ class AsyncJob(Generic[T]):
@property
def id(self) -> Optional[str]:
"""Identifies the operation on the server that is running it.
"""The operation's opaque ID, interpreted according to :attr:`job_type`.
Returned for correlating with server logs or the jobs API. Operations
that run in this process have no server id and return `None`. The value
is opaque: parsing it or storing it to resume the job later is not
supported.
Registry jobs expose a registry job ID; SQL-backed jobs expose a
connection-scoped query ID as a string. Local operations return ``None``.
Query IDs cannot be used with registry job APIs or to resume after a
client restart.
"""
return self._inner.id if self._inner is not None else None
@@ -94,9 +94,11 @@ class AsyncJob(Generic[T]):
"""Ask the backend for this job's current state, and for a server-side
job its full record, then cache it for the properties below.
The properties are all `None` until this runs, because submitting an
operation returns only a job id. `status` fetches the whole record too;
`wait` records only the terminal state it establishes.
Registry details are `None` until this runs, because submitting an
operation does not fetch its full registry record. Query-backed jobs
expose only lifecycle state through this method. For registry jobs,
`status` fetches the whole record too; `wait` records only the terminal
state it establishes.
"""
if self._inner is None:
return
@@ -114,11 +116,13 @@ class AsyncJob(Generic[T]):
@property
def job_type(self) -> Optional[str]:
"""The job's type, as the server names it.
"""The job's kind of work, without contacting the backend.
`None` for an in-process job, which has no server-side record.
SQL-backed handles report ``"query"`` immediately. Registry handles
report the server's job type after :meth:`refresh` or :meth:`status`.
In-process handles return ``None``.
"""
return self._field("job_type")
return self._inner.job_type if self._inner is not None else None
@property
def creation_ms(self) -> Optional[int]:
@@ -294,7 +298,7 @@ class Job(Generic[T]):
@property
def job_type(self) -> Optional[str]:
"""The job's type. See :attr:`AsyncJob.job_type`."""
return self._field("job_type")
return self._inner.job_type if self._inner is not None else None
@property
def creation_ms(self) -> Optional[int]:
+171
View File
@@ -5,6 +5,7 @@
from __future__ import annotations
import asyncio
from typing import Any, AsyncIterator, Callable, Generic, Iterator, Optional, TypeVar
from .background_loop import LOOP
@@ -83,3 +84,173 @@ class Listing(Iterator[T], Generic[T]):
for initial, cached, final-page, and error behavior.
"""
return self._inner.page_token()
def validate_page_limit(limit: Optional[int]) -> None:
"""Validate the legacy limit argument before passing it to the bindings."""
if limit is not None and (
not isinstance(limit, int) or limit < 0 or limit > 2**31 - 1
):
raise ValueError("limit must be between 0 and 2147483647")
class _AsyncPagedNameListing(AsyncIterator[str]):
"""An async listing with an inspectable copy of its current backend page."""
def __init__(
self, inner: Any, page_token: Optional[str] = None, *, empty: bool = False
):
self._inner = inner
self._token = page_token
self._page: Optional[list[str]] = [] if empty else None
self._position = 0
self._done = empty
self._failure: Optional[BaseException] = None
self._lock = asyncio.Lock()
@property
def page_token(self) -> Optional[str]:
"""The continuation token for the next backend request."""
return self._token
def num_page_results(self) -> int:
"""Count the remaining names in the current page."""
return 0 if self._page is None else len(self._page) - self._position
def _page_copy(self) -> list[str]:
if self._page is None:
raise RuntimeError("Await or iterate the listing before reading its page")
return self._page.copy()
async def _fetch_page(self) -> None:
try:
page = await self._inner.next_page()
except asyncio.CancelledError:
raise
except BaseException as exc:
self._failure = exc
self._done = True
raise
if page is None:
self._page = []
self._token = None
self._done = True
else:
self._page = page["items"]
self._token = page["page_token"]
self._done = self._token is None
self._position = 0
async def _load_initial(self):
async with self._lock:
if self._failure is not None:
raise self._failure
if self._page is None:
await self._fetch_page()
return self
def __await__(self):
return self._load_initial().__await__()
def __aiter__(self):
return self
async def __anext__(self) -> str:
async with self._lock:
if self._failure is not None:
raise StopAsyncIteration
if self._page is None:
await self._fetch_page()
assert self._page is not None
while self._position >= len(self._page):
if self._done:
raise StopAsyncIteration
await self._fetch_page()
assert self._page is not None
item = self._page[self._position]
self._position += 1
return item
class AsyncTableListing(_AsyncPagedNameListing):
"""Awaitable table iterator; awaiting loads one page.
For example::
page = await db.list_tables(limit=20)
print(page.tables, page.page_token)
async for name in db.list_tables(limit=20):
print(name)
"""
@property
def tables(self) -> list[str]:
"""A copy of every table name in the most recently fetched page."""
return self._page_copy()
class AsyncNamespaceListing(_AsyncPagedNameListing):
"""Awaitable namespace iterator; awaiting loads one page.
For example::
page = await db.list_namespaces(limit=20)
print(page.namespaces, page.page_token)
"""
@property
def namespaces(self) -> list[str]:
"""A copy of every name in the most recently fetched page."""
return self._page_copy()
class _PagedNameListing(Iterator[str]):
def __init__(self, inner: _AsyncPagedNameListing):
self._inner = inner
LOOP.run(inner._load_initial())
def __iter__(self):
return self
def __next__(self) -> str:
try:
return LOOP.run(self._inner.__anext__())
except StopAsyncIteration:
raise StopIteration from None
@property
def page_token(self) -> Optional[str]:
return self._inner.page_token
def num_page_results(self) -> int:
return self._inner.num_page_results()
class TableListing(_PagedNameListing):
"""Table iterator retaining the current page as a list.
For example::
page = db.list_tables(limit=20)
print(page.tables, page.page_token)
for name in db.list_tables(limit=20):
print(name)
"""
@property
def tables(self) -> list[str]:
return self._inner.tables
class NamespaceListing(_PagedNameListing):
"""Namespace iterator retaining the current page as a list.
For example::
page = db.list_namespaces(limit=20)
print(page.namespaces, page.page_token)
"""
@property
def namespaces(self) -> list[str]:
return self._inner.namespaces
+23 -18
View File
@@ -52,7 +52,14 @@ from lancedb.background_loop import LOOP
from lancedb.arrow import AsyncRecordBatchReader
from lancedb.db import AsyncConnection, DBConnection
from lancedb.job import AsyncJob, Job
from lancedb.listing import AsyncListing, Listing
from lancedb.listing import (
AsyncListing,
AsyncNamespaceListing,
AsyncTableListing,
Listing,
NamespaceListing,
TableListing,
)
from lancedb.sql import AsyncQuery as AsyncSqlQuery
from lancedb.sql import QueryDescription, QueryParameters
from lance_namespace import (
@@ -61,8 +68,6 @@ from lance_namespace import (
CreateNamespaceResponse,
DescribeNamespaceResponse,
DropNamespaceResponse,
ListNamespacesResponse,
ListTablesResponse,
NamespaceExistsRequest,
TableExistsRequest,
)
@@ -875,7 +880,7 @@ class LanceNamespaceDBConnection(DBConnection):
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListNamespacesResponse:
) -> NamespaceListing:
"""
List child namespaces under the given namespace.
@@ -892,12 +897,12 @@ class LanceNamespaceDBConnection(DBConnection):
Returns
-------
ListNamespacesResponse
NamespaceListing
Response containing namespace names and optional page_token for pagination.
"""
if namespace_path is None:
namespace_path = []
return LOOP.run(
return NamespaceListing(
self._inner.list_namespaces(
namespace_path=namespace_path, page_token=page_token, limit=limit
)
@@ -1044,7 +1049,7 @@ class LanceNamespaceDBConnection(DBConnection):
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListTablesResponse:
) -> TableListing:
"""
List all tables in this database with pagination support.
@@ -1061,12 +1066,12 @@ class LanceNamespaceDBConnection(DBConnection):
Returns
-------
ListTablesResponse
TableListing
Response containing table names and optional page_token for pagination.
"""
if namespace_path is None:
namespace_path = []
return LOOP.run(
return TableListing(
self._inner.list_tables(
namespace_path=namespace_path, page_token=page_token, limit=limit
)
@@ -1483,12 +1488,12 @@ class AsyncLanceNamespaceDBConnection:
namespace_path = []
await self._inner.drop_all_tables(namespace_path=namespace_path)
async def list_namespaces(
def list_namespaces(
self,
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListNamespacesResponse:
) -> AsyncNamespaceListing:
"""
List child namespaces under the given namespace.
@@ -1505,12 +1510,12 @@ class AsyncLanceNamespaceDBConnection:
Returns
-------
ListNamespacesResponse
NamespaceListing
Response containing namespace names and optional page_token for pagination.
"""
if namespace_path is None:
namespace_path = []
return await self._inner.list_namespaces(
return self._inner.list_namespaces(
namespace_path=namespace_path, page_token=page_token, limit=limit
)
@@ -1640,12 +1645,12 @@ class AsyncLanceNamespaceDBConnection:
except TableNotFoundError:
return False
async def list_tables(
def list_tables(
self,
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListTablesResponse:
) -> AsyncTableListing:
"""
List all tables in this database with pagination support.
@@ -1662,12 +1667,12 @@ class AsyncLanceNamespaceDBConnection:
Returns
-------
ListTablesResponse
TableListing
Response containing table names and optional page_token for pagination.
"""
if namespace_path is None:
namespace_path = []
return await self._inner.list_tables(
return self._inner.list_tables(
namespace_path=namespace_path, page_token=page_token, limit=limit
)
@@ -1702,7 +1707,7 @@ class AsyncLanceNamespaceDBConnection:
parameters=parameters,
)
async def describe_query(self, query_id: UUID) -> QueryDescription:
async def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription:
"""Describe a submitted SQL query when supported."""
return await self._inner.describe_query(query_id)
+5 -3
View File
@@ -916,9 +916,11 @@ class Permutation:
f"Cannot rename column `{old_name}` to `{new_name}` because a column "
"with that name already exists"
)
new_selection = self.selection.copy()
new_selection[new_name] = new_selection[old_name]
del new_selection[old_name]
# Rebuild the mapping so the renamed column keeps its position
new_selection = {
(new_name if name == old_name else name): value
for name, value in self.selection.items()
}
return self._with_selection(new_selection)
def rename_columns(self, column_map: dict[str, str]) -> "Permutation":
+17 -12
View File
@@ -1561,9 +1561,10 @@ class LanceQueryBuilder(ABC):
ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance]
GlobalLimitExec: skip=0, fetch=10
FilterExec: _distance@2 IS NOT NULL
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
KNNVectorDistance: metric=l2
LanceRead: uri=..., projection=[vector], ...
LanceRead: uri=..., projection=[vector], num_fragments=..., ...
<BLANKLINE>
Parameters
----------
@@ -1619,8 +1620,8 @@ class LanceQueryBuilder(ABC):
GlobalLimitExec: elapsed=..., skip=0, fetch=10,
metrics=[output_rows=..., elapsed_compute=..., output_bytes=...]
FilterExec: elapsed=..., _distance@2 IS NOT NULL, metrics=[...]
SortExec: elapsed=..., TopK(fetch=10), expr=[...],
preserve_partitioning=[...],
SortExec: elapsed=..., TopK(fetch=1), expr=[...],
preserve_partitioning=[...], filter=[...],
metrics=[output_rows=..., elapsed_compute=...,
output_bytes=..., row_replacements=...]
KNNVectorDistance: elapsed=..., metric=l2,
@@ -1633,6 +1634,7 @@ class LanceQueryBuilder(ABC):
metrics=[output_rows=..., elapsed_compute=..., output_bytes=...,
fragments_scanned=..., ranges_scanned=1, rows_scanned=1,
bytes_read=..., iops=..., requests=..., task_wait_time=...]
<BLANKLINE>
Parameters
----------
@@ -2959,9 +2961,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance]
GlobalLimitExec: skip=0, fetch=10
FilterExec: _distance@2 IS NOT NULL
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
KNNVectorDistance: metric=l2
LanceRead: uri=..., projection=[vector], ...
LanceRead: uri=..., projection=[vector], num_fragments=..., ...
<BLANKLINE>
Parameters
----------
@@ -3601,9 +3604,10 @@ class AsyncQueryBase(object):
ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance]
GlobalLimitExec: skip=0, fetch=10
FilterExec: _distance@2 IS NOT NULL
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
KNNVectorDistance: metric=l2
LanceRead: uri=..., projection=[vector], ...
LanceRead: uri=..., projection=[vector], num_fragments=..., ...
<BLANKLINE>
<BLANKLINE>
Parameters
@@ -4521,9 +4525,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
LanceRead: uri=..., projection=[text], source=stream(_rowid)
GlobalLimitExec: skip=0, fetch=10
FilterExec: _distance@2 IS NOT NULL
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
KNNVectorDistance: metric=l2
LanceRead: uri=..., projection=[vector], ...
LanceRead: uri=..., projection=[vector], num_fragments=..., ...
ProjectionExec: expr=[vector@2 as vector, text@3 as text, _score@1 as _score, _rowid@0 as _rowid]
LanceRead: uri=..., projection=[vector, text], source=stream(_rowid)
GlobalLimitExec: skip=0, fetch=10
@@ -4877,9 +4881,10 @@ class BaseQueryBuilder(object):
ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance]
GlobalLimitExec: skip=0, fetch=10
FilterExec: _distance@2 IS NOT NULL
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
KNNVectorDistance: metric=l2
LanceRead: uri=..., projection=[vector], ...
LanceRead: uri=..., projection=[vector], num_fragments=..., ...
<BLANKLINE>
<BLANKLINE>
Parameters
+67 -25
View File
@@ -7,7 +7,6 @@ from __future__ import annotations
from dataclasses import replace
from datetime import timedelta
import json
import logging
from concurrent.futures import ThreadPoolExecutor
import sys
from typing import (
@@ -36,7 +35,7 @@ from lancedb.remote import ClientConfig, RetryConfig, TimeoutConfig, TlsConfig
import pyarrow as pa
from ..common import DATA
from lancedb.listing import Listing
from lancedb.listing import Listing, NamespaceListing, TableListing
from ..db import DBConnection, LOOP
from ..functions import FunctionVersion, UdfDefinition
from ..job import AsyncJob, Job
@@ -47,15 +46,13 @@ from ..secrets import EnvVarSecret, SecretInfo
from ..view import ViewDescription
if TYPE_CHECKING:
from .._lancedb import JobInfo
from .._lancedb import JobInfo, Session
from ..embeddings import EmbeddingFunctionConfig
from lance_namespace import (
LanceNamespace,
CreateNamespaceResponse,
DescribeNamespaceResponse,
DropNamespaceResponse,
ListNamespacesResponse,
ListTablesResponse,
)
from ..pydantic import LanceModel
from ..table import Table
@@ -242,6 +239,23 @@ class RemoteDBConnection(DBConnection):
def __repr__(self) -> str:
return f"RemoteConnect(name={self.db_name})"
@property
@override
def uri(self) -> str:
return self._conn.uri
@property
def session(self) -> Optional["Session"]:
"""Remote connections do not expose a local cache session."""
raise NotImplementedError("session is not supported by remote connections")
@property
def read_consistency_interval(self) -> Optional[timedelta]:
"""Remote connections do not expose a local read consistency interval."""
raise NotImplementedError(
"read_consistency_interval is not supported by remote connections"
)
@override
def serialize(self) -> str:
if getattr(self, "_catalog_oauth", False):
@@ -268,7 +282,7 @@ class RemoteDBConnection(DBConnection):
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListNamespacesResponse:
) -> NamespaceListing:
"""List immediate child namespace names in the given namespace.
Parameters
@@ -284,12 +298,12 @@ class RemoteDBConnection(DBConnection):
Returns
-------
ListNamespacesResponse
NamespaceListing
Response containing namespace names and optional page_token for pagination.
"""
if namespace_path is None:
namespace_path = []
return LOOP.run(
return NamespaceListing(
self._conn.list_namespaces(
namespace_path=namespace_path, page_token=page_token, limit=limit
)
@@ -379,7 +393,7 @@ class RemoteDBConnection(DBConnection):
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListTablesResponse:
) -> TableListing:
"""List a page of tables in this database.
Parameters
@@ -396,12 +410,12 @@ class RemoteDBConnection(DBConnection):
Returns
-------
ListTablesResponse
TableListing
Response containing table names and optional page_token for pagination.
"""
if namespace_path is None:
namespace_path = []
return LOOP.run(
return TableListing(
self._conn.list_tables(
namespace_path=namespace_path, page_token=page_token, limit=limit
)
@@ -487,18 +501,14 @@ class RemoteDBConnection(DBConnection):
if namespace_path is None:
namespace_path = []
if storage_options is not None:
logging.info(
"storage_options is ignored for remote tables"
" (storage is managed; set storage_options on connect() instead)"
table = LOOP.run(
self._conn.open_table(
name,
namespace_path=namespace_path,
storage_options=storage_options,
index_cache_size=index_cache_size,
)
if index_cache_size is not None:
logging.info(
"index_cache_size is ignored for remote tables"
" (there is no local cache to configure)"
)
table = LOOP.run(self._conn.open_table(name, namespace_path=namespace_path))
)
tbl = RemoteTable(
table,
self.db_name,
@@ -571,13 +581,16 @@ class RemoteDBConnection(DBConnection):
name: str,
data: DATA = None,
schema: Optional[Union[pa.Schema, LanceModel]] = None,
on_bad_vectors: str = "error",
fill_value: float = 0.0,
mode: Optional[str] = None,
exist_ok: bool = False,
on_bad_vectors: str = "error",
fill_value: float = 0.0,
embedding_functions: Optional[List[EmbeddingFunctionConfig]] = None,
*,
namespace_path: Optional[List[str]] = None,
storage_options: Optional[Dict[str, str]] = None,
data_storage_version: Optional[str] = None,
enable_v2_manifest_paths: Optional[bool] = None,
) -> Table:
"""Create a [Table][lancedb.table.Table] in the database.
@@ -617,6 +630,12 @@ class RemoteDBConnection(DBConnection):
embedding_functions: list of EmbeddingFunctionConfig, optional
The embedding functions to store in the table schema and use to
generate vectors when creating the table or adding data.
storage_options: dict, optional
Not supported for remote table creation. Raises NotImplementedError.
data_storage_version: str, optional
Not supported for remote table creation. Raises NotImplementedError.
enable_v2_manifest_paths: bool, optional
Not supported for remote table creation. Raises NotImplementedError.
Returns
-------
@@ -683,6 +702,14 @@ class RemoteDBConnection(DBConnection):
LanceTable(table4)
"""
for option, value in (
("data_storage_version", data_storage_version),
("enable_v2_manifest_paths", enable_v2_manifest_paths),
):
if value is not None:
raise NotImplementedError(
f"{option} is not supported for remote table creation"
)
if exist_ok:
if mode == "create":
mode = "exist_ok"
@@ -704,6 +731,7 @@ class RemoteDBConnection(DBConnection):
on_bad_vectors=on_bad_vectors,
fill_value=fill_value,
embedding_functions=embedding_functions,
storage_options=storage_options,
)
)
return RemoteTable(
@@ -843,6 +871,20 @@ class RemoteDBConnection(DBConnection):
job = LOOP.run(self._conn.drop_table_async(name, namespace_path=namespace_path))
return Job(job if isinstance(job, AsyncJob) else AsyncJob(job))
@override
def drop_all_tables(self, namespace_path: Optional[List[str]] = None):
"""Drop all tables from the database.
The remote API does not currently support this operation. Use
``list_tables()`` and ``drop_table()`` to drop tables individually.
"""
LOOP.run(self._conn.drop_all_tables(namespace_path=namespace_path))
@override
def drop_database(self):
"""Deprecated alias for ``drop_all_tables()``."""
LOOP.run(self._conn.drop_database())
@override
def rename_table(
self,
@@ -1122,7 +1164,7 @@ class RemoteDBConnection(DBConnection):
)
@override
def describe_query(self, query_id: UUID) -> QueryDescription:
def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription:
"""Describe a submitted SQL query by its connection-scoped id."""
return LOOP.run(
self._conn.describe_query(
+255 -150
View File
@@ -4,9 +4,10 @@
from datetime import timedelta
import deprecation
import logging
from functools import cached_property
from functools import cached_property, wraps
import os
from typing import (
TYPE_CHECKING,
Any,
Callable,
Dict,
@@ -41,19 +42,12 @@ from lancedb.index import (
FTS,
BTree,
Bitmap,
HnswFlat,
HnswPq,
HnswSq,
IvfFlat,
IvfPq,
IvfRq,
IvfSq,
LabelList,
)
from lancedb.job import Job
from lancedb.functions import FunctionApplication, RefreshColumnResult
from lancedb.remote.db import LOOP
from lancedb.table import IndexConfigType, KNOWN_METRICS
from lancedb.table import IndexConfigType, _maybe_add_fts_error_note
import pyarrow as pa
from lancedb.common import DATA, VEC, VECTOR_COLUMN_NAME
@@ -68,7 +62,6 @@ from ..query import (
FullTextQuery,
LanceQueryBuilder,
LanceTakeQueryBuilder,
LanceVectorQueryBuilder,
)
from ..table import (
AsyncTable,
@@ -79,9 +72,35 @@ from ..table import (
Table,
Tags,
)
from ..types import BaseTokenizerType
from ..types import BaseTokenizerType, ScalarIndexType
from ..util import infer_vector_column_name
if TYPE_CHECKING:
import lance
import polars as pl
def _deprecated_keyword(old_name: str, new_name: str):
"""Keep old remote-only keyword names without changing the Table signature."""
def decorate(method):
@wraps(method)
def wrapped(self, *args, **kwargs):
if old_name in kwargs:
if new_name in kwargs:
raise TypeError(f"Cannot specify both {old_name} and {new_name}")
warnings.warn(
f"{old_name} is deprecated; use {new_name} instead",
DeprecationWarning,
stacklevel=2,
)
kwargs[new_name] = kwargs.pop(old_name)
return method(self, *args, **kwargs)
return wrapped
return decorate
class RemoteTable(Table):
def __init__(
@@ -173,6 +192,18 @@ class RemoteTable(Table):
"""The name of the table"""
return self._name
@property
def namespace(self) -> List[str]:
return self._namespace_path
@deprecation.deprecated(
deprecated_in="0.33.1",
current_version=__version__,
details="Use update_field_metadata() instead.",
)
def replace_field_metadata(self, field_name: str, new_metadata: Dict[str, str]):
LOOP.run(self._table.replace_field_metadata(field_name, new_metadata))
def __repr__(self) -> str:
return f"RemoteTable({self.db_name}.{self.name})"
@@ -243,12 +274,32 @@ class RemoteTable(Table):
def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table:
"""to_arrow() is not supported for remote tables."""
validate_blob_mode(blob_mode)
raise NotImplementedError("to_arrow() is not supported for remote tables.")
raise NotImplementedError(
"to_arrow() is not supported for remote tables. "
"Use table.search().to_arrow() instead."
)
def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs):
"""to_pandas() is not supported for remote tables."""
validate_blob_mode(blob_mode)
raise NotImplementedError("to_pandas() is not supported for remote tables.")
raise NotImplementedError(
"to_pandas() is not supported for remote tables. "
"Use table.search().to_pandas() instead."
)
def to_lance(self, **kwargs) -> "lance.LanceDataset":
"""Direct Lance dataset access is not supported for remote tables.
Use ``table.search().to_arrow()`` to read query results through the server.
"""
return LOOP.run(self._table.to_lance(**kwargs))
def to_polars(self, batch_size=None) -> "pl.LazyFrame":
"""Use ``table.search().to_polars()`` to read remote query results."""
raise NotImplementedError(
"to_polars() is not supported for remote tables. "
"Use table.search().to_polars() instead."
)
def checkout(self, version: Union[int, str]):
result = LOOP.run(self._table.checkout(version))
@@ -290,9 +341,10 @@ class RemoteTable(Table):
self._table.tokenize(query, column=column, index_name=index_name)
)
def index_stats(self, index_uuid: str) -> Optional[IndexStatistics]:
@_deprecated_keyword("index_uuid", "index_name")
def index_stats(self, index_name: str) -> Optional[IndexStatistics]:
"""List all the stats of a specified index"""
return LOOP.run(self._table.index_stats(index_uuid))
return LOOP.run(self._table.index_stats(index_name))
@deprecation.deprecated(
deprecated_in="0.25.0",
@@ -302,43 +354,38 @@ class RemoteTable(Table):
def create_scalar_index(
self,
column: str,
index_type: Literal["BTREE", "BITMAP", "LABEL_LIST", "scalar"] = "scalar",
*,
replace: bool = False,
replace: bool = True,
index_type: ScalarIndexType = "BTREE",
wait_timeout: Optional[timedelta] = None,
name: Optional[str] = None,
):
"""Creates a scalar index.
"""Create a scalar index on a column.
.. deprecated:: 0.25.0
Use :meth:`create_index` with a BTree, Bitmap, or LabelList config instead.
Example: ``table.create_index("column", config=BTree())``
Parameters
----------
column : str
The column to be indexed. Must be a boolean, integer, float,
or string column.
index_type : str
The index type of the scalar index. Must be "scalar" (BTREE),
"BTREE", "BITMAP", or "LABEL_LIST",
replace : bool
If True, replace the existing index with the new one.
"""
if index_type == "scalar" or index_type == "BTREE":
if index_type == "scalar":
warnings.warn(
'index_type="scalar" is deprecated; use "BTREE" instead',
DeprecationWarning,
stacklevel=2,
)
index_type = "BTREE"
if index_type == "BTREE":
config = BTree()
elif index_type == "BITMAP":
config = Bitmap()
elif index_type == "LABEL_LIST":
config = LabelList()
else:
raise ValueError(f"Unknown index type: {index_type}")
LOOP.run(
raise ValueError(f"Unknown index type {index_type}")
return LOOP.run(
self._table.create_index(
column,
config=config,
replace=replace,
config=config,
wait_timeout=wait_timeout,
name=name,
)
@@ -349,12 +396,16 @@ class RemoteTable(Table):
current_version=__version__,
details="Use create_index() with config=FTS() instead.",
)
@_deprecated_keyword("column", "field_names")
def create_fts_index(
self,
column: str,
field_names: Union[str, List[str]],
*,
ordering_field_names: Optional[Union[str, List[str]]] = None,
replace: bool = False,
wait_timeout: Optional[timedelta] = None,
writer_heap_size: Optional[int] = 1024 * 1024 * 1024,
use_tantivy: bool = False,
tokenizer_name: Optional[str] = None,
with_position: bool = False,
# tokenizer configs:
base_tokenizer: BaseTokenizerType = "simple",
@@ -370,6 +421,7 @@ class RemoteTable(Table):
prefix_only: bool = False,
block_size: int = 128,
document_granularity: DocumentGranularity = DocumentGranularity.ROW,
wait_timeout: Optional[timedelta] = None,
name: Optional[str] = None,
):
"""Create a full-text search index on a column.
@@ -378,31 +430,70 @@ class RemoteTable(Table):
Use :meth:`create_index` with an FTS config instead.
Example: ``table.create_index("text_column", config=FTS())``
"""
if use_tantivy:
raise ValueError(
"Tantivy-based FTS has been removed. "
"Remove use_tantivy and recreate the index with native FTS."
)
if ordering_field_names is not None:
raise ValueError(
"ordering_field_names was only supported by the removed "
"Tantivy-based FTS implementation."
)
if writer_heap_size != 1024 * 1024 * 1024:
raise ValueError(
"writer_heap_size was only supported by the removed "
"Tantivy-based FTS implementation."
)
if not isinstance(field_names, str):
raise ValueError(
"Native FTS indexes can only be created on a single field "
"at a time. To search over multiple text fields, create a "
"separate FTS index for each field."
)
if tokenizer_name is None:
tokenizer_configs = {
"base_tokenizer": base_tokenizer,
"language": language,
"with_position": with_position,
"max_token_length": max_token_length,
"lower_case": lower_case,
"stem": stem,
"remove_stop_words": remove_stop_words,
"custom_stop_words": custom_stop_words,
"ascii_folding": ascii_folding,
"ngram_min_length": ngram_min_length,
"ngram_max_length": ngram_max_length,
"prefix_only": prefix_only,
}
else:
tokenizer_configs = self.infer_tokenizer_configs(tokenizer_name)
tokenizer_configs["custom_stop_words"] = custom_stop_words
config = FTS(
with_position=with_position,
base_tokenizer=base_tokenizer,
language=language,
max_token_length=max_token_length,
lower_case=lower_case,
stem=stem,
remove_stop_words=remove_stop_words,
custom_stop_words=custom_stop_words,
ascii_folding=ascii_folding,
ngram_min_length=ngram_min_length,
ngram_max_length=ngram_max_length,
prefix_only=prefix_only,
block_size=block_size,
document_granularity=document_granularity,
**tokenizer_configs,
)
LOOP.run(
self._table.create_index(
column,
config=config,
replace=replace,
wait_timeout=wait_timeout,
name=name,
try:
LOOP.run(
self._table.create_index(
field_names,
replace=replace,
config=config,
name=name,
wait_timeout=wait_timeout,
)
)
)
except (ValueError, RuntimeError) as e:
_maybe_add_fts_error_note(
e,
base_tokenizer=config.base_tokenizer,
language=config.language,
)
raise e
# New unified API overload
@overload
@@ -412,6 +503,7 @@ class RemoteTable(Table):
/,
*,
config: IndexConfigType,
replace: bool = ...,
wait_timeout: Optional[timedelta] = ...,
name: Optional[str] = ...,
train: bool = ...,
@@ -422,38 +514,56 @@ class RemoteTable(Table):
def create_index(
self,
metric: Literal["l2", "cosine", "dot", "hamming"] = ...,
vector_column_name: str = ...,
index_cache_size: Optional[int] = ...,
num_partitions: Optional[int] = ...,
num_sub_vectors: Optional[int] = ...,
replace: Optional[bool] = ...,
vector_column_name: str = ...,
replace: bool = ...,
accelerator: Optional[str] = ...,
index_type: Literal[
"VECTOR", "IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ"
] = ...,
wait_timeout: Optional[timedelta] = ...,
*,
index_cache_size: Optional[int] = ...,
num_bits: int = ...,
index_type: Literal[
"IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_RQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ"
] = ...,
max_iterations: int = ...,
sample_rate: int = ...,
m: int = ...,
ef_construction: int = ...,
*,
wait_timeout: Optional[timedelta] = ...,
name: Optional[str] = ...,
train: bool = ...,
target_partition_size: Optional[int] = ...,
) -> None: ...
def create_index(
self,
metric: str = "l2",
vector_column_name: str = VECTOR_COLUMN_NAME,
index_cache_size: Optional[int] = None,
num_partitions: Optional[int] = None,
num_sub_vectors: Optional[int] = None,
replace: Optional[bool] = None,
vector_column_name: str = VECTOR_COLUMN_NAME,
replace: bool = True,
accelerator: Optional[str] = None,
index_type="vector",
wait_timeout: Optional[timedelta] = None,
*,
index_cache_size: Optional[int] = None,
num_bits: int = 8,
index_type: Literal[
"IVF_FLAT",
"IVF_SQ",
"IVF_PQ",
"IVF_RQ",
"IVF_HNSW_SQ",
"IVF_HNSW_PQ",
"IVF_HNSW_FLAT",
] = "IVF_PQ",
max_iterations: int = 50,
sample_rate: int = 256,
m: int = 20,
ef_construction: int = 300,
*,
config: Optional[IndexConfigType] = None,
wait_timeout: Optional[timedelta] = None,
name: Optional[str] = None,
train: bool = True,
target_partition_size: Optional[int] = None,
):
"""Create an index on a column.
@@ -464,6 +574,24 @@ class RemoteTable(Table):
argument plus separate ``vector_column_name`` / ``num_partitions`` /
etc. parameters, and emits a ``DeprecationWarning``.
Parameters
----------
metric : str
For new API: the column name to index.
For legacy API: the distance metric ("l2", "cosine", "dot", "hamming").
config : IndexConfigType, optional
The index configuration object. If provided, uses the new unified API.
Can be one of: IvfFlat, IvfPq, IvfSq, IvfRq, HnswPq, HnswSq,
BTree, Bitmap, LabelList, Fm, FTS.
replace : bool, default True
Whether to replace an existing index on this column.
wait_timeout : timedelta, optional
Timeout to wait for async indexing to complete.
name : str, optional
Custom name for the index.
train : bool, default True
Whether to train the index with existing data.
Examples
--------
New API (recommended):
@@ -489,7 +617,6 @@ class RemoteTable(Table):
vector_column_name,
accelerator,
index_cache_size,
replace,
)
if is_legacy:
@@ -505,6 +632,7 @@ class RemoteTable(Table):
stacklevel=2,
)
# Legacy API: first arg is the distance metric
column = vector_column_name
if accelerator is not None:
@@ -513,45 +641,27 @@ class RemoteTable(Table):
"and will be ignored."
)
idx_type = index_type.upper()
if idx_type == "VECTOR" or idx_type == "IVF_PQ":
config = IvfPq(
distance_type=metric,
num_partitions=num_partitions,
num_sub_vectors=num_sub_vectors,
num_bits=num_bits,
)
elif idx_type == "IVF_RQ":
config = IvfRq(
distance_type=metric,
num_partitions=num_partitions,
num_bits=num_bits,
)
elif idx_type == "IVF_SQ":
config = IvfSq(distance_type=metric, num_partitions=num_partitions)
elif idx_type == "IVF_HNSW_PQ":
config = HnswPq(
distance_type=metric,
num_partitions=num_partitions,
num_sub_vectors=num_sub_vectors,
num_bits=num_bits,
)
elif idx_type == "IVF_HNSW_SQ":
config = HnswSq(distance_type=metric, num_partitions=num_partitions)
elif idx_type == "IVF_HNSW_FLAT":
config = HnswFlat(distance_type=metric, num_partitions=num_partitions)
elif idx_type == "IVF_FLAT":
config = IvfFlat(distance_type=metric, num_partitions=num_partitions)
else:
raise ValueError(
f"Unknown vector index type: {idx_type}. Valid options are"
" 'IVF_FLAT', 'IVF_PQ', 'IVF_RQ', 'IVF_SQ',"
" 'IVF_HNSW_PQ', 'IVF_HNSW_SQ', 'IVF_HNSW_FLAT'"
)
# Build config from legacy parameters
config = self._build_vector_config_from_legacy_params(
metric=metric,
index_type="IVF_PQ"
if index_type.upper() == "VECTOR"
else index_type.upper(),
num_partitions=num_partitions,
num_sub_vectors=num_sub_vectors,
num_bits=num_bits,
max_iterations=max_iterations,
sample_rate=sample_rate,
m=m,
ef_construction=ef_construction,
target_partition_size=target_partition_size,
accelerator=None,
)
else:
column = metric
LOOP.run(
return LOOP.run(
self._table.create_index(
column,
replace=replace,
@@ -590,37 +700,6 @@ class RemoteTable(Table):
)
)
def _is_legacy_create_index_call(
self,
first_arg: str,
config: Optional[IndexConfigType],
num_partitions: Optional[int],
num_sub_vectors: Optional[int],
vector_column_name: str,
accelerator: Optional[str],
index_cache_size: Optional[int],
replace: Optional[bool],
) -> bool:
"""Detect if this is a legacy create_index call."""
if config is not None:
return False
if any(
x is not None
for x in (
num_partitions,
num_sub_vectors,
accelerator,
index_cache_size,
replace,
)
):
return True
if vector_column_name != VECTOR_COLUMN_NAME:
return True
if first_arg.lower() in KNOWN_METRICS:
return True
return False
def add(
self,
data: DATA,
@@ -692,9 +771,10 @@ class RemoteTable(Table):
query: Union[VEC, str] = None,
vector_column_name: Optional[str] = None,
query_type="auto",
ordering_field_name: Optional[str] = None,
fts_columns: Optional[Union[str, List[str]]] = None,
fast_search: Optional[bool] = None,
) -> LanceVectorQueryBuilder:
) -> LanceQueryBuilder:
"""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/)
@@ -785,6 +865,7 @@ class RemoteTable(Table):
query,
query_type,
vector_column_name=vector_column_name,
ordering_field_name=ordering_field_name,
fts_columns=fts_columns or [],
fast_search=fast_search,
)
@@ -844,7 +925,8 @@ class RemoteTable(Table):
self._table._do_merge(merge, new_data, on_bad_vectors, fill_value)
)
def delete(self, predicate: str) -> DeleteResult:
@_deprecated_keyword("predicate", "where")
def delete(self, where: Union[str, Expr]) -> DeleteResult:
"""Delete rows from the table.
This can be used to delete a single row, many rows, all rows, or
@@ -852,7 +934,7 @@ class RemoteTable(Table):
Parameters
----------
predicate: str
where: Union[str, Expr]
The SQL where clause to use when deleting rows.
- For example, 'x = 2' or 'x IN (1, 2, 3)'.
@@ -896,7 +978,7 @@ class RemoteTable(Table):
x vector _distance # doctest: +SKIP
0 2 [3.0, 4.0] 85.0 # doctest: +SKIP
"""
return LOOP.run(self._table.delete(predicate))
return LOOP.run(self._table.delete(where))
def update(
self,
@@ -958,7 +1040,12 @@ class RemoteTable(Table):
self._table.update(where=where, updates=values, updates_sql=values_sql)
)
def cleanup_old_versions(self, *_):
def cleanup_old_versions(
self,
older_than: Optional[timedelta] = None,
*,
delete_unverified: bool = False,
):
"""
cleanup_old_versions() is a no-op for remote tables.
"""
@@ -968,7 +1055,7 @@ class RemoteTable(Table):
)
pass
def compact_files(self, *_):
def compact_files(self, *args, **kwargs):
"""
compact_files() is a no-op for remote tables.
"""
@@ -983,12 +1070,14 @@ class RemoteTable(Table):
*,
cleanup_older_than: Optional[timedelta] = None,
delete_unverified: bool = False,
retrain: bool = False,
):
"""optimize() is not supported for remote tables."""
return LOOP.run(
self._table.optimize(
cleanup_older_than=cleanup_older_than,
delete_unverified=delete_unverified,
retrain=retrain,
)
)
@@ -1096,8 +1185,9 @@ class RemoteTable(Table):
"""No-op for remote tables (no local shard writers)."""
return LOOP.run(self._table.close_rte_writers())
def drop_index(self, index_name: str):
return LOOP.run(self._table.drop_index(index_name))
@_deprecated_keyword("index_name", "name")
def drop_index(self, name: str) -> None:
return LOOP.run(self._table.drop_index(name))
def prewarm_index(self, name: str) -> None:
"""Prewarm an index in the table.
@@ -1154,11 +1244,21 @@ class RemoteTable(Table):
"""
return LOOP.run(self._table.uri())
def take_offsets(self, offsets: list[int]) -> LanceTakeQueryBuilder:
return LanceTakeQueryBuilder(self._table.take_offsets(offsets))
def take_offsets(
self, offsets: list[int], *, with_row_id: bool = False
) -> LanceTakeQueryBuilder:
builder = LanceTakeQueryBuilder(self._table.take_offsets(offsets))
if with_row_id:
builder.with_row_id()
return builder
def take_row_ids(self, row_ids: list[int]) -> LanceTakeQueryBuilder:
return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids))
def take_row_ids(
self, row_ids: list[int], *, with_row_id: bool = False
) -> LanceTakeQueryBuilder:
builder = LanceTakeQueryBuilder(self._table.take_row_ids(row_ids))
if with_row_id:
builder.with_row_id()
return builder
def uses_v2_manifest_paths(self) -> bool:
return LOOP.run(self._table.uses_v2_manifest_paths())
@@ -1186,6 +1286,11 @@ class RemoteTable(Table):
) -> "list[Optional[BlobFile]]":
return LOOP.run(self._table.fetch_blob_files(column, row_ids))
def open_blob_file(
self, descriptor: Union[dict, pa.StructScalar, None]
) -> Optional[BlobFile]:
return LOOP.run(self._table.open_blob_file(descriptor))
def head(self, n=5) -> pa.Table:
"""
Return the first `n` rows of the table.
+3 -4
View File
@@ -21,7 +21,7 @@ class OpenaiReranker(Reranker):
Parameters
----------
model_name : str, default "gpt-4-turbo-preview"
model_name : str, default "gpt-5.6-sol"
The name of the cross encoder model to use.
column : str, default "text"
The name of the column to use as input to the cross encoder model.
@@ -33,7 +33,7 @@ class OpenaiReranker(Reranker):
def __init__(
self,
model_name: str = "gpt-4-turbo-preview",
model_name: str = "gpt-5.6-sol",
column: str = "text",
return_score="relevance",
api_key: Optional[str] = None,
@@ -54,7 +54,6 @@ class OpenaiReranker(Reranker):
response = self._client.chat.completions.create(
model=self.model_name,
response_format={"type": "json_object"},
temperature=0,
messages=[
{
"role": "system",
@@ -132,4 +131,4 @@ class OpenaiReranker(Reranker):
"OPENAI_API_KEY not set. Either set it in your environment or \
pass it as `api_key` argument to the CohereReranker."
)
return openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY") or self.api_key)
return openai.OpenAI(api_key=self.api_key or os.environ.get("OPENAI_API_KEY"))
+38 -24
View File
@@ -4,6 +4,7 @@
"""Schema helpers for Lance blob columns."""
import ctypes
import importlib
from typing import TYPE_CHECKING, Optional
@@ -16,6 +17,7 @@ if TYPE_CHECKING:
_BLOB_EXTENSION_NAME = "lance.blob.v2"
_BLOB_V1_KEY = "lance-encoding:blob"
_ARROW_EXT_NAME_KEY = "ARROW:extension:name"
_MAX_RUST_USIZE = ctypes.c_size_t(-1).value
_BLOB_V2_STORAGE_TYPE = pa.struct(
[
pa.field("data", pa.large_binary(), nullable=True),
@@ -209,32 +211,44 @@ def blob(
Maximum size in bytes of a single packed blob sidecar file before
a new one is started. Must be positive.
The thresholds are forwarded to ``lance.blob_field``, which validates
them and stores them as field metadata, so pylance is required whenever
any threshold is set. When all thresholds are omitted the behavior is
unchanged.
Thresholds are validated and stored as Arrow field metadata, with the
same format as ``lance.blob_field``. They do not require pylance.
"""
if (
inline_size_threshold is None
and dedicated_size_threshold is None
and pack_file_size_threshold is None
metadata = {}
for threshold, value, allow_zero, key in (
(
"inline_size_threshold",
inline_size_threshold,
True,
"lance-encoding:blob-inline-size-threshold",
),
(
"dedicated_size_threshold",
dedicated_size_threshold,
False,
"lance-encoding:blob-dedicated-size-threshold",
),
(
"pack_file_size_threshold",
pack_file_size_threshold,
False,
"lance-encoding:blob-pack-file-size-threshold",
),
):
blob_type = _resolve_blob_type()
return pa.field(name, blob_type(), nullable=nullable)
try:
lance = importlib.import_module("lance")
except ModuleNotFoundError as err:
raise ValueError(
"Setting blob tier thresholds requires pylance to be installed"
) from err
_resolve_blob_type()
return lance.blob_field(
name,
nullable=nullable,
inline_size_threshold=inline_size_threshold,
dedicated_size_threshold=dedicated_size_threshold,
pack_file_size_threshold=pack_file_size_threshold,
)
if value is None:
continue
if isinstance(value, bool) or not isinstance(value, int):
raise TypeError(f"{threshold} must be an int, got {type(value).__name__}")
if value < 0 or (value == 0 and not allow_zero):
bound = "non-negative" if allow_zero else "positive"
raise ValueError(f"{threshold} must be {bound}")
if value > _MAX_RUST_USIZE:
raise OverflowError(f"{threshold} must fit in a Rust usize")
metadata[key] = str(value)
blob_type = _resolve_blob_type()
field = pa.field(name, blob_type(), nullable=nullable)
return field.with_metadata(metadata) if metadata else field
def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType:
File diff suppressed because it is too large Load Diff
+9
View File
@@ -452,6 +452,15 @@ def _(value: pa.Scalar):
return value_to_sql(value.as_py())
@value_to_sql.register(pa.MapScalar)
def _(value: pa.MapScalar):
if not value.is_valid:
return value_to_sql(None)
keys = value_to_sql(list(value.values.field(0)))
values = value_to_sql(list(value.values.field(1)))
return f"map({keys}, {values})"
@value_to_sql.register(np.ndarray)
def _(value: np.ndarray):
return value_to_sql(value.tolist())
+2 -14
View File
@@ -91,19 +91,13 @@ def test_quickstart(tmp_path):
}
)
# --8<-- [end:alter_columns_vector]
# Change it back since we can get a panic with fp16
tbl.alter_columns(
{
"path": "vector",
"data_type": pa.list_(pa.float32(), list_size=2),
}
)
# --8<-- [start:drop_columns]
tbl.drop_columns(["dbl_price"])
# --8<-- [end:drop_columns]
# --8<-- [start:create_index]
tbl.create_index(num_sub_vectors=1)
# --8<-- [end:create_index]
tbl.search([100, 100]).limit(2).to_pandas()
# --8<-- [start:delete_rows]
tbl.delete("item = 'fizz'")
# --8<-- [end:delete_rows]
@@ -185,13 +179,6 @@ async def test_quickstart_async(tmp_path):
}
)
# --8<-- [end:alter_columns_async_vector]
# Change it back since we can get a panic with fp16
await tbl.alter_columns(
{
"path": "vector",
"data_type": pa.list_(pa.float32(), list_size=2),
}
)
# --8<-- [start:drop_columns_async]
await tbl.drop_columns(["dbl_price"])
# --8<-- [end:drop_columns_async]
@@ -200,6 +187,7 @@ async def test_quickstart_async(tmp_path):
# --8<-- [start:create_index_async]
await tbl.create_index("vector")
# --8<-- [end:create_index_async]
await tbl.vector_search([100, 100]).limit(2).to_pandas()
# --8<-- [start:delete_rows_async]
await tbl.delete("item = 'fizz'")
# --8<-- [end:delete_rows_async]
+169
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import ctypes
import io
import subprocess
import sys
@@ -14,6 +15,7 @@ import pytest
from lance.blob import Blob, BlobType as LanceBlobType
import lancedb
from lancedb import schema as schema_module
from lancedb._blob import (
_set_blob_column,
blob_v2_projection_sources,
@@ -139,6 +141,41 @@ def test_blob_without_thresholds_has_no_threshold_metadata():
assert not any(b"threshold" in key for key in metadata)
@pytest.mark.parametrize("nullable", [False, True])
@pytest.mark.parametrize(
"thresholds",
[
{},
{"inline_size_threshold": 0},
{"inline_size_threshold": 8192},
{"dedicated_size_threshold": 1 << 20},
{"pack_file_size_threshold": 1 << 26},
{
"inline_size_threshold": 8192,
"dedicated_size_threshold": 1 << 20,
"pack_file_size_threshold": 1 << 26,
},
{
"inline_size_threshold": ctypes.c_size_t(-1).value,
"dedicated_size_threshold": ctypes.c_size_t(-1).value,
"pack_file_size_threshold": ctypes.c_size_t(-1).value,
},
],
)
def test_blob_thresholds_match_pylance(monkeypatch, nullable, thresholds):
expected = lance.blob_field("payload", nullable=nullable, **thresholds)
actual = lancedb.blob("payload", nullable=nullable, **thresholds)
assert type(actual.type) is LanceBlobType
assert actual.equals(expected, check_metadata=True)
# Both extension classes must encode the same Arrow field on the wire.
monkeypatch.setattr(
schema_module, "_resolve_blob_type", lambda: schema_module._FallbackBlobType
)
fallback = lancedb.blob("payload", nullable=nullable, **thresholds)
assert pa.schema([fallback]).serialize() == pa.schema([expected]).serialize()
def test_blob_thresholds_are_validated():
with pytest.raises(ValueError, match="inline_size_threshold must be non-negative"):
lancedb.blob("payload", inline_size_threshold=-1)
@@ -2659,3 +2696,135 @@ def test_create_table_with_python_blob_dict_and_object():
b"from-dict",
b"from-blob",
]
def _openable_descriptors(table, *, by_address):
hits = (
table.search()
.select(["id", "image", "_rowaddr"])
.with_row_id(True)
.limit(100)
.to_arrow(blob_mode="descriptions")
)
version = table.version
committed_at = next(
item["timestamp"]
for item in table.list_versions()
if item["version"] == version
)
descriptors = {}
for row in hits.to_pylist():
descriptor = row["image"]
if descriptor is not None and descriptor.get("kind") is not None:
descriptor = dict(descriptor)
descriptor["source"] = {
"uri": table.to_lance().uri,
"version": version,
"committed_at_micros": round(committed_at.timestamp() * 1_000_000),
"branch": None,
"column": "image",
"row_reference": row["_rowaddr" if by_address else "_rowid"],
"row_is_address": by_address,
}
else:
descriptor = None
descriptors[row["id"]] = descriptor
return descriptors
@pytest.mark.parametrize("by_address", [False, True])
def test_open_blob_files_preserve_snapshot_order_nulls_and_empty_values(
tmp_path, by_address
):
db = lancedb.connect(str(tmp_path))
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("videos", schema=schema)
table.add(
[
{"id": 1, "image": b"alpha"},
{"id": 2, "image": None},
{"id": 3, "image": b"beta"},
{"id": 4, "image": b""},
]
)
descriptors = _openable_descriptors(table, by_address=by_address)
table.add([{"id": 1, "image": b"replacement"}], mode="overwrite")
latest_version = table.version
values = pa.array(
[descriptors[3], descriptors[2], descriptors[1], descriptors[3], descriptors[4]]
)
files = table.open_blob_files(
pa.chunked_array([values.slice(0, 2), values.slice(2)])
)
assert iter(files) is files
actual = []
for file in files:
if file is None:
actual.append(None)
else:
with file:
actual.append(file.read())
assert file.closed
assert actual == [b"beta", None, b"alpha", b"beta", b""]
with table.open_blob_file(descriptors[1]) as file:
file.seek(2)
assert file.read_range(0, 2) == b"al"
assert file.tell() == 2
assert file.read(2) == b"ph"
assert table.version == latest_version
assert table.open_blob_file(pa.scalar(None, type=values.type)) is None
other = db.create_table("other", schema=schema)
with pytest.raises(ValueError, match="different table"):
other.open_blob_file(descriptors[1])
stale = dict(descriptors[1], source=dict(descriptors[1]["source"]))
stale["source"]["committed_at_micros"] += 1
with pytest.raises(ValueError, match="snapshot has been replaced"):
table.open_blob_file(stale)
raw = {key: value for key, value in descriptors[1].items() if key != "source"}
with pytest.raises(ValueError, match="source"):
table.open_blob_file(raw)
with pytest.raises(TypeError, match="descriptor"):
table.open_blob_file("not a descriptor")
def test_open_blob_files_does_not_open_ahead(monkeypatch):
table = _blob_table("lazy_descriptor_iteration", [{"id": 1, "image": b"x"}])
opened = []
def open_file(descriptor):
opened.append(descriptor)
return descriptor
monkeypatch.setattr(table, "open_blob_file", open_file)
files = table.open_blob_files(iter([1, None, 2]))
assert opened == []
assert next(files) == 1
assert opened == [1]
assert list(files) == [None, 2]
assert opened == [1, None, 2]
@pytest.mark.asyncio
async def test_open_blob_files_async_iterates_snapshot_handles(tmp_path):
db = lancedb.connect(str(tmp_path))
table = db.create_table(
"videos",
schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]),
)
table.add([{"id": 1, "image": b"alpha"}, {"id": 2, "image": None}])
descriptors = _openable_descriptors(table, by_address=True)
async_db = await lancedb.connect_async(str(tmp_path))
videos = await async_db.open_table("videos")
file = await videos.open_blob_file(pa.scalar(descriptors[1]))
with file:
assert file.read(2) == b"al"
actual = []
async for file in videos.open_blob_files(
[descriptors[2], descriptors[1], descriptors[1]]
):
if file is None:
actual.append(None)
else:
with file:
actual.append(await file.aread())
assert actual == [None, b"alpha", b"alpha"]
+59
View File
@@ -0,0 +1,59 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import io
import pytest
from lancedb import BlobReader
def test_reader_is_lazy_and_seekable():
requests = []
size = 64 * 1024 * 1024
def read_range(offset, length):
requests.append((offset, length))
return bytes((index % 251 for index in range(offset, offset + length)))
with BlobReader(read_range, size) as reader:
assert requests == []
assert reader.size() == size
reader.seek(-4, io.SEEK_END)
assert reader.read(10) == bytes(index % 251 for index in range(size - 4, size))
assert reader.read(1) == b""
assert requests == [(size - 4, 4)]
assert reader.read_range(1, 3) == b"\x01\x02\x03"
assert reader.tell() == size
reader.seek(2)
target = bytearray(3)
assert reader.readinto(target) == 3
assert target == b"\x02\x03\x04"
reader.seek(-1, io.SEEK_CUR)
assert reader.tell() == 4
with pytest.raises(ValueError, match="closed"):
reader.read(1)
def test_reader_supports_buffered_io_and_bounds_range_requests():
requests = []
payload = b"x" * (3 * 1024 * 1024 + 7)
def read_range(offset, length):
requests.append((offset, length))
return payload[offset : offset + length]
with io.BufferedReader(BlobReader(read_range, len(payload))) as reader:
assert reader.read() == payload
assert requests
assert all(length <= 1024 * 1024 for _, length in requests)
def test_reader_rejects_invalid_ranges_and_short_reads():
reader = BlobReader(lambda offset, length: b"", 10)
with pytest.raises(ValueError):
reader.seek(-1)
with pytest.raises(ValueError):
reader.read_range(-1, 1)
with pytest.raises(OSError, match="payload length"):
reader.read(1)
+33
View File
@@ -1039,6 +1039,39 @@ def test_open_clip_opens_percent_encoded_file_uri(tmp_path):
assert image.size == (4, 4)
def test_open_clip_model_in_eval_mode():
"""
open_clip.create_model_and_transforms returns the model in training mode,
which leaves dropout, stochastic depth and batch norm in their training
behaviour. OpenCLIP must switch the model to eval mode before embedding.
"""
from lancedb.embeddings.open_clip import OpenClipEmbeddings
class FakeModel:
training = True
def to(self, device):
return self
def eval(self):
self.training = False
return self
open_clip = MagicMock()
open_clip.create_model_and_transforms.return_value = (FakeModel(), None, None)
with patch(
"lancedb.embeddings.open_clip.attempt_import_or_raise",
return_value=open_clip,
):
func = OpenClipEmbeddings(name="ViTamin-B", pretrained="datacomp1b")
open_clip.create_model_and_transforms.assert_called_once_with(
"ViTamin-B", pretrained="datacomp1b"
)
assert func._model.training is False
def test_jina_generate_image_input_dict_local_path(tmp_path):
"""
JinaEmbeddings._generate_image_input_dict must accept a local image path
+21
View File
@@ -3,7 +3,28 @@
import pickle
import pytest
from lancedb.db import DBConnection
from lancedb.remote.errors import HttpError, LanceDBClientError, RetryError
from lancedb.table import Table
@pytest.mark.parametrize(
"base, method",
[
(Table, "to_lance"),
(Table, "to_polars"),
(DBConnection, "drop_all_tables"),
(DBConnection, "drop_database"),
],
)
def test_base_unsupported_methods_have_messages(base, method):
with pytest.raises(NotImplementedError) as exc_info:
getattr(base, method)(None)
message = str(exc_info.value)
assert method in message
assert "not supported" in message
def test_pickle_lancedb_client_error():
@@ -23,6 +23,7 @@ import pyarrow as pa
import pytest
import lancedb
from lancedb import BlobReader
from lancedb.functions import (
PythonRuntimeSpec,
SecretBinding,
@@ -966,6 +967,60 @@ def test_blob_fields_use_the_scalar_function_semantic_type():
assert signature.output.arrow_type == "blob_v2"
def test_blob_input_access_is_inferred_from_annotations():
@udf
def read_blob(video: BlobReader, timestamp_s: float) -> bytes:
video.seek(int(timestamp_s))
return video.read(4)
@udf
def read_bytes(video: bytes) -> int:
return len(video)
reader = read_blob.registration_request.signature
assert reader.inputs[0].arrow_type == "blob_v2"
assert reader.inputs[0].blob_access == "reader"
assert reader.inputs[1].blob_access is None
assert reader.output.arrow_type == "binary"
eager = read_bytes.registration_request.signature.inputs[0]
assert eager.arrow_type == "binary"
assert eager.blob_access == "bytes"
encoded = json.loads(read_blob.registration_request.to_canonical_json())
assert encoded["signature"]["inputs"][0]["blob_access"] == "reader"
def test_nullable_blob_reader_and_explicit_blob_bytes():
@udf
def nullable(video: Optional[BlobReader]) -> int:
return 0 if video is None else video.size()
@udf(
input_schema=pa.schema([lancedb.blob("video", nullable=False)]),
output_schema=pa.field("size", pa.int64(), nullable=False),
)
def eager(video: bytes) -> int:
return len(video)
parameter = nullable.registration_request.signature.inputs[0]
assert parameter.nullable is True
assert parameter.blob_access == "reader"
parameter = eager.registration_request.signature.inputs[0]
assert parameter.arrow_type == "blob_v2"
assert parameter.blob_access == "bytes"
def test_blob_reader_rejects_an_incompatible_explicit_schema():
def read(video: BlobReader) -> int:
return video.size()
with pytest.raises(TypeError, match="requires a Blob V2 field"):
udf(
read,
input_schema=pa.schema([pa.field("video", pa.int64(), nullable=False)]),
output_schema=pa.field("size", pa.int64(), nullable=False),
)
def test_whole_named_struct_function_can_include_a_blob_result_field():
@udf(
input_schema=pa.schema([lancedb.blob("image", nullable=False)]),
@@ -1674,7 +1729,11 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip():
)
registration = db.create_function_async(normalize_score)
assert registration.id == "job-register"
assert registration.job_type is None
assert registration.status() == "finished"
assert registration.job_type == "create_function"
created = registration.wait()
assert registration.job_type == "create_function"
reopened = db.get_function("normalize_score", version=created.version)
assert created == reopened
+14
View File
@@ -391,6 +391,20 @@ def test_custom_stop_words_list(table):
)
def test_list_indices_with_missing_fts_files(table):
# Regression test for https://github.com/lancedb/lancedb/issues/4250
table.create_fts_index("text")
table.create_scalar_index("id")
fts = next(i for i in table.list_indices() if i.index_type == "FTS")
shutil.rmtree(Path(table.uri) / "_indices" / fts.index_uuid)
indices = {i.name: i for i in table.list_indices()}
assert set(indices) == {"text_idx", "id_idx"}
assert indices["text_idx"].index_type == "FTS"
assert indices["text_idx"].columns == ["text"]
def test_search_fts(table):
table.create_fts_index("text")
results = table.search("puppy").select(["id", "text"]).limit(5).to_list()
+12 -2
View File
@@ -421,13 +421,23 @@ async def test_create_4bit_ivfpq_index(some_table: AsyncTable):
@pytest.mark.asyncio
async def test_create_ivfrq_index(some_table: AsyncTable):
await some_table.create_index("vector", config=IvfRq(num_bits=1))
@pytest.mark.parametrize(
"config, expected_bits",
[
pytest.param(IvfRq(), 5, id="default"),
pytest.param(IvfRq(num_bits=None), 5, id="none"),
pytest.param(IvfRq(num_bits=1), 1, id="one-bit"),
pytest.param(IvfRq(num_bits=4), 4, id="multi-bit"),
],
)
async def test_create_ivfrq_index(some_table: AsyncTable, config, expected_bits):
await some_table.create_index("vector", config=config)
indices = await some_table.list_indices()
assert len(indices) == 1
assert indices[0].index_type == "IvfRq"
assert indices[0].columns == ["vector"]
assert indices[0].name == "vector_idx"
assert indices[0].index_details["compression"]["num_bits"] == expected_bits
@pytest.mark.asyncio
+137
View File
@@ -88,6 +88,143 @@ async def advance(items, asynchronous):
KINDS = ["secrets", "views", "jobs", "functions", "materialized_views"]
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_compatible_name_listing_pages(listing_server, kind, asynchronous):
endpoint, requests, responses = listing_server
responses.extend(
[
(200, response_page(kind, ["a", "b"], "next")),
(200, response_page(kind, ["c"], None)),
]
)
db = await connection(endpoint, asynchronous)
listing = getattr(db, "list_" + kind)(
namespace_path=["team"], page_token="start", limit=2
)
field = kind
assert len(requests) == (0 if asynchronous else 1)
if asynchronous:
with pytest.raises(RuntimeError, match="Await or iterate"):
getattr(listing, field)
assert listing.page_token == "start"
assert await listing is listing
assert len(requests) == 1
page_names = getattr(listing, field)
assert type(page_names) is list
assert page_names == ["a", "b"]
page_names.append("changed")
assert getattr(listing, field) == ["a", "b"]
assert listing.page_token == "next"
assert listing.num_page_results() == 2
assert await advance(listing, asynchronous) == "a"
assert getattr(listing, field) == ["a", "b"]
assert await advance(listing, asynchronous) == "b"
assert len(requests) == 1
assert await advance(listing, asynchronous) == "c"
assert getattr(listing, field) == ["c"]
assert listing.page_token is None
assert len(requests) == 2
assert ([item async for item in listing] if asynchronous else list(listing)) == []
if asynchronous:
assert await listing is listing
assert len(requests) == 2
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
async def test_name_listing_skips_empty_page_only_during_iteration(
listing_server, kind
):
endpoint, requests, responses = listing_server
responses.extend(
[
(200, response_page(kind, [], "next")),
(200, response_page(kind, ["a"], None)),
]
)
db = await connection(endpoint, True)
listing = getattr(db, "list_" + kind)()
await listing
assert getattr(listing, kind) == []
assert listing.page_token == "next"
assert len(requests) == 1
assert [name async for name in listing] == ["a"]
assert len(requests) == 2
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_name_listing_zero_limit_uses_no_request(
listing_server, kind, asynchronous
):
endpoint, requests, _ = listing_server
db = await connection(endpoint, asynchronous)
listing = getattr(db, "list_" + kind)(limit=0)
if asynchronous:
assert await listing is listing
assert [name async for name in listing] == []
else:
assert list(listing) == []
assert getattr(listing, kind) == []
assert listing.page_token is None
assert requests == []
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
async def test_name_listing_failure_retains_token(listing_server, kind):
endpoint, requests, responses = listing_server
responses.append((401, {}))
db = await connection(endpoint, True)
listing = getattr(db, "list_" + kind)(page_token="retry")
with pytest.raises(HttpError):
await listing
assert listing.page_token == "retry"
with pytest.raises(HttpError):
await listing
assert [name async for name in listing] == []
assert len(requests) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
async def test_name_listing_direct_async_iteration(listing_server, kind):
endpoint, requests, responses = listing_server
responses.extend(
[
(200, response_page(kind, ["a"], "next")),
(200, response_page(kind, ["b"], None)),
]
)
db = await connection(endpoint, True)
listing = getattr(db, "list_" + kind)(limit=1)
assert requests == []
assert [name async for name in listing] == ["a", "b"]
assert getattr(listing, kind) == ["b"]
assert len(requests) == 2
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
async def test_name_listing_serializes_concurrent_advances(listing_server, kind):
endpoint, requests, responses = listing_server
responses.append((200, response_page(kind, ["a", "b"], None)))
db = await connection(endpoint, True)
listing = getattr(db, "list_" + kind)()
results = await asyncio.gather(
listing.__anext__(),
listing.__anext__(),
listing.__anext__(),
return_exceptions=True,
)
assert results[:2] == ["a", "b"]
assert isinstance(results[2], StopAsyncIteration)
assert len(requests) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", KINDS)
@pytest.mark.parametrize("asynchronous", [False, True])
+63 -24
View File
@@ -8,6 +8,7 @@ import threading
import lancedb
import pytest
from lancedb.db import AsyncConnection
from lancedb.materialized_view import MaterializedViewDefinition
from lancedb.remote.db import RemoteDBConnection
@@ -56,7 +57,7 @@ def mock_remote_materialized_views():
@contextlib.contextmanager
def mock_remote_materialized_view_create():
def mock_remote_materialized_view_drop():
requests = []
class Handler(http.server.BaseHTTPRequestHandler):
@@ -67,7 +68,8 @@ def mock_remote_materialized_view_create():
length = int(self.headers.get("Content-Length", 0))
body = json.loads(self.rfile.read(length) or b"{}")
requests.append((self.path, body))
job_id = "mv-drop-123" if self.path.endswith("/drop") else "mv-create-123"
assert self.path.endswith("/drop")
job_id = "mv-drop-123"
encoded = json.dumps({"job_id": job_id}).encode()
self.send_response(202)
self.send_header("Content-Type", "application/json")
@@ -98,7 +100,7 @@ def test_remote_list_uses_namespace_route():
def test_remote_drop_async_returns_server_job():
with mock_remote_materialized_view_create() as (host, requests):
with mock_remote_materialized_view_drop() as (host, requests):
db = lancedb.connect(
"db://dev",
api_key="fake",
@@ -192,11 +194,13 @@ def test_create_and_refresh_jobs(tmp_path):
"adults", "people", where="age >= 18", with_no_data=True
)
assert create_job.id is None
assert create_job.job_type is None
assert create_job.wait() is None
view = db.open_materialized_view("adults")
refresh_job = view.refresh_async()
assert refresh_job.id is None
assert refresh_job.job_type is None
result = refresh_job.wait()
assert result.mode == "rebuild"
assert result.rows_written == 2
@@ -337,11 +341,13 @@ async def test_async_create_and_refresh_jobs(tmp_path):
"adults", "people", with_no_data=True
)
assert create_job.id is None
assert create_job.job_type is None
assert await create_job.wait() is None
view = await db.open_materialized_view("adults")
refresh_job = await view.refresh_async()
assert refresh_job.id is None
assert refresh_job.job_type is None
result = await refresh_job.wait()
assert result.mode == "rebuild"
assert result.rows_written == 1
@@ -567,6 +573,27 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path):
assert list(resumed) == []
class QueryBackedCreateJob:
id = "0198f1b2-c3d4-7e5f-8123-456789abcdef"
job_type = "query"
async def wait(self):
return None
async def status(self):
return "finished"
class CapturingCreateConnection:
def __init__(self):
self.calls = []
async def create_materialized_view_async(self, name, source, **options):
self.calls.append((name, source, options))
return QueryBackedCreateJob()
@pytest.mark.parametrize(
"options",
[
@@ -576,16 +603,28 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path):
],
)
def test_remote_create_forwards_refresh_options_without_interval_default(options):
with mock_remote_materialized_view_create() as (host, requests):
db = lancedb.connect("db://dev", api_key="fake", host_override=host)
job = db.create_materialized_view_async("adults", "people", **options)
assert job.id == "mv-create-123"
assert requests == [
native = CapturingCreateConnection()
db = RemoteDBConnection.__new__(RemoteDBConnection)
db._conn = AsyncConnection(native)
job = db.create_materialized_view_async("adults", "people", **options)
assert job.id == QueryBackedCreateJob.id
assert job.job_type == "query"
assert not hasattr(job, "query_id")
assert job.wait() is None
assert job.status() == "finished"
assert native.calls == [
(
"/v1/materialized_view/adults/create",
"adults",
"people",
{
"query": 'SELECT * FROM "people"',
**options,
"projections": None,
"filter": None,
"limit": None,
"with_no_data": False,
"auto_refresh_enabled": options.get("auto_refresh_enabled"),
"auto_refresh_interval_seconds": options.get(
"auto_refresh_interval_seconds"
),
},
)
]
@@ -593,18 +632,18 @@ def test_remote_create_forwards_refresh_options_without_interval_default(options
@pytest.mark.asyncio
async def test_async_remote_create_keeps_interval_unspecified():
with mock_remote_materialized_view_create() as (host, requests):
db = await lancedb.connect_async("db://dev", api_key="fake", host_override=host)
job = await db.create_materialized_view_async(
"adults",
"people",
auto_refresh_enabled=True,
)
assert job.id == "mv-create-123"
assert requests[0][1] == {
"query": 'SELECT * FROM "people"',
"auto_refresh_enabled": True,
}
native = CapturingCreateConnection()
db = AsyncConnection(native)
job = await db.create_materialized_view_async(
"adults", "people", auto_refresh_enabled=True
)
assert job.id == QueryBackedCreateJob.id
assert job.job_type == "query"
assert not hasattr(job, "query_id")
assert await job.wait() is None
assert await job.status() == "finished"
assert native.calls[0][2]["auto_refresh_enabled"] is True
assert native.calls[0][2]["auto_refresh_interval_seconds"] is None
@pytest.mark.parametrize(
@@ -625,7 +664,7 @@ def test_embedded_create_refuses_automatic_refresh_before_creating_view(
@pytest.mark.parametrize("seconds", [0, -1, float("nan"), float("inf")])
def test_remote_create_rejects_invalid_interval_without_request(seconds):
with mock_remote_materialized_view_create() as (host, requests):
with mock_remote_materialized_view_drop() as (host, requests):
db = lancedb.connect("db://dev", api_key="fake", host_override=host)
with pytest.raises(ValueError, match="positive, finite"):
db.create_materialized_view_async(
@@ -3,7 +3,6 @@
import json
import httpx
import numpy as np
import pyarrow as pa
import pytest
@@ -14,6 +13,7 @@ from lancedb.embeddings import get_registry
@pytest.fixture
def ollama_embeddings():
httpx = pytest.importorskip("httpx")
pytest.importorskip("ollama")
requests = []
@@ -0,0 +1,38 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import pytest
from lancedb.rerankers import OpenaiReranker
@pytest.mark.parametrize(
"explicit_key, environment_key, expected_key",
[
("explicit-test-key", "environment-test-key", "explicit-test-key"),
("explicit-test-key", None, "explicit-test-key"),
(None, "environment-test-key", "environment-test-key"),
],
)
def test_openai_reranker_api_key_precedence(
monkeypatch, explicit_key, environment_key, expected_key
):
pytest.importorskip("openai")
if environment_key is None:
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
else:
monkeypatch.setenv("OPENAI_API_KEY", environment_key)
reranker = OpenaiReranker(api_key=explicit_key)
client = reranker._client
try:
assert client.api_key == expected_key
finally:
client.close()
def test_openai_reranker_requires_api_key(monkeypatch):
pytest.importorskip("openai")
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
with pytest.raises(ValueError, match="OPENAI_API_KEY not set"):
OpenaiReranker()._client
+7
View File
@@ -811,6 +811,13 @@ def test_rename_column(some_permutation: Permutation):
some_permutation.rename_column("non_existent", "new_value")
def test_rename_column_keeps_position(some_permutation: Permutation):
renamed = some_permutation.rename_column("id", "new_id")
assert renamed.schema == pa.schema([("new_id", pa.int64()), ("value", pa.int64())])
batch = renamed.with_format("arrow").__getitems__([0])
assert batch.schema.names == ["new_id", "value"]
def test_rename_columns(some_permutation: Permutation):
assert some_permutation.rename_columns({"value": "new_value"}).schema == pa.schema(
[("id", pa.int64()), ("new_value", pa.int64())]
+27
View File
@@ -2324,6 +2324,33 @@ def test_take_queries(tmp_path):
]
def test_take_queries_with_row_id(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("test", pa.table({"idx": range(100)}))
by_row_id = table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow()
assert "_rowid" in by_row_id.column_names
by_offset = table.take_offsets([5, 2, 17], with_row_id=True).to_arrow()
assert "_rowid" in by_offset.column_names
# The column is opt-in and must stay out of the default projection.
assert "_rowid" not in table.take_row_ids([5, 2, 17]).to_arrow().column_names
@pytest.mark.asyncio
async def test_take_queries_with_row_id_async(tmp_path):
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("test", pa.table({"idx": range(100)}))
by_row_id = await table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow()
assert "_rowid" in by_row_id.column_names
by_offset = await table.take_offsets([5, 2, 17], with_row_id=True).to_arrow()
assert "_rowid" in by_offset.column_names
default = await table.take_row_ids([5, 2, 17]).to_arrow()
assert "_rowid" not in default.column_names
def test_take_queries_to_batches(tmp_path):
# Regression test for the sync take-query path: `to_batches` previously
# raised ``AttributeError: 'AsyncTakeQuery' object has no attribute
+217 -19
View File
@@ -5,6 +5,7 @@ import contextlib
from datetime import timedelta
import http.server
import io
import inspect
import json
import multiprocessing as mp
import pickle
@@ -765,6 +766,15 @@ def test_create_table_embedding_functions(with_schema):
]
def test_create_table_positional_parameters_match_base_contract():
from lancedb.db import DBConnection
from lancedb.remote.db import RemoteDBConnection
expected = list(inspect.signature(DBConnection.create_table).parameters)[:9]
actual = list(inspect.signature(RemoteDBConnection.create_table).parameters)[:9]
assert actual == expected
def test_create_table_exist_ok():
def handler(request):
if request.path == "/v1/table/test/create/?mode=exist_ok":
@@ -857,10 +867,116 @@ def test_table_unimplemented_functions():
with mock_lancedb_connection(handler) as db:
table = db.create_table("test", [{"id": 1}])
with pytest.raises(NotImplementedError):
table.to_arrow()
with pytest.raises(NotImplementedError):
table.to_pandas()
for method in ["to_arrow", "to_pandas", "to_polars"]:
with pytest.raises(NotImplementedError) as exc_info:
getattr(table, method)()
message = str(exc_info.value)
assert method in message
assert "not" in message and "supported" in message
assert "remote tables" in message
assert "cloud" not in message.lower()
assert f"search().{method}()" in message
@pytest.mark.parametrize("method", ["to_arrow", "to_pandas"])
@pytest.mark.parametrize("blob_mode", ["lazy", "descriptions", "bytes", "invalid"])
def test_remote_conversion_blob_mode_errors(method, blob_mode):
with blob_remote_table() as table:
if blob_mode == "invalid":
with pytest.raises(ValueError, match="blob_mode must be one of"):
getattr(table, method)(blob_mode=blob_mode)
else:
with pytest.raises(NotImplementedError) as error:
getattr(table, method)(blob_mode=blob_mode)
assert f"search().{method}()" in str(error.value)
@pytest.mark.asyncio
@pytest.mark.parametrize("version", [None, 1])
@pytest.mark.parametrize("expose_location", [False, True])
async def test_remote_to_lance_rejection(
tmp_path, version, expose_location, monkeypatch
):
import lance
location = tmp_path / "test.lance"
lance.write_dataset(pa.table({"id": [1]}), location)
lance.write_dataset(pa.table({"id": [2]}), location, mode="append")
requests = []
def fail_dataset(*args, **kwargs):
raise AssertionError("remote storage must not be opened")
monkeypatch.setattr(lance, "dataset", fail_dataset)
def handler(request):
requests.append(request.path)
assert request.path == "/v1/table/test/describe/"
content_len = int(request.headers.get("Content-Length", 0))
body = json.loads(request.rfile.read(content_len)) if content_len else {}
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
response = {
"version": body.get("version") or 2,
"schema": {"fields": []},
}
if expose_location:
response["location"] = str(location)
request.wfile.write(json.dumps(response).encode())
options = {"default_scan_options": {"with_row_id": True}}
with mock_lancedb_connection(handler) as db:
table = db.open_table("test", version=version)
opened_requests = list(requests)
with pytest.raises(NotImplementedError) as sync_error:
table.to_lance(**options)
assert requests == opened_requests
async with mock_lancedb_connection_async(handler) as db:
table = await db.open_table("test", version=version)
opened_requests = list(requests)
with pytest.raises(NotImplementedError) as async_error:
await table.to_lance(**options)
assert requests == opened_requests
assert "to_lance() is not supported for remote tables" in str(sync_error.value)
assert "query the server" in str(sync_error.value)
assert str(sync_error.value) == str(async_error.value)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"method, kwargs",
[
("drop_all_tables", {}),
("drop_all_tables", {"namespace_path": []}),
("drop_all_tables", {"namespace_path": ["team"]}),
("drop_database", {}),
],
)
async def test_remote_drop_all_tables_not_supported(method, kwargs):
requests = []
def handler(request):
requests.append(request.path)
request.send_response(500)
request.end_headers()
with mock_lancedb_connection(handler) as db:
with pytest.raises(NotImplementedError) as sync_error:
getattr(db, method)(**kwargs)
async with mock_lancedb_connection_async(handler) as db:
with pytest.raises(NotImplementedError) as async_error:
await getattr(db, method)(**kwargs)
assert "Dropping all tables is not currently supported in the remote API" in str(
sync_error.value
)
assert "drop_table" in str(sync_error.value)
assert str(sync_error.value) == str(async_error.value)
assert requests == []
def test_table_to_pandas_not_supported():
@@ -1039,8 +1155,9 @@ def test_table_create_indices():
# Test create_fts_index with custom name (legacy method)
with pytest.warns(DeprecationWarning, match="create_fts_index"):
table.create_fts_index(
"text",
field_names="text",
wait_timeout=timedelta(seconds=2),
tokenizer_name="en_stem",
block_size=256,
custom_stop_words=["cloud"],
name="custom_fts_idx",
@@ -1049,6 +1166,8 @@ def test_table_create_indices():
# Test create_index with custom name (legacy form: vector_column_name kwarg)
with pytest.warns(DeprecationWarning, match="create_index"):
table.create_index(
"cosine",
2,
vector_column_name="vector",
wait_timeout=timedelta(seconds=10),
name="custom_vector_idx",
@@ -1061,7 +1180,8 @@ def test_table_create_indices():
scalar_req = received_requests[0]
assert "name" in scalar_req
assert scalar_req["name"] == "custom_scalar_idx"
assert scalar_req["replace"] is False
# The wire protocol omits replace=True (the server default).
assert "replace" not in scalar_req
# Check FTS index request has custom name
fts_req = received_requests[1]
@@ -1070,18 +1190,39 @@ def test_table_create_indices():
assert fts_req["replace"] is False
assert fts_req["block_size"] == 256
assert fts_req["custom_stop_words"] == ["cloud"]
assert fts_req["language"] == "English"
assert fts_req["stem"] is True
# Check vector index request has custom name
vector_req = received_requests[2]
assert "name" in vector_req
assert vector_req["name"] == "custom_vector_idx"
assert "replace" not in vector_req
assert vector_req["metric_type"] == "cosine"
assert vector_req["num_partitions"] == 2
with pytest.warns(DeprecationWarning, match="create_index"):
table.create_index(
index_type="IVF_HNSW_SQ",
num_partitions=1,
m=10,
ef_construction=50,
max_iterations=12,
sample_rate=64,
)
tuned_req = received_requests[-1]
assert tuned_req["index_type"] == "IVF_HNSW_SQ"
assert tuned_req["num_partitions"] == 1
assert tuned_req["m"] == 10
assert tuned_req["ef_construction"] == 50
assert tuned_req["max_iterations"] == 12
assert tuned_req["sample_rate"] == 64
table.wait_for_index(["custom_scalar_idx"], timedelta(seconds=2))
table.wait_for_index(
["custom_fts_idx", "custom_vector_idx"], timedelta(seconds=2)
)
table.drop_index("custom_vector_idx")
table.drop_index(name="custom_vector_idx")
table.drop_index("custom_scalar_idx")
table.drop_index("custom_fts_idx")
@@ -1477,7 +1618,10 @@ def test_remote_json_fts_config(server_version):
assert request["block_size"] == 256
def test_remote_create_index_new_api(caplog):
@pytest.mark.parametrize(
"rq_kwargs", [{}, {"num_bits": None}, {"num_bits": 1}, {"num_bits": 4}]
)
def test_remote_create_index_new_api(caplog, rq_kwargs):
received_requests = []
def handler(request):
@@ -1550,7 +1694,7 @@ def test_remote_create_index_new_api(caplog):
table.create_index("category", config=BTree())
table.create_index("text", config=FTS(block_size=256))
# IvfRq via new API
table.create_index("vector", config=IvfRq(distance_type="l2"))
table.create_index("vector", config=IvfRq(distance_type="l2", **rq_kwargs))
table.create_index(
"vector", config=IvfPq(distance_type="l2"), replace=False
)
@@ -1575,6 +1719,10 @@ def test_remote_create_index_new_api(caplog):
"vector",
]
assert received_requests[2]["block_size"] == 256
if rq_kwargs.get("num_bits") is None:
assert "num_bits" not in received_requests[3]
else:
assert received_requests[3]["num_bits"] == rq_kwargs["num_bits"]
assert received_requests[4]["replace"] is False
assert received_requests[5]["replace"] is False
assert "replace is not supported" not in caplog.text
@@ -1675,16 +1823,19 @@ def test_remote_unsupported_legacy_hnsw_pq():
)
def test_remote_python_only_messages(caplog):
def test_remote_python_only_messages():
with mock_lancedb_connection(remote_unsupported_handler) as db:
with caplog.at_level("INFO"):
table = db.open_table("test", storage_options={}, index_cache_size=1)
assert "storage_options is ignored for remote tables" in caplog.text
assert "index_cache_size is ignored for remote tables" in caplog.text
for method in ["to_arrow", "to_pandas"]:
for options in [{"storage_options": {}}, {"index_cache_size": 1}]:
with pytest.raises(NotImplementedError, match=next(iter(options))):
db.open_table("test", **options)
table = db.open_table("test")
for method in ["to_arrow", "to_pandas", "to_polars"]:
with pytest.raises(NotImplementedError) as error:
getattr(table, method)()
assert str(error.value) == f"{method}() is not supported for remote tables."
assert str(error.value) == (
f"{method}() is not supported for remote tables. "
f"Use table.search().{method}() instead."
)
for method in ["compact_files", "cleanup_old_versions"]:
with pytest.warns(UserWarning) as warning:
getattr(table, method)()
@@ -2225,6 +2376,23 @@ def test_head():
assert data == pa.table({"id": [1, 2, 3]})
def test_remote_take_row_ids_with_row_id():
bodies = []
def handler(body):
bodies.append(body)
return pa.table({"id": [1, 2], "_rowid": pa.array([0, 1], pa.uint64())})
with query_test_table(handler) as table:
result = table.take_row_ids([0, 1], with_row_id=True).to_arrow()
assert "_rowid" in result.column_names
assert bodies[-1]["with_row_id"] is True
with query_test_table(handler) as table:
table.take_row_ids([0, 1]).to_arrow()
assert "with_row_id" not in bodies[-1]
def test_query_sync_minimal():
def handler(body):
assert body == {
@@ -3473,14 +3641,17 @@ def blob_descriptor_take_table():
@contextlib.contextmanager
def blob_remote_table(*, server_version=Version("0.5.0")):
def blob_remote_table(*, server_version=Version("0.5.0"), location=None):
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(json.dumps(BLOB_DESCRIBE_RESPONSE).encode())
response = dict(BLOB_DESCRIBE_RESPONSE)
if location is not None:
response["location"] = str(location)
request.wfile.write(json.dumps(response).encode())
elif request.path.startswith("/v1/table/test/blob/image/"):
path = request.path.partition("?")[0]
row_id = int(path.split("/")[-2])
@@ -3515,7 +3686,6 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
if columns == ["id"]:
response_table = blob_query_response_table().select(["id"])
else:
assert body["with_row_id"] is True
if columns == ["image"]:
# Blob handles are sized from a descriptor take.
assert body["filter"].startswith("_rowid IN")
@@ -3528,6 +3698,8 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
)
else:
response_table = blob_query_response_table()
if not body.get("with_row_id", False):
response_table = response_table.drop(["_rowid"])
request.send_response(200)
request.send_header("Content-Type", "application/vnd.apache.arrow.file")
request.send_header("x-lancedb-version", "1")
@@ -3562,6 +3734,32 @@ def test_remote_blob_columns_and_fetch():
assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
@pytest.mark.parametrize("blob_mode", ["descriptions", "bytes", "lazy"])
def test_sync_remote_blob_pandas_avoids_storage_access(
tmp_path, monkeypatch, blob_mode
):
import lance
location = tmp_path / "remote.lance"
lance.write_dataset(pa.table({"id": [7], "image": [b"server-owned"]}), location)
def fail_dataset(*args, **kwargs):
raise AssertionError("remote storage must not be opened")
monkeypatch.setattr(lance, "dataset", fail_dataset)
with blob_remote_table(location=location) as table:
monkeypatch.setattr(table, "to_lance", fail_dataset)
query = table.search().select(["id", "image"])
if blob_mode == "descriptions":
df = query.to_pandas(blob_mode=blob_mode)
assert df["id"].tolist() == [1, 2, 3]
assert "_rowid" not in df.columns
assert df["image"][0]["size"] == 5
else:
with pytest.raises(NotImplementedError, match="stable table snapshot"):
query.to_pandas(blob_mode=blob_mode)
@pytest.mark.asyncio
@pytest.mark.parametrize("use_async", [False, True])
async def test_remote_fetch_blob_ranges(use_async):
+154
View File
@@ -0,0 +1,154 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import inspect
from urllib.parse import parse_qs, urlparse
import lancedb
import pyarrow as pa
import pytest
from test_remote_db import mock_lancedb_connection, mock_lancedb_connection_async
@pytest.mark.parametrize("uri", ["memory://", "db://dev"])
def test_unknown_connect_options(uri):
with pytest.raises(ValueError, match="Unknown keyword arguments.*bogus"):
lancedb.connect(uri, api_key="fake", bogus=1)
@pytest.mark.parametrize("option", ["session", "manifest_enabled"])
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_remote_connect_rejects_local_options(option, asynchronous):
value = lancedb.Session() if option == "session" else True
with pytest.raises(NotImplementedError, match=option):
if asynchronous:
await lancedb.connect_async("db://dev", api_key="fake", **{option: value})
else:
lancedb.connect("db://dev", api_key="fake", **{option: value})
def test_remote_create_signature():
from lancedb.db import DBConnection
from lancedb.remote.db import RemoteDBConnection
expected = inspect.signature(DBConnection.create_table).parameters
actual = inspect.signature(RemoteDBConnection.create_table).parameters
assert list(actual) == list(expected)
for name in expected:
assert actual[name].kind == expected[name].kind
if name == "mode":
# Remote connections retain their existing server-default mode.
assert actual[name].default is None
else:
assert actual[name].default == expected[name].default
def test_remote_create_positional_mode():
modes = []
def handler(request):
assert urlparse(request.path).path == "/v1/table/ct_pos/create/"
modes.append(parse_qs(urlparse(request.path).query)["mode"])
request.send_response(200)
request.end_headers()
with mock_lancedb_connection(handler) as db:
data = [{"value": 1}]
db.create_table("ct_pos", data, None, "overwrite")
db.create_table("ct_pos", data, None, "overwrite")
db.create_table("ct_pos", data, None, "create", True)
with pytest.raises(ValueError, match="Invalid mode"):
db.create_table("ct_pos", data, None, "drop")
assert modes == [["overwrite"], ["overwrite"], ["exist_ok"]]
CREATE_OPTIONS = [
{"new_table_data_storage_version": "2.0"},
{"new_table_enable_v2_manifest_paths": "false"},
{"bogus_option_xyz": "1"},
]
@pytest.mark.parametrize("storage_options", CREATE_OPTIONS)
@pytest.mark.parametrize("empty", [False, True])
@pytest.mark.asyncio
async def test_remote_create_rejects_write_options(storage_options, empty):
requests = []
def handler(request):
requests.append(request.path)
request.send_response(200)
request.end_headers()
kwargs = {"schema": pa.schema([("value", pa.int64())])}
if not empty:
kwargs["data"] = [{"value": 1}]
async with mock_lancedb_connection_async(handler) as db:
with pytest.raises(NotImplementedError, match="write_options.*remote"):
await db.create_table("t", storage_options=storage_options, **kwargs)
assert requests == []
@pytest.mark.parametrize(
"options",
[
{"data_storage_version": "2.0"},
{"enable_v2_manifest_paths": False},
{"storage_options": {"x": "1"}},
],
)
def test_sync_remote_create_rejects_write_options(options):
requests = []
def handler(request):
requests.append(request.path)
request.send_response(200)
request.end_headers()
with mock_lancedb_connection(handler) as db:
with pytest.raises(NotImplementedError, match="remote"):
db.create_table("t", [{"value": 1}], **options)
assert requests == []
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize(
"options", [{"index_cache_size": 10}, {"storage_options": {"x": "y"}}]
)
@pytest.mark.asyncio
async def test_remote_open_rejects_local_options(options, asynchronous):
requests = []
def handler(request):
requests.append(request.path)
request.send_response(200)
request.end_headers()
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
if asynchronous:
async with mock_lancedb_connection_async(handler) as db:
with pytest.raises(NotImplementedError, match="remote"):
await db.open_table("t", **options)
else:
with mock_lancedb_connection(handler) as db:
with pytest.raises(NotImplementedError, match="remote"):
db.open_table("t", **options)
assert requests == []
@pytest.mark.asyncio
async def test_remote_open_rejects_location():
requests = []
def handler(request):
requests.append(request.path)
request.send_response(200)
request.end_headers()
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
async with mock_lancedb_connection_async(handler) as db:
with pytest.raises(NotImplementedError, match="location.*remote"):
await db.open_table("t", location="/tmp/zzz")
assert requests == []
+19
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import os
from unittest.mock import MagicMock
import random
import lancedb
@@ -547,6 +548,24 @@ def test_openai_reranker(tmp_path):
_run_test_reranker(reranker, table, "single player experience", None, schema)
def test_openai_reranker_does_not_send_temperature():
reranker = OpenaiReranker(api_key="test-key")
create = MagicMock()
create.return_value.choices = [
MagicMock(
message=MagicMock(
content='{"documents": [{"content": "a", "relevance_score": 0.5}]}'
)
)
]
reranker.__dict__["_client"] = MagicMock()
reranker.__dict__["_client"].chat.completions.create = create
reranker._rerank(pa.table({"text": ["a"]}), "query")
assert "temperature" not in create.call_args.kwargs
@pytest.mark.skipif(
os.environ.get("JINA_API_KEY") is None, reason="JINA_API_KEY not set"
)
+17 -3
View File
@@ -137,9 +137,23 @@ async def test_async_namespace_connection_rejects_sql(tmp_path):
await connection.describe_query(NIL_QUERY_ID)
def test_describe_query_requires_uuid():
with pytest.raises(TypeError, match="UUID"):
remote_connection().describe_query(str(NIL_QUERY_ID))
@pytest.mark.asyncio
@pytest.mark.parametrize("identity", [NIL_QUERY_ID, str(NIL_QUERY_ID)])
async def test_describe_query_accepts_uuid_and_job_id_string(identity):
calls = []
class NativeConnection:
async def describe_query(self, query_id):
calls.append(query_id)
connection = AsyncConnection(NativeConnection())
await connection.describe_query(identity)
assert calls == [NIL_QUERY_ID]
def test_describe_query_rejects_registry_job_id():
with pytest.raises(ValueError, match="UUID"):
remote_connection().describe_query("registry-job-id")
@pytest.mark.parametrize(
+170
View File
@@ -0,0 +1,170 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from datetime import timedelta
import json
from urllib.parse import parse_qs, urlparse
import lancedb
from lancedb.table import Table
import pytest
from test_remote_db import mock_lancedb_connection
@pytest.fixture
def legacy_table():
"""Implement the pre-existing abstract API without the two new members."""
def unsupported(self, *args, **kwargs):
raise NotImplementedError
def update_field_metadata(self, updates):
self.metadata_updates.append(updates)
return object() # The deprecated wrapper must not expose this result.
# Supply concrete implementations through normal ABC class creation. Leave
# namespace and replace_field_metadata inherited, as an old subclass would.
implementations = {
name: unsupported
for name in Table.__abstractmethods__
if name not in {"namespace", "replace_field_metadata"}
}
implementations.update(
name=property(lambda self: "legacy"),
update_field_metadata=update_field_metadata,
)
legacy_class = type("LegacyTable", (Table,), implementations)
table = legacy_class()
table.metadata_updates = []
return table
def test_legacy_table_namespace_and_id(legacy_table):
assert legacy_table.namespace == []
assert legacy_table.id == "legacy"
namespace = legacy_table.namespace
namespace.append("changed")
assert legacy_table.namespace == []
assert legacy_table.id == "legacy"
def test_legacy_table_replace_field_metadata(legacy_table):
with pytest.warns(DeprecationWarning, match="update_field_metadata"):
result = legacy_table.replace_field_metadata("nested.field", {"a": "b"})
assert result is None
assert legacy_table.metadata_updates == [
{"path": "nested.field", "metadata": {"a": "b"}, "replace": True}
]
def respond(request, body):
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(json.dumps(body).encode())
def test_remote_uri():
with mock_lancedb_connection(lambda request: respond(request, {})) as db:
assert db.uri == "db://dev"
@pytest.mark.parametrize("attribute", ["session", "read_consistency_interval"])
def test_remote_unsupported_properties(attribute):
with mock_lancedb_connection(lambda request: respond(request, {})) as db:
with pytest.raises(NotImplementedError, match=attribute):
getattr(db, attribute)
@pytest.mark.parametrize("tables", [[], ["first", "last"]])
def test_remote_len_and_contains_follow_pages(tables):
requests = []
def handler(request):
requests.append(request.path)
query = parse_qs(urlparse(request.path).query)
if tables and "page_token" not in query:
respond(request, {"tables": [tables[0]], "page_token": "opaque-next"})
else:
if tables:
assert query["page_token"] == ["opaque-next"]
respond(request, {"tables": tables[1:]})
with mock_lancedb_connection(handler) as db:
assert len(db) == len(tables)
assert ("last" in db) == bool(tables)
assert "missing" not in db
requests.clear()
assert ("first" in db) == bool(tables)
assert len(requests) == 1
@pytest.mark.parametrize("namespace", [[], ["parent", "child"]])
def test_remote_table_identifiers(namespace):
def handler(request):
respond(request, {"version": 1, "schema": {"fields": []}})
with mock_lancedb_connection(handler) as db:
table = db.open_table("rfm", namespace_path=namespace)
assert table.namespace == namespace
assert table.id == "$".join(namespace + ["rfm"])
def test_remote_replace_field_metadata():
updates = []
def handler(request):
if request.path.endswith("/describe/"):
respond(request, {"version": 1, "schema": {"fields": []}})
else:
assert request.path == "/v1/table/rfm/update_field_metadata/"
body = request.rfile.read(int(request.headers["Content-Length"]))
updates.append(json.loads(body))
respond(request, {"version": 2})
with mock_lancedb_connection(handler) as db:
table = db.open_table("rfm")
with pytest.warns(DeprecationWarning, match="update_field_metadata"):
assert table.replace_field_metadata("id", {"a": "b"}) is None
assert updates == [
{"updates": [{"path": "id", "metadata": {"a": "b"}, "replace": True}]}
]
@pytest.mark.parametrize("custom_session", [False, True])
def test_local_properties_and_metadata(tmp_path, custom_session):
supplied = lancedb.Session() if custom_session else None
db = lancedb.connect(
tmp_path, session=supplied, read_consistency_interval=timedelta(seconds=5)
)
session = db.session
assert session is supplied
before = session.size_bytes if session is not None else 0
table = db.create_table("rfm", [{"id": i} for i in range(3)])
assert table.to_arrow().num_rows == 3
if supplied is not None:
assert supplied.size_bytes > before
assert db.read_consistency_interval == timedelta(seconds=5)
assert table.namespace == []
assert table.id == "rfm"
with pytest.warns(DeprecationWarning):
table.replace_field_metadata("id", {"old": "value"})
with pytest.warns(DeprecationWarning):
table.replace_field_metadata("id", {"a": "b"})
assert table.schema.field("id").metadata == {b"a": b"b"}
def test_namespace_connection_properties_remain_writable(tmp_path):
session = lancedb.Session()
interval = timedelta(seconds=5)
with lancedb.connect_namespace(
"dir",
{"root": str(tmp_path)},
session=session,
read_consistency_interval=interval,
) as db:
assert db.session is session
assert db.read_consistency_interval == interval
db.read_consistency_interval = timedelta(0)
assert db.read_consistency_interval == timedelta(0)
+277 -5
View File
@@ -60,6 +60,124 @@ def _blob_v2_table(db: DBConnection, name: str):
return table
@pytest.mark.parametrize(
"threshold",
[
"inline_size_threshold",
"dedicated_size_threshold",
"pack_file_size_threshold",
],
)
@pytest.mark.parametrize(
"value,error",
[
(-1, ValueError),
(True, TypeError),
(False, TypeError),
(1.5, TypeError),
("8192", TypeError),
(ctypes.c_size_t(-1).value + 1, OverflowError),
],
)
def test_blob_thresholds_reject_invalid_values(threshold, value, error):
# Also runs in CI with pylance actually uninstalled.
with pytest.raises(error, match=f"{threshold} must"):
lancedb.blob("blob", **{threshold: value})
@pytest.mark.parametrize(
"threshold", ["dedicated_size_threshold", "pack_file_size_threshold"]
)
def test_blob_thresholds_reject_zero_for_positive_thresholds(threshold):
with pytest.raises(ValueError, match=f"{threshold} must be positive"):
lancedb.blob("blob", **{threshold: 0})
def test_blob_thresholds_accept_integer_boundaries():
maximum = ctypes.c_size_t(-1).value
field = lancedb.blob(
"blob",
nullable=False,
inline_size_threshold=0,
dedicated_size_threshold=maximum,
pack_file_size_threshold=maximum,
)
assert not field.nullable
assert field.metadata == {
b"lance-encoding:blob-inline-size-threshold": b"0",
b"lance-encoding:blob-dedicated-size-threshold": str(maximum).encode(),
b"lance-encoding:blob-pack-file-size-threshold": str(maximum).encode(),
}
def test_blob_thresholds_control_writes_after_reopen(tmp_path):
# Keep this in test_table.py: CI runs it without pylance or pandas installed.
field = lancedb.blob(
"blob",
inline_size_threshold=4,
dedicated_size_threshold=12,
pack_file_size_threshold=8,
)
metadata = {
b"lance-encoding:blob-inline-size-threshold": b"4",
b"lance-encoding:blob-dedicated-size-threshold": b"12",
b"lance-encoding:blob-pack-file-size-threshold": b"8",
}
assert field.metadata == metadata
schema = pa.schema([pa.field("id", pa.int64()), field])
db = lancedb.connect(tmp_path)
table = db.create_table("blobs", schema=schema)
# Every payload fits in the default inline tier; custom thresholds must
# instead send the larger values to packed and dedicated storage.
rows = [
{"id": 0, "blob": b"a" * 2},
{"id": 1, "blob": b"b" * 8},
{"id": 2, "blob": b"c" * 16},
{"id": 3, "blob": b"d" * 8},
]
def check_writes(table, expected):
assert table.schema.field("blob").metadata.items() >= metadata.items()
hits = table.search().with_row_id(True).limit(10).to_arrow()
actual = {row["id"]: row for row in hits.to_pylist()}
assert actual.keys() == {row["id"] for row in expected}
for row in expected:
size = len(row["blob"])
descriptor = actual[row["id"]]["blob"]
assert descriptor["size"] == size
assert descriptor["kind"] == (2 if size > 12 else 1 if size > 4 else 0)
row_ids = [actual[row["id"]]["_rowid"] for row in expected]
assert table.fetch_blobs("blob", row_ids).to_pylist() == [
row["blob"] for row in expected
]
# Two packed values reach the configured sidecar limit independently.
packed = [
actual[row["id"]]["blob"] for row in expected if len(row["blob"]) == 8
]
assert len({desc["blob_id"] for desc in packed}) == len(packed)
table.add(rows)
check_writes(table, rows)
table = lancedb.connect(tmp_path).open_table("blobs")
check_writes(table, rows)
updates = [
{"id": 0, "blob": b"e" * 16},
{"id": 1, "blob": b"f" * 2},
{"id": 2, "blob": b"g" * 8},
{"id": 3, "blob": b"h" * 8},
{"id": 4, "blob": b"i" * 16},
]
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(updates)
)
assert result.num_updated_rows == 4
assert result.num_inserted_rows == 1
check_writes(lancedb.connect(tmp_path).open_table("blobs"), updates)
async def _blob_v2_table_async(db: AsyncConnection, name: str):
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
table = await db.create_table(name, schema=schema)
@@ -1634,6 +1752,28 @@ def test_create_index_async_returns_done_job(mem_db: DBConnection):
job.cancel()
@pytest.mark.parametrize(
"rq_kwargs", [{}, {"num_bits": None}, {"num_bits": 1}, {"num_bits": 4}]
)
@patch("lancedb.table.LanceTable.to_lance")
def test_create_ivfrq_index_accelerator(mock_to_lance, mem_db: DBConnection, rq_kwargs):
from lancedb.index import IvfRq
table = mem_db.create_table("test", [{"vector": [1.0] * 8}])
torch = MagicMock()
torch.cuda.is_available.return_value = True
with patch.dict(sys.modules, {"torch": torch}):
table.create_index("vector", config=IvfRq(accelerator="cuda", **rq_kwargs))
kwargs = mock_to_lance.return_value.create_index.call_args.kwargs
assert kwargs["index_type"] == "IVF_RQ"
assert kwargs["accelerator"] == "cuda"
if rq_kwargs.get("num_bits") is None:
assert "num_bits" not in kwargs
else:
assert kwargs["num_bits"] == rq_kwargs["num_bits"]
@pytest.mark.parametrize("method", ["create_index", "create_index_async"])
@pytest.mark.parametrize("wait_timeout", ["invalid", 1, []])
def test_create_index_invalid_wait_timeout(mem_db: DBConnection, method, wait_timeout):
@@ -2053,6 +2193,71 @@ def test_add_with_nans(mem_db: DBConnection):
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
@pytest.mark.parametrize("operation", ["create", "append", "overwrite"])
@pytest.mark.parametrize("input_type", ["pandas", "rows", "arrow"])
@pytest.mark.parametrize("on_bad_vectors", [None, "error", "drop", "fill", "null"])
def test_write_vectors_with_null_elements(
mem_db: DBConnection, operation, input_type, on_bad_vectors
):
schema = pa.schema({"vector": pa.list_(pa.float32(), 2), "id": pa.int64()})
vectors = [[3.0, 4.0], [np.nan, 1.0], [np.nan, np.nan]]
if input_type == "pandas":
pd = pytest.importorskip("pandas")
data = pd.DataFrame(
{
"vector": [np.array(v, dtype=np.float32) for v in vectors],
"id": [0, 1, 2],
}
)
else:
# Pandas turns the NaNs above into null child elements. Explicit nulls
# in other input formats must follow the same bad-vector policy.
data = [
{"vector": [3.0, 4.0], "id": 0},
{"vector": [None, 1.0], "id": 1},
{"vector": [None, None], "id": 2},
]
if input_type == "arrow":
data = pa.Table.from_pylist(data, schema=schema)
options = {} if on_bad_vectors is None else {"on_bad_vectors": on_bad_vectors}
if on_bad_vectors == "fill":
options["fill_value"] = 42.0
table = None
if operation != "create":
table = mem_db.create_table("test", schema=schema)
def write():
if operation == "create":
return mem_db.create_table("test", data=data, schema=schema, **options)
table.add(data, mode=operation, **options)
return table
if on_bad_vectors in (None, "error"):
# Create/overwrite can wrap errors from lazy Arrow readers in RuntimeError.
expected_error = (
ValueError if operation == "append" else (ValueError, RuntimeError)
)
with pytest.raises(expected_error, match="NaN|null"):
write()
if table is not None:
assert table.count_rows() == 0
return
table = write()
if on_bad_vectors == "drop":
expected = [{"vector": [3.0, 4.0], "id": 0}]
else:
expected = [
{"vector": [3.0, 4.0], "id": 0},
{"vector": [42.0, 1.0] if on_bad_vectors == "fill" else None, "id": 1},
{"vector": [42.0, 42.0] if on_bad_vectors == "fill" else None, "id": 2},
]
assert table.to_arrow().to_pylist() == expected
def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection):
class Schema(LanceModel):
text: str
@@ -2979,6 +3184,57 @@ def test_update_with_arrow_scalar(mem_db: DBConnection):
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0, 3.0, 4.0]]
@pytest.mark.asyncio
@pytest.mark.parametrize("async_api", [False, True])
@pytest.mark.parametrize("entries", [[("o'brien", 4), ("nil", None)], [], None])
async def test_update_map_arrow_scalar(tmp_path, async_api, entries):
map_type = pa.map_(pa.string(), pa.int64())
data = pa.table(
{
"id": [1, 2],
"attrs": pa.array([[("original", 1)], entries], type=map_type),
}
)
if async_api:
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("maps", data)
source = await table.query().where("id = 2").to_arrow()
result = await table.update({"attrs": source["attrs"][0]}, where="id = 1")
actual = await table.query().to_arrow()
else:
db = lancedb.connect(tmp_path)
table = db.create_table("maps", data)
source = table.search().where("id = 2").to_arrow()
result = table.update(values={"attrs": source["attrs"][0]}, where="id = 1")
actual = table.to_arrow()
assert result.rows_updated == 1
assert actual.schema.field("attrs").type == map_type
assert actual.sort_by("id")["attrs"].to_pylist() == [entries, entries]
@pytest.mark.parametrize(
"map_type, entries",
[
(pa.map_(pa.int64(), pa.string()), []),
(
pa.map_(pa.string(), pa.map_(pa.string(), pa.int64())),
[("outer", [("inner", 3)]), ("null", None)],
),
],
)
def test_update_typed_map_arrow_scalar(mem_db: DBConnection, map_type, entries):
table = mem_db.create_table(
"maps",
pa.table({"id": [1, 2], "attrs": pa.array([None, entries], type=map_type)}),
)
scalar = table.search().where("id = 2").to_arrow()["attrs"][0]
result = table.update(values={"attrs": scalar}, where="id = 1")
actual = table.to_arrow().sort_by("id")["attrs"]
assert result.rows_updated == 1
assert actual.type == map_type
assert actual.to_pylist() == [entries, entries]
@pytest.mark.parametrize("use_arrow_scalar", [False, True])
@pytest.mark.parametrize(
"decimal_type, value",
@@ -3964,26 +4220,42 @@ def test_create_f16_table_from_arrow_data(mem_db: DBConnection):
assert "s-2" in expected["text"].to_pylist()
def test_create_f16_table(mem_db: DBConnection):
@pytest.mark.parametrize("accelerator", [None, "cuda"])
def test_create_f16_table(tmp_path, accelerator):
if accelerator == "cuda":
torch = pytest.importorskip("torch")
if not torch.cuda.is_available():
pytest.skip("CUDA not available")
class MyTable(LanceModel):
text: str
vector: Vector(32, value_type=pa.float16())
rng = np.random.default_rng(42)
vectors = rng.standard_normal((512, 32)).astype(np.float16)
df = pa.table(
{
"text": [f"s-{i}" for i in range(512)],
"vector": [rng.standard_normal(32).astype(np.float16) for _ in range(512)],
"vector": list(vectors),
}
)
table = mem_db.create_table(
db = lancedb.connect(tmp_path)
table = db.create_table(
"f16_tbl",
schema=MyTable,
)
table.add(df)
table.create_index(num_partitions=2, num_sub_vectors=2)
table.create_index(
"vector",
config=IvfPq(
num_partitions=2,
num_sub_vectors=2,
accelerator=accelerator,
),
)
query = df["vector"][2].as_py()
# Match the issue's float64 query against an explicitly typed float16 column.
query = vectors[2].astype(np.float64)
expected = table.search(query).limit(2).to_arrow()
assert "s-2" in expected["text"].to_pylist()
@@ -0,0 +1,253 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from datetime import timedelta
import inspect
from unittest.mock import AsyncMock, MagicMock
import lancedb
from lancedb.expr import col
from lancedb.index import BTree, FTS, HnswSq, IvfPq
from lancedb.remote.table import RemoteTable
from lancedb.table import LanceTable, Table
import pyarrow as pa
import pytest
@pytest.mark.parametrize(
"method",
[
"delete",
"drop_index",
"index_stats",
"create_index",
"create_scalar_index",
"create_fts_index",
"search",
"cleanup_old_versions",
"compact_files",
"optimize",
"to_polars",
],
)
def test_sync_table_signatures(method):
"""A Table call must bind the same arguments on either backend."""
def parameters(cls):
return [
(p.name, p.kind, p.default)
for p in inspect.signature(getattr(cls, method)).parameters.values()
]
assert parameters(Table) == parameters(LanceTable) == parameters(RemoteTable)
@pytest.fixture(params=[LanceTable, RemoteTable], ids=["local", "remote"])
def sync_table(request, monkeypatch):
"""Exercise real sync wrappers with a recording async backend."""
handle = MagicMock()
handle.name = "test"
handle.create_index = AsyncMock()
handle.delete = AsyncMock()
handle.drop_index = AsyncMock()
handle.index_stats = AsyncMock()
if request.param is RemoteTable:
table = RemoteTable(handle, "test_db")
else:
table = LanceTable.__new__(LanceTable)
table._table = handle
monkeypatch.setattr(table, "_ensure_no_legacy_fts_index", lambda: None)
return table, handle
def test_sync_keyword_names(sync_table):
table, handle = sync_table
predicate = col("id") == 1
table.delete(where=predicate)
handle.delete.assert_awaited_once()
expected = predicate._inner if isinstance(table, LanceTable) else predicate
assert handle.delete.await_args.args[0] is expected
table.drop_index(name="vector_idx")
handle.drop_index.assert_awaited_once_with("vector_idx")
table.index_stats(index_name="vector_idx")
handle.index_stats.assert_awaited_once_with("vector_idx")
def test_legacy_vector_positional_arguments(sync_table):
table, handle = sync_table
with pytest.warns(DeprecationWarning):
table.create_index("cosine", 2)
args, kwargs = handle.create_index.await_args
assert args == ("vector",)
assert isinstance(kwargs["config"], IvfPq)
assert kwargs["config"].distance_type == "cosine"
assert kwargs["config"].num_partitions == 2
assert kwargs["replace"] is True
def test_vector_tuning_options(sync_table):
table, handle = sync_table
timeout = timedelta(seconds=2)
with pytest.warns(DeprecationWarning):
table.create_index(
index_type="IVF_HNSW_SQ",
num_partitions=1,
m=10,
ef_construction=50,
max_iterations=12,
sample_rate=64,
target_partition_size=1000,
wait_timeout=timeout,
name="custom_idx",
replace=False,
train=False,
)
kwargs = handle.create_index.await_args.kwargs
config = kwargs.pop("config")
assert isinstance(config, HnswSq)
for name, value in {
"num_partitions": 1,
"m": 10,
"ef_construction": 50,
"max_iterations": 12,
"sample_rate": 64,
"target_partition_size": 1000,
}.items():
assert getattr(config, name) == value
assert kwargs == dict(
wait_timeout=timeout, name="custom_idx", replace=False, train=False
)
def test_scalar_defaults_and_wait_timeout(sync_table):
table, handle = sync_table
timeout = timedelta(seconds=2)
with pytest.warns(DeprecationWarning):
table.create_scalar_index("id", wait_timeout=timeout)
kwargs = handle.create_index.await_args.kwargs
assert isinstance(kwargs["config"], BTree)
assert kwargs["replace"] is True
assert kwargs["wait_timeout"] == timeout
def test_fts_tokenizer_and_wait_timeout(sync_table):
table, handle = sync_table
timeout = timedelta(seconds=2)
with pytest.warns(DeprecationWarning):
table.create_fts_index(
field_names="text",
tokenizer_name="en_stem",
replace=True,
wait_timeout=timeout,
custom_stop_words=["example"],
)
args, kwargs = handle.create_index.await_args
assert args == ("text",)
assert kwargs["replace"] is True
assert kwargs["wait_timeout"] == timeout
config = kwargs["config"]
assert isinstance(config, FTS)
assert config.language == "English"
assert config.stem is True
assert config.custom_stop_words == ["example"]
@pytest.mark.parametrize(
"kwargs, message",
[
({"use_tantivy": True}, "Tantivy-based FTS has been removed"),
({"ordering_field_names": "id"}, "ordering_field_names"),
({"writer_heap_size": 123}, "writer_heap_size"),
({"tokenizer_name": "invalid"}, "Invalid tokenizer name"),
],
)
def test_fts_unsupported_options(sync_table, kwargs, message):
table, handle = sync_table
with pytest.warns(DeprecationWarning), pytest.raises(ValueError, match=message):
table.create_fts_index("text", **kwargs)
handle.create_index.assert_not_awaited()
@pytest.mark.parametrize(
"method, kwargs",
[
(
"cleanup_old_versions",
{"older_than": timedelta(days=1), "delete_unverified": True},
),
("compact_files", {"target_rows_per_fragment": 1000}),
],
)
def test_remote_maintenance_options_warn(method, kwargs):
table = RemoteTable(MagicMock(name="test"), "test_db")
with pytest.warns(UserWarning, match="no-op"):
getattr(table, method)(**kwargs)
def test_remote_optimize_options():
handle = MagicMock()
handle.optimize = AsyncMock(side_effect=NotImplementedError("remote optimize"))
table = RemoteTable(handle, "test_db")
age = timedelta(days=1)
with pytest.raises(NotImplementedError, match="remote optimize"):
table.optimize(cleanup_older_than=age, delete_unverified=True, retrain=True)
handle.optimize.assert_awaited_once_with(
cleanup_older_than=age, delete_unverified=True, retrain=True
)
@pytest.mark.parametrize(
"method, old, new, value",
[
("delete", "predicate", "where", "id = 1"),
("drop_index", "index_name", "name", "vector_idx"),
("index_stats", "index_uuid", "index_name", "vector_idx"),
("create_fts_index", "column", "field_names", "text"),
],
)
def test_deprecated_remote_keywords(method, old, new, value):
handle = MagicMock()
handle.name = "test"
target = "create_index" if method == "create_fts_index" else method
setattr(handle, target, AsyncMock())
table = RemoteTable(handle, "test_db")
with pytest.warns(DeprecationWarning, match=f"{old} is deprecated"):
getattr(table, method)(**{old: value})
assert getattr(handle, target).await_args.args == (value,)
with pytest.raises(TypeError, match="Cannot specify both"):
getattr(table, method)(**{old: value, new: value})
@pytest.mark.parametrize("remote_wrapper", [False, True], ids=["local", "remote"])
def test_sync_search_and_polars(tmp_path, remote_wrapper):
"""Use the native engine through both sync wrappers, without a cloud service."""
local = lancedb.connect(tmp_path).create_table(
"test",
[{"id": 1, "text": "quick fox", "vector": [1.0, 1.0]}],
)
table = RemoteTable(local._table, "test_db") if remote_wrapper else local
table.create_index("text", config=FTS())
assert table.search("quick", query_type="fts").to_arrow()["id"].to_pylist() == [1]
# Both backends must accept the deprecated ordering argument in the same slot.
with pytest.warns(DeprecationWarning, match="ordering_field_name"):
table.search("quick", None, "fts", "id", ["text"])
fast = table.search([1.0, 1.0], fast_search=True).to_arrow()
assert len(fast) == 0 # No vector index: fast_search excludes unindexed rows.
assert table.search([1.0, 1.0]).to_arrow()["id"].to_pylist() == [1]
if remote_wrapper:
with pytest.raises(NotImplementedError, match=r"table.search\(\).to_polars"):
table.to_polars(batch_size=1)
assert table.search().to_polars()["id"].to_list() == [1]
else:
assert table.to_polars(batch_size=1).collect()["id"].to_list() == [1]
table.delete(where="id = 1")
assert table.count_rows() == 0
def test_index_replacement_defaults(tmp_path):
local = lancedb.connect(tmp_path).create_table("test", pa.table({"id": [1]}))
table = RemoteTable(local._table, "test_db")
for sync in [local, table]:
with pytest.warns(DeprecationWarning):
sync.create_scalar_index("id")
assert len(list(table.list_indices())) == 1
+47 -7
View File
@@ -406,7 +406,7 @@ def test_handle_bad_vectors_jagged(on_bad_vectors):
).read_all()
if on_bad_vectors == "drop":
expected = pa.array([[1.0, 2.0], [4.0, 5.0], [None, 9.0]])
expected = pa.array([[1.0, 2.0], [4.0, 5.0]])
elif on_bad_vectors == "fill":
expected = pa.array(
[
@@ -414,12 +414,12 @@ def test_handle_bad_vectors_jagged(on_bad_vectors):
[3.0, 42.0],
[4.0, 5.0],
[6.0, 7.0],
[None, 9.0],
[42.0, 9.0],
[42.0, 42.0],
]
)
elif on_bad_vectors == "null":
expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, [None, 9.0], None])
expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, None, None])
assert output["vector"].combine_chunks() == expected
@@ -430,17 +430,17 @@ def test_handle_bad_vectors_jagged(on_bad_vectors):
(
pa.list_(pa.float64()),
[[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]],
[[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]],
[[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [42.0, 3.0], [4.0, 5.0]],
),
(
pa.large_list(pa.float64()),
[[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]],
[[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]],
[[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [42.0, 3.0], [4.0, 5.0]],
),
(
pa.list_(pa.float64(), 2),
[[1.0, float("nan")], None, [None, 3.0]],
[[1.0, 42.0], [42.0, 42.0], [None, 3.0]],
[[1.0, 42.0], [42.0, 42.0], [42.0, 3.0]],
),
],
)
@@ -467,7 +467,7 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
).read_all()
output = exception_output(e)
assert output == (
"ValueError: Vector column 'vector' has NaNs. Set "
"ValueError: Vector column 'vector' has NaNs or null values. Set "
"on_bad_vectors='drop' to remove them, set on_bad_vectors='fill' "
"and fill_value=<value> to replace them, or set on_bad_vectors='null' "
"to replace them with null."
@@ -490,6 +490,46 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
assert output["vector"].combine_chunks() == expected
@pytest.mark.parametrize("value_type", [pa.float16(), pa.float32(), pa.float64()])
@pytest.mark.parametrize(
"list_type", [pa.list_, pa.large_list, lambda t: pa.list_(t, 2)]
)
@pytest.mark.parametrize("chunked", [False, True])
def test_has_nan_values_with_null_elements(value_type, list_type, chunked):
arr = pa.array(
[[None, None], [1.0, 2.0], [None, 3.0], None, [4.0, float("nan")], [5.0, 6.0]],
type=list_type(value_type),
).slice(1, 5)
if chunked:
arr = pa.chunked_array([arr.slice(0, 2), arr.slice(2)])
# Whole null vectors have no components to inspect. Their handling is
# separate from missing or NaN components in a valid vector.
assert has_nan_values(arr).to_pylist() == [False, True, False, True, False]
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
def test_sanitize_vectors_with_only_null_elements(on_bad_vectors):
data = pa.table({"vector": [[None, None]]})
schema = pa.schema({"vector": pa.list_(pa.float32(), 2)})
def sanitize():
return _sanitize_data(
data, schema, on_bad_vectors=on_bad_vectors, fill_value=42.0
).read_all()
if on_bad_vectors == "error":
with pytest.raises(ValueError, match="null"):
sanitize()
return
expected = {
"drop": [],
"fill": [{"vector": [42.0, 42.0]}],
"null": [{"vector": None}],
}
assert sanitize().to_pylist() == expected[on_bad_vectors]
@pytest.mark.parametrize("dtype", [pa.float16(), pa.float32(), pa.float64()])
@pytest.mark.parametrize("large", [False, True])
@pytest.mark.parametrize("chunked", [False, True])
+16 -40
View File
@@ -590,29 +590,17 @@ impl Connection {
#[pyo3(signature = (namespace_path=None, page_token=None, limit=None))]
pub fn list_namespaces(
self_: PyRef<'_, Self>,
&self,
namespace_path: Option<Vec<String>>,
page_token: Option<String>,
limit: Option<u32>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let py = self_.py();
future_into_py(py, async move {
use lance_namespace::models::ListNamespacesRequest;
let request = ListNamespacesRequest {
id: namespace_path,
page_token,
limit: limit.map(|l| l as i32),
..Default::default()
};
let response = inner.list_namespaces(request).await.infer_error()?;
Python::attach(|py| -> PyResult<Py<PyDict>> {
let dict = PyDict::new(py);
dict.set_item("namespaces", response.namespaces)?;
dict.set_item("page_token", response.page_token)?;
Ok(dict.unbind())
})
})
) -> PyResult<crate::listing::NameListing> {
Ok(crate::listing::NameListing::new(
self.get_inner()?.list_namespaces(
&namespace_path.unwrap_or_default(),
crate::listing::options(page_token, limit),
),
))
}
#[pyo3(signature = (namespace_path, mode=None, properties=None))]
@@ -725,29 +713,17 @@ impl Connection {
#[pyo3(signature = (namespace_path=None, page_token=None, limit=None))]
pub fn list_tables(
self_: PyRef<'_, Self>,
&self,
namespace_path: Option<Vec<String>>,
page_token: Option<String>,
limit: Option<u32>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let py = self_.py();
future_into_py(py, async move {
use lance_namespace::models::ListTablesRequest;
let request = ListTablesRequest {
id: namespace_path,
page_token,
limit: limit.map(|l| l as i32),
..Default::default()
};
let response = inner.list_tables(request).await.infer_error()?;
Python::attach(|py| -> PyResult<Py<PyDict>> {
let dict = PyDict::new(py);
dict.set_item("tables", response.tables)?;
dict.set_item("page_token", response.page_token)?;
Ok(dict.unbind())
})
})
) -> PyResult<crate::listing::NameListing> {
Ok(crate::listing::NameListing::new(
self.get_inner()?.list_tables(
&namespace_path.unwrap_or_default(),
crate::listing::options(page_token, limit),
),
))
}
/// Get the configuration for constructing an equivalent namespace client.
+5 -3
View File
@@ -136,8 +136,10 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
let mut ivf_rq_builder = IvfRqIndexBuilder::default()
.distance_type(distance_type)
.max_iterations(params.max_iterations)
.sample_rate(params.sample_rate)
.num_bits(params.num_bits);
.sample_rate(params.sample_rate);
if let Some(num_bits) = params.num_bits {
ivf_rq_builder = ivf_rq_builder.num_bits(num_bits);
}
if let Some(num_partitions) = params.num_partitions {
ivf_rq_builder = ivf_rq_builder.num_partitions(num_partitions);
}
@@ -265,7 +267,7 @@ struct IvfSqParams {
struct IvfRqParams {
distance_type: String,
num_partitions: Option<u32>,
num_bits: u32,
num_bits: Option<u32>,
max_iterations: u32,
sample_rate: u32,
target_partition_size: Option<u32>,
+5
View File
@@ -81,6 +81,11 @@ impl Job {
self.inner.id().map(str::to_string)
}
#[getter]
pub fn job_type(&self) -> Option<String> {
self.inner.job_type()
}
pub fn status(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(
+28 -2
View File
@@ -3,10 +3,11 @@
use crate::error::PythonErrorExt;
use crate::runtime::future_into_py;
use futures::{StreamExt, future::poll_fn};
use futures::{Future, StreamExt, future::poll_fn};
use lancedb::listing::{Listing, ListingOptions};
use pyo3::exceptions::PyStopAsyncIteration;
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pymethods};
use pyo3::types::{PyDict, PyDictMethods};
use pyo3::{Bound, Py, PyAny, PyRef, PyResult, Python, pyclass, pymethods};
use std::sync::{Arc, Mutex};
pub fn options(page_token: Option<String>, page_limit: Option<u32>) -> ListingOptions {
@@ -52,6 +53,31 @@ impl NameListing {
Ok(value)
})
}
fn next_page(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
let next_lock = self_.next_lock.clone();
future_into_py(self_.py(), async move {
let _guard = next_lock.lock().await;
let page = futures::future::poll_fn(|cx| {
let mut inner = inner.lock().unwrap();
let future = inner.next_page();
futures::pin_mut!(future);
future.poll(cx)
})
.await
.infer_error()?;
Python::attach(|py| -> PyResult<Option<Py<PyDict>>> {
page.map(|page| {
let dict = PyDict::new(py);
dict.set_item("items", page.items)?;
dict.set_item("page_token", page.page_token)?;
Ok(dict.unbind())
})
.transpose()
})
})
}
}
#[pyclass]
+18 -1
View File
@@ -24,7 +24,7 @@ use arrow::{
ffi_stream::ArrowArrayStreamReader,
pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
};
use lancedb::blob::{BlobFile, BlobRangeRequest};
use lancedb::blob::{BlobDescriptor, BlobFile, BlobRangeRequest};
use lancedb::index::scalar::FtsIndexBuilder;
use lancedb::table::{
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken,
@@ -1551,6 +1551,23 @@ impl Table {
})
}
#[pyo3(signature = (descriptor))]
pub fn open_blob_file<'py>(
self_: PyRef<'py, Self>,
descriptor: &str,
) -> PyResult<Bound<'py, PyAny>> {
let descriptor: BlobDescriptor = serde_json::from_str(descriptor)
.map_err(|error| PyValueError::new_err(format!("invalid blob descriptor: {error}")))?;
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let file = inner.open_blob_file(&descriptor).await.infer_error()?;
Ok(PyBlobFile {
inner: Arc::new(file),
closed: AtomicBool::new(false),
})
})
}
/// Optimize the on-disk data by compacting and pruning old data, for better performance.
#[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))]
pub fn optimize(
+9 -9
View File
@@ -2006,8 +2006,8 @@ requires-dist = [
{ name = "pyarrow", marker = "extra == 'tests'", specifier = "<25" },
{ name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" },
{ name = "pydantic", specifier = ">=2.7.4,<3" },
{ name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" },
{ name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b12" },
{ name = "pylance", marker = "extra == 'pylance'", specifier = ">=15.0.0b1" },
{ name = "pylance", marker = "extra == 'tests'", specifier = "==15.0.0b1" },
{ name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" },
{ name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" },
{ name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" },
@@ -3938,7 +3938,7 @@ crypto = [
[[package]]
name = "pylance"
version = "14.0.0b12"
version = "15.0.0b1"
source = { registry = "https://pypi.fury.io/lance-format/" }
dependencies = [
{ name = "lance-namespace" },
@@ -3948,12 +3948,12 @@ dependencies = [
{ name = "pyarrow" },
]
wheels = [
{ url = "https://pypi.fury.io/lance-format/-/ver_1Ixasc/pylance-14.0.0b12-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:84350a3c5c1eea3398f4c4fad7cfe85b831d831a5e43b5ca2a783dbdb0e11969" },
{ url = "https://pypi.fury.io/lance-format/-/ver_sCEkG/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:76b386c182a6a8e7b86fbef28f84d1dd39e4626ac8681a0cf36f40a45dd7cb87" },
{ url = "https://pypi.fury.io/lance-format/-/ver_1vdT1C/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6aa548c7937fcdfa8f1943a2d59041a42f8b500199880dbb509c7195fa253381" },
{ url = "https://pypi.fury.io/lance-format/-/ver_2fXu1V/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:79f82d6a6c50f37aab338470e3733436e6a27fbb5972d29fcf453ece56d3b3c6" },
{ url = "https://pypi.fury.io/lance-format/-/ver_24UavN/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:de0305999257002785f60169e6926896584d66f1da050bb5afd1d48c6e3d9bb9" },
{ url = "https://pypi.fury.io/lance-format/-/ver_XNynX/pylance-14.0.0b12-cp310-abi3-win_amd64.whl", hash = "sha256:4d0cdb689e898ce71e8e434d00a3b97639cea8fb97da21f1c921543be363268f" },
{ url = "https://pypi.fury.io/lance-format/-/ver_23L1Nd/pylance-15.0.0b1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:19605d3daca1224cc0db079054be5d4eec403dcf286bf4617430bfc259ab0a52" },
{ url = "https://pypi.fury.io/lance-format/-/ver_1Ho3RK/pylance-15.0.0b1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:be22b0b6f1e1d250ffaa82b449e982c817db33103c483d1005020285310623cf" },
{ url = "https://pypi.fury.io/lance-format/-/ver_1QmyF0/pylance-15.0.0b1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a5b2761d60b1b9db112058213f1f7b2de82b08cbe65888fab08b0a56e702482e" },
{ url = "https://pypi.fury.io/lance-format/-/ver_1in2Ph/pylance-15.0.0b1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:294a0fedae312ed32d1b66cddbf7e69770ec9b4c82a873af60c619ed88450718" },
{ url = "https://pypi.fury.io/lance-format/-/ver_21xpLH/pylance-15.0.0b1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ff20a16703d4d3f3bc675b25ce1871ba1d54137929dc46305bffdd07f7b35a6b" },
{ url = "https://pypi.fury.io/lance-format/-/ver_AAGzF/pylance-15.0.0b1-cp310-abi3-win_amd64.whl", hash = "sha256:3b134e24a5becfc940ca134cf78aa3bae8b80ec40f8927a1b856273ab76213f2" },
]
[[package]]