mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-28 06:26:15 +00:00
Merge remote-tracking branch 'origin/main' into xuanwo/pr3851-merge
This commit is contained in:
+1
-2
@@ -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
|
||||
|
||||
@@ -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"])
|
||||
```
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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`.
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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:
|
||||
|
||||
+446
-272
File diff suppressed because it is too large
Load Diff
@@ -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())
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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])
|
||||
|
||||
@@ -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
|
||||
@@ -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())]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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 == []
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
@@ -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
@@ -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>,
|
||||
|
||||
@@ -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
@@ -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
@@ -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(
|
||||
|
||||
Generated
+9
-9
@@ -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]]
|
||||
|
||||
Reference in New Issue
Block a user