rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-09-03 03:58:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into gatekeeper/fix-2468-1

Browse Source 
# Conflicts:
#	rust/lancedb/src/table/query.rs
This commit is contained in:
Gatefixer
2026-08-26 02:43:02 +00:00
parent de54965ece 9b825c5f29
commit 6b90ddd1b6
251 changed files with 34807 additions and 3764 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/Cargo.toml
+14 -15
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb-python"
version = "0.37.1-beta.0"
version = "0.38.0-beta.10"
publish = false
edition.workspace = true
description = "Python bindings for LanceDB"
@@ -15,10 +15,10 @@ name = "_lancedb"
crate-type = ["cdylib"]
[dependencies]
arrow = { version = "58.0.0", features = ["pyarrow"] }
async-trait = "0.1"
bytes = "1"
lancedb = { path = "../rust/lancedb", default-features = false }
arrow = { workspace = true, features = ["pyarrow"] }
async-trait.workspace = true
bytes.workspace = true
lancedb.workspace = true
datafusion-common.workspace = true
lance-core.workspace = true
lance-namespace.workspace = true
@@ -26,25 +26,24 @@ lance-namespace-impls.workspace = true
lance-io.workspace = true
env_logger.workspace = true
log.workspace = true
pyo3 = { version = "0.28", features = ["extension-module", "abi3-py39", "chrono"] }
chrono = { version = "0.4", default-features = false, features = ["clock"] }
# Maturin enables extension-module mode for Python builds. Keeping it out of
# Cargo features lets Rust unit tests link against libpython.
pyo3 = { version = "0.28", features = ["abi3-py310", "chrono"] }
chrono.workspace = true
pyo3-async-runtimes = { version = "0.28", features = [
"attributes",
"tokio-runtime",
] }
pin-project = "1.1.5"
pin-project.workspace = true
futures.workspace = true
serde = "1"
serde_json = "1"
serde.workspace = true
serde_json.workspace = true
snafu.workspace = true
tokio = { version = "1.40", features = ["sync", "rt-multi-thread"] }
tokio.workspace = true
libc = "0.2"
[build-dependencies]
pyo3-build-config = { version = "0.28", features = [
"extension-module",
"abi3-py39",
] }
pyo3-build-config = { version = "0.28", features = ["abi3-py310"] }
[features]
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/cos", "lancedb/goosefs", "lancedb/metrics-otel"]
python/README.md
+19
View File
@@ -38,6 +38,25 @@ 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
```
### Threading in CPU-limited containers
LanceDB uses separate pools for compute work and storage I/O. On a container with
two visible CPUs, current releases intentionally use one compute worker by default;
no manual configuration is needed. If every query logs an I/O core reservation
warning on a two-CPU container, upgrade from LanceDB 0.21.1 or earlier.
The two commonly tuned environment variables control different resources:
- `LANCE_CPU_THREADS` overrides the number of compute workers. One worker is the
appropriate setting for a two-CPU container when an explicit override is needed.
- `LANCE_IO_THREADS` controls concurrent storage operations, not reserved CPU
cores. Its default can be greater than the number of CPUs because I/O workers
spend much of their time waiting for storage.
Keep the defaults unless measurements show that the workload benefits from an
override. See the [Lance threading model](https://lance.org/guide/performance/#threading-model)
for the current defaults and tuning guidance.
## Usage
### Basic Example
python/pyproject.toml
+7 -3
View File
@@ -8,7 +8,7 @@ dependencies = [
"overrides>=0.7; python_version<'3.12'",
"packaging>=23.0",
"pyarrow>=16",
"pydantic>=1.10",
"pydantic>=2.7.4,<3",
"tqdm>=4.27.0",
"lance-namespace>=0.3.2"
]
@@ -60,7 +60,7 @@ tests = [
"pytest-asyncio>=0.21",
"duckdb>=0.9.0",
"pytz>=2023.3",
"polars>=0.19, <=1.3.0",
"polars>=0.19, <=1.32.3",
"pyarrow<25",
"pyarrow-stubs>=16.0",
"pylance==9.0.0rc1",
@@ -101,9 +101,12 @@ azure = ["adlfs>=2024.2.0"]
[tool.maturin]
python-source = "python"
module-name = "lancedb._lancedb"
# uv installs the project as an editable package before `uv run`, so keep that
# bootstrap build consistent with `maturin develop`.
editable-profile = "dev"
[build-system]
requires = ["maturin>=1.4"]
requires = ["maturin>=1.10"]
build-backend = "maturin"
[tool.ruff.lint]
@@ -140,6 +143,7 @@ include = [
"python/lancedb/remote/errors.py",
"python/lancedb/embeddings/__init__.py",
"python/lancedb/_lancedb.pyi",
"python/type_tests/connect.py",
]
exclude = ["python/tests/"]
pythonVersion = "3.13"
python/python/lancedb/__init__.py
+42
View File
@@ -12,6 +12,7 @@ __version__ = importlib.metadata.version("lancedb")
from ._lancedb import connect as lancedb_connect
from ._lancedb import FtsToken
from ._lancedb import LsmWriteSpec
from ._lancedb import tokenize as _tokenize
from .common import URI, sanitize_uri
from urllib.parse import urlparse
@@ -21,6 +22,22 @@ from .remote.db import RemoteDBConnection
from .expr import Expr, col, lit, func
from .schema import blob, vector, BlobType
from .job import AsyncJob, Job
from .functions import (
FunctionArtifactRequest as FunctionArtifactRequest,
FunctionApplication as FunctionApplication,
FunctionBinding as FunctionBinding,
FunctionRegistrationRequest as FunctionRegistrationRequest,
FunctionVersion as FunctionVersion,
PythonRuntimeSpec as PythonRuntimeSpec,
RefreshColumnResult as RefreshColumnResult,
UdfDefinition as UdfDefinition,
udf as udf,
)
from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
MaterializedViewDefinition,
)
from .table import AsyncTable, Table
from .types import BaseTokenizerType
from ._lancedb import Session
@@ -162,6 +179,18 @@ def connect(
... },
... )
For Azure Blob Storage, credentials can be passed directly without setting
environment variables:
>>> azure_storage_options = {
... "account_name": "some-account",
... "account_key": "some-key",
... }
>>> db = lancedb.connect( # doctest: +SKIP
... "az://my-container/my-database",
... storage_options=azure_storage_options,
... )
For tests and temporary data, use an in-memory database:
>>> db = lancedb.connect("memory://")
@@ -448,6 +477,10 @@ async def connect_async(
--------
>>> import lancedb
>>> azure_storage_options = {
... "account_name": "some-account",
... "account_key": "some-key",
... }
>>> async def doctest_example():
... # For a local directory, provide a path to the database
... db = await lancedb.connect_async("~/.lancedb")
@@ -455,6 +488,11 @@ async def connect_async(
... db = await lancedb.connect_async("s3://my-bucket/lancedb",
... storage_options={
... "aws_access_key_id": "***"})
... # Azure credentials can also be passed directly
... db = await lancedb.connect_async(
... "az://my-container/my-database",
... storage_options=azure_storage_options,
... )
... # For tests and temporary data, use an in-memory database
... db = await lancedb.connect_async("memory://")
... # Connect to LanceDB cloud
@@ -495,6 +533,9 @@ async def connect_async(
__all__ = [
"AsyncMaterializedView",
"MaterializedView",
"MaterializedViewDefinition",
"connect",
"connect_async",
"tokenize",
@@ -518,6 +559,7 @@ __all__ = [
"Job",
"LanceDBConnection",
"LanceNamespaceDBConnection",
"LsmWriteSpec",
"RemoteDBConnection",
"Session",
"Table",
python/python/lancedb/_lancedb.pyi
+49 -6
View File
@@ -147,6 +147,8 @@ class Connection(object):
limit: Optional[int],
) -> list[str]: ... # Deprecated: Use list_tables instead
def job(self, job_id: str) -> Job: ...
async def create_function_async(self, request_json: str) -> Job: ...
async def get_function(self, name: str, version: str) -> str: ...
async def list_jobs(self) -> List[JobInfo]: ...
async def get_job(self, job_id: str) -> Optional[JobDescription]: ...
async def cancel_job(self, job_id: str) -> bool: ...
@@ -195,9 +197,21 @@ class Connection(object):
cur_namespace_path: Optional[List[str]] = None,
new_namespace_path: Optional[List[str]] = None,
) -> None: ...
async def create_materialized_view(
self,
name: str,
source: str,
projections: Optional[List[Tuple[str, str]]] = None,
filter: Optional[str] = None,
limit: Optional[int] = None,
) -> Table: ...
async def list_materialized_views(self) -> List[str]: ...
async def drop_table(
self, name: str, namespace_path: Optional[List[str]] = None
) -> None: ...
async def drop_table_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Job: ...
async def drop_all_tables(
self, namespace_path: Optional[List[str]] = None
) -> None: ...
@@ -220,7 +234,7 @@ class Job:
@property
def id(self) -> Optional[str]: ...
async def status(self) -> str: ...
async def wait(self) -> None: ...
async def wait(self) -> Optional[str]: ...
async def cancel(self) -> None: ...
class JobInfo:
@@ -269,6 +283,7 @@ class Table:
mode: Literal["append", "overwrite"],
progress: Optional[Any] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult: ...
async def update(
self, updates: Dict[str, str], where: Optional[str]
@@ -335,6 +350,17 @@ class Table:
) -> list[FtsToken]: ...
async def delete(self, filter: Union[str, PyExpr]) -> DeleteResult: ...
async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ...
async def add_computed_columns(
self, columns: list[tuple[str, str]]
) -> AddColumnsResult: ...
async def add_function_columns(
self, application_json: str, output_name: Optional[str]
) -> AddColumnsResult: ...
async def refresh_column(self, column: str) -> RefreshColumnResult: ...
async def refresh_column_async(self, column: str) -> Job: ...
async def refresh_materialized_view(
self, full: bool = False, source_version: Optional[int] = None
) -> RefreshMaterializedViewResult: ...
async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ...
async def alter_columns(
self, columns: list[dict[str, Any]]
@@ -355,6 +381,10 @@ class Table:
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
async def unset_lsm_write_spec(self) -> None: ...
async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ...
async def checkpoint_lsm(self) -> None: ...
async def flush_lsm(self) -> None: ...
async def compact_lsm(self) -> None: ...
async def get_lsm_stats(self, include_generation_rows: bool) -> Optional[dict]: ...
async def close_lsm_writers(self) -> None: ...
@property
def tags(self) -> Tags: ...
@@ -396,7 +426,7 @@ class Branches:
async def checkout(self, name: str, version: Optional[int] = None) -> Table: ...
async def delete(self, name: str) -> None: ...
async def diff(self, from_branch: str) -> Dict[str, Any]: ...
async def merge(
async def cherry_pick(
self, from_branch: str, dry_run: bool = False
) -> Dict[str, Any]: ...
@@ -649,9 +679,10 @@ class LsmWriteSpec:
def identity(column: str) -> "LsmWriteSpec": ...
@staticmethod
def unsharded() -> "LsmWriteSpec": ...
def with_maintained_indexes(self, indexes: List[str]) -> "LsmWriteSpec":
"""Return a copy of this spec asking the MemWAL to keep the named
indexes up to date as rows are appended."""
def with_maintained_indexes(self, indexes: Optional[List[str]]) -> "LsmWriteSpec":
"""Set which indexes the MemWAL keeps up to date. None resolves every
index on the table at install, failing if one cannot be maintained;
a list is verbatim, empty means none."""
...
def with_writer_config_defaults(self, defaults: Dict[str, str]) -> "LsmWriteSpec":
"""Return a copy of this spec recording the given default
@@ -666,13 +697,25 @@ class LsmWriteSpec:
@property
def num_buckets(self) -> Optional[int]: ...
@property
def maintained_indexes(self) -> List[str]: ...
def maintained_indexes(self) -> Optional[List[str]]:
"""Indexes the MemWAL keeps up to date, or None for every supported one."""
...
@property
def writer_config_defaults(self) -> Dict[str, str]: ...
class AddColumnsResult:
version: int
class RefreshColumnResult:
rows_filled: int
version: int
class RefreshMaterializedViewResult:
mode: str
rows_written: int
source_version: int
version: int
class AlterColumnsResult:
version: int
python/python/lancedb/db.py
+275 -4
View File
@@ -45,7 +45,14 @@ from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
from . import __version__
from ._lancedb import connect as lancedb_connect # type: ignore
from .job import AsyncJob, Job
from .functions import FunctionVersion, UdfDefinition
from .job import AsyncJob, Job, _typed_job
from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
SelectArg,
normalize_select,
)
from .table import (
AsyncTable,
LanceTable,
@@ -509,6 +516,70 @@ class DBConnection(EnforceOverrides):
"""
raise NotImplementedError
def create_materialized_view(
self,
name: str,
source: str,
*,
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
) -> MaterializedView:
"""Define a materialized view named ``name`` over the table ``source``.
The view is created empty, with the query recorded in its schema
metadata; ``view.refresh()`` computes the rows. The view is a normal
table: it can be queried, indexed and searched, and it appears in
``table_names``. Local databases only.
The source table must have stable row ids (create it with the
``new_table_enable_stable_row_ids`` storage option): they keep the
view's provenance valid across source compactions, and cannot be
enabled after a table exists.
Parameters
----------
name: str
The name of the view.
source: str
The name of the source table, in this database.
select: list or dict, optional
The view's columns: column names, ``(alias, SQL expression)``
pairs, or a dict of the same. Omitting it selects every source
column, expanded against the source schema at creation time.
where: str, optional
SQL predicate; only matching source rows appear in the view.
limit: int, optional
Cap the view at this many rows, in materialization order.
Returns
-------
MaterializedView
"""
raise NotImplementedError(
"materialized views are not supported on this connection type"
)
def open_materialized_view(self, name: str) -> MaterializedView:
"""Open the materialized view named ``name``.
Raises ``ValueError`` if the table exists but is not a materialized
view.
"""
raise NotImplementedError(
"materialized views are not supported on this connection type"
)
def list_materialized_views(self) -> List[str]:
"""The names of the materialized views in this database.
Found by reading every table's schema, so this costs an open per
table.
"""
raise NotImplementedError(
"materialized views are not supported on this connection type"
)
def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
"""Drop a table from the database.
@@ -524,6 +595,12 @@ class DBConnection(EnforceOverrides):
namespace_path = []
raise NotImplementedError
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
def rename_table(
self,
cur_name: str,
@@ -610,6 +687,31 @@ class DBConnection(EnforceOverrides):
"""
raise NotImplementedError("serialize is not supported for this connection type")
def create_function(self, definition: UdfDefinition) -> FunctionVersion:
"""Register a scalar Python UDF and wait for its immutable version.
This is the blocking counterpart of :meth:`create_function_async`.
Local connections raise ``NotImplementedError``.
"""
return self.create_function_async(definition).wait()
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
"""Register a scalar Python UDF through the remote Function catalog.
Submission returns a typed job. The immutable Function version becomes
available only when :meth:`Job.wait` succeeds. Local connections raise
``NotImplementedError``.
"""
raise NotImplementedError(
"Function catalog operations are not supported for this connection type"
)
def get_function(self, name: str, *, version: str) -> FunctionVersion:
"""Open one exact immutable Function version from the remote catalog."""
raise NotImplementedError(
"Function catalog operations are not supported for this connection type"
)
def job(self, job_id: str) -> Job:
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
@@ -707,6 +809,9 @@ class LanceDBConnection(DBConnection):
self._namespace_client_properties = namespace_client_properties
if _inner is not None:
self._conn = _inner
# Native-derived wrappers resolve this in their async reconstruction
# path so construction never synchronously re-enters LOOP.
self._read_consistency_interval = read_consistency_interval
self._cached_namespace_client = None
return
@@ -756,11 +861,14 @@ class LanceDBConnection(DBConnection):
# storage_options. Also, this class really shouldn't be holding any state
# beyond _conn.
self._conn = AsyncConnection(LOOP.run(do_connect()))
# Keep property access synchronous so debugger introspection cannot wait on
# the background loop while that thread is suspended at a breakpoint.
self._read_consistency_interval = read_consistency_interval
self._cached_namespace_client: Optional[LanceNamespace] = None
@property
def read_consistency_interval(self) -> Optional[timedelta]:
return LOOP.run(self._conn.get_read_consistency_interval())
return self._read_consistency_interval
@property
def session(self) -> Optional[Session]:
@@ -771,8 +879,16 @@ class LanceDBConnection(DBConnection):
return self._conn.uri
@classmethod
def from_inner(cls, inner: LanceDbConnection):
return cls(None, _inner=inner)
def from_inner(
cls,
inner: LanceDbConnection,
read_consistency_interval: Optional[timedelta],
):
return cls(
None,
read_consistency_interval=read_consistency_interval,
_inner=inner,
)
def __repr__(self) -> str:
return f"{self.__class__.__name__}(uri={self._conn.uri!r})"
@@ -1090,6 +1206,58 @@ class LanceDBConnection(DBConnection):
tbl.checkout(version)
return tbl
@override
def create_materialized_view(
self,
name: str,
source: str,
*,
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
) -> MaterializedView:
"""Define a materialized view named ``name`` over the table ``source``.
See
[DBConnection.create_materialized_view][lancedb.DBConnection.create_materialized_view].
Examples
--------
>>> import lancedb
>>> db = lancedb.connect(
... "./.lancedb",
... storage_options={"new_table_enable_stable_row_ids": "true"},
... )
>>> data = [{"name": "ada", "age": 36}, {"name": "kid", "age": 7}]
>>> table = db.create_table("people", data)
>>> view = db.create_materialized_view(
... "adults",
... "people",
... select=["name", ("shout", "upper(name)")],
... where="age >= 18",
... )
>>> result = view.refresh()
>>> result.rows_written
1
"""
LOOP.run(
self._conn.create_materialized_view(
name, source, select=select, where=where, limit=limit
)
)
return MaterializedView(self.open_table(name))
@override
def open_materialized_view(self, name: str) -> MaterializedView:
"""Open the materialized view named ``name``."""
view = MaterializedView(self.open_table(name))
view.definition
return view
@override
def list_materialized_views(self) -> List[str]:
"""The names of the materialized views in this database."""
return LOOP.run(self._conn.list_materialized_views())
def clone_table(
self,
target_table_name: str,
@@ -1172,6 +1340,20 @@ class LanceDBConnection(DBConnection):
)
)
@override
def drop_table_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Job:
"""Start dropping a table and return its cleanup job.
The table may become unavailable before its data files are removed.
Call :meth:`Job.wait` to wait for cleanup to finish.
"""
if namespace_path is None:
namespace_path = []
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):
if namespace_path is None:
@@ -1222,6 +1404,15 @@ class LanceDBConnection(DBConnection):
"""
return Job(self._conn.job(job_id))
@override
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
job = LOOP.run(self._conn.create_function_async(definition))
return Job(job)
@override
def get_function(self, name: str, *, version: str) -> FunctionVersion:
return LOOP.run(self._conn.get_function(name, version=version))
@override
def list_jobs(self) -> List[JobInfo]:
"""List server-side jobs across the database's tables."""
@@ -1837,6 +2028,50 @@ class AsyncConnection(object):
await tbl.checkout(version)
return tbl
async def create_materialized_view(
self,
name: str,
source: str,
*,
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
) -> AsyncMaterializedView:
"""Define a materialized view named ``name`` over the table ``source``.
See
[DBConnection.create_materialized_view][lancedb.DBConnection.create_materialized_view].
"""
inner = await self._inner.create_materialized_view(
name,
source,
projections=normalize_select(select),
filter=where,
limit=limit,
)
return AsyncMaterializedView(AsyncTable(inner))
async def open_materialized_view(self, name: str) -> AsyncMaterializedView:
"""Open the materialized view named ``name``.
Raises ``ValueError`` if the table exists but is not a materialized
view.
"""
if self.uri.startswith("db://"):
raise NotImplementedError(
"materialized views are supported only on local databases"
)
view = AsyncMaterializedView(await self.open_table(name))
await view.definition()
return view
async def list_materialized_views(self) -> List[str]:
"""The names of the materialized views in this database.
Found by reading every table's schema, so this costs an open per
table.
"""
return await self._inner.list_materialized_views()
async def clone_table(
self,
target_table_name: str,
@@ -1949,6 +2184,23 @@ class AsyncConnection(object):
if f"Table '{name}' was not found" not in str(e):
raise e
async def drop_table_async(
self,
name: str,
*,
namespace_path: Optional[List[str]] = None,
) -> AsyncJob:
"""Start dropping a table and return its cleanup job.
The table may become unavailable before its data files are removed.
Await :meth:`AsyncJob.wait` to wait for cleanup to finish.
"""
if namespace_path is None:
namespace_path = []
return AsyncJob(
await self._inner.drop_table_async(name, namespace_path=namespace_path)
)
async def drop_all_tables(self, namespace_path: Optional[List[str]] = None):
"""Drop all tables from the database.
@@ -1972,6 +2224,25 @@ class AsyncConnection(object):
"""
return AsyncJob(self._inner.job(job_id))
async def create_function_async(
self, definition: UdfDefinition
) -> AsyncJob[FunctionVersion]:
"""Register a scalar Python UDF through the remote Function catalog.
The returned typed job resolves to the immutable Function version.
Local connections raise ``NotImplementedError``.
"""
if not isinstance(definition, UdfDefinition):
raise TypeError("create_function_async requires a @udf definition")
inner = await self._inner.create_function_async(
definition.registration_request.to_canonical_json()
)
return _typed_job(inner, FunctionVersion.from_json)
async def get_function(self, name: str, *, version: str) -> FunctionVersion:
"""Open one exact immutable Function version from the remote catalog."""
return FunctionVersion.from_json(await self._inner.get_function(name, version))
async def list_jobs(self) -> List[JobInfo]:
"""List server-side jobs across the database's tables."""
return await self._inner.list_jobs()
python/python/lancedb/embeddings/base.py
-1
View File
@@ -26,7 +26,6 @@ class EmbeddingFunction(BaseModel, ABC):
3. ndims() which returns the number of dimensions of the vector column
"""
__slots__ = ("__weakref__",) # pydantic 1.x compatibility
max_retries: int = (
7 # Setting 0 disables retires. Maybe this should not be enabled by default,
)
python/python/lancedb/embeddings/bedrock.py
+2 -9
View File
@@ -7,8 +7,7 @@ from functools import cached_property
from typing import List, Union
import numpy as np
from lancedb.pydantic import PYDANTIC_VERSION
from pydantic import ConfigDict
from ..util import attempt_import_or_raise
from .base import TextEmbeddingFunction
@@ -67,13 +66,7 @@ class BedRockText(TextEmbeddingFunction):
source_input_type: str = "search_document"
query_input_type: str = "search_query"
if PYDANTIC_VERSION.major < 2: # Pydantic 1.x compat
class Config:
keep_untouched = (cached_property,)
else:
model_config = dict()
model_config["ignored_types"] = (cached_property,)
model_config = ConfigDict(ignored_types=(cached_property,))
def ndims(self):
# return len(self._generate_embedding("test"))
python/python/lancedb/embeddings/gemini_text.py
+2 -9
View File
@@ -7,8 +7,7 @@ from functools import cached_property
from typing import List, Optional, Union
import numpy as np
from lancedb.pydantic import PYDANTIC_VERSION
from pydantic import ConfigDict
from ..util import attempt_import_or_raise
from .base import TextEmbeddingFunction
@@ -87,13 +86,7 @@ class GeminiText(TextEmbeddingFunction):
query_task_type: str = "retrieval_query"
source_task_type: str = "retrieval_document"
if PYDANTIC_VERSION.major < 2: # Pydantic 1.x compat
class Config:
keep_untouched = (cached_property,)
else:
model_config = dict()
model_config["ignored_types"] = (cached_property,)
model_config = ConfigDict(ignored_types=(cached_property,))
def ndims(self):
if self.dim:
python/python/lancedb/embeddings/imagebind.py
+2 -9
View File
@@ -7,14 +7,13 @@ from typing import List, Union
import numpy as np
import pyarrow as pa
from pydantic import ConfigDict
from ..util import attempt_import_or_raise
from .base import EmbeddingFunction
from .registry import register
from .utils import AUDIO, IMAGES, TEXT
from lancedb.pydantic import PYDANTIC_VERSION
@register("imagebind")
class ImageBindEmbeddings(EmbeddingFunction):
@@ -31,13 +30,7 @@ class ImageBindEmbeddings(EmbeddingFunction):
device: str = "cpu"
normalize: bool = False
if PYDANTIC_VERSION.major < 2: # Pydantic 1.x compat
class Config:
keep_untouched = (cached_property,)
else:
model_config = dict()
model_config["ignored_types"] = (cached_property,)
model_config = ConfigDict(ignored_types=(cached_property,))
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
python/python/lancedb/embeddings/instructor.py
+1 -2
View File
@@ -101,8 +101,7 @@ class InstructorEmbeddingFunction(TextEmbeddingFunction):
@weak_lru(maxsize=1)
def ndims(self):
model = self.get_model()
return model.encode("foo").shape[0]
return len(self.generate_embeddings([[self.source_instruction, "foo"]])[0])
def compute_query_embeddings(self, query: str, *args, **kwargs) -> List[np.array]:
return self.generate_embeddings([[self.query_instruction, query]])
python/python/lancedb/embeddings/jinaai.py
+4 -3
View File
@@ -87,12 +87,13 @@ class JinaEmbeddings(EmbeddingFunction):
if isinstance(image, bytes):
image_dict = {"image": base64.b64encode(image).decode("utf-8")}
elif isinstance(image, (str, Path)):
parsed = urlparse.urlparse(image)
# TODO handle drive letter on windows.
parsed = urlparse(str(image))
PIL_Image = attempt_import_or_raise("PIL.Image", "pillow")
if parsed.scheme == "file":
pil_image = PIL_Image.open(parsed.path)
elif parsed.scheme == "":
elif parsed.scheme == "" or (os.name == "nt" and len(parsed.scheme) == 1):
# A Windows drive letter parses as a one-character scheme
# ("C:\\img.png" -> scheme="c"), so treat it as a local path.
pil_image = PIL_Image.open(image if os.name == "nt" else parsed.path)
elif parsed.scheme.startswith("http"):
pil_image = PIL_Image.open(io.BytesIO(url_retrieve(image)))
python/python/lancedb/embeddings/transformers.py
+2 -9
View File
@@ -7,8 +7,7 @@ from typing import List, Any
import numpy as np
from pydantic import PrivateAttr
from lancedb.pydantic import PYDANTIC_VERSION
from pydantic import ConfigDict, PrivateAttr
from ..util import attempt_import_or_raise
from .base import EmbeddingFunction
@@ -59,13 +58,7 @@ class TransformersEmbeddingFunction(EmbeddingFunction):
)
self._model.to(self.device)
if PYDANTIC_VERSION.major < 2: # Pydantic 1.x compat
class Config:
keep_untouched = (cached_property,)
else:
model_config = dict()
model_config["ignored_types"] = (cached_property,)
model_config = ConfigDict(ignored_types=(cached_property,))
def ndims(self):
self._ndims = self._model.config.hidden_size
python/python/lancedb/expr.py
+3 -2
View File
@@ -85,8 +85,9 @@ class Expr:
# for dict keys / set membership.
__hash__ = None # type: ignore[assignment]
def __init__(self, inner: PyExpr) -> None:
def __init__(self, inner: PyExpr, *, column_path: str | None = None) -> None:
self._inner = inner
self._column_path = column_path
# ── comparisons ──────────────────────────────────────────────────────────
@@ -273,7 +274,7 @@ def col(name: str) -> Expr:
>>> col("age") > lit(18)
Expr((age > 18))
"""
return Expr(expr_col(name))
return Expr(expr_col(name), column_path=name)
def lit(value: Union[bool, int, float, str, bytes, date, datetime, Decimal]) -> Expr:
python/python/lancedb/functions.py
+1081
View File
File diff suppressed because it is too large Load Diff
python/python/lancedb/index.py
+11
View File
@@ -163,6 +163,15 @@ class FTS:
The number of documents per compressed posting block. Supported values
are 128 and 256. A value of 256 uses the experimental FTS V3 format
and may introduce breaking changes.
memory_limit : int, optional
The total memory limit in MiB for the local FTS build stage. The limit
is divided evenly among indexing workers. This build-only setting is
not persisted with the index and does not apply to remote tables.
num_workers : int, optional
The number of workers for a local FTS build. By default Lance uses
roughly half of the available CPU cores. The effective value is
limited by the available compute capacity. This build-only setting is
not persisted with the index and does not apply to remote tables.
Notes
-----
@@ -185,6 +194,8 @@ class FTS:
prefix_only: bool = False
block_size: int = 128
custom_stop_words: Optional[List[str]] = None
memory_limit: Optional[int] = None
num_workers: Optional[int] = None
@dataclass
python/python/lancedb/job.py
+39 -14
View File
@@ -5,21 +5,30 @@
import asyncio
from datetime import timedelta
from typing import Optional
from typing import Any, Callable, Generic, Optional, TypeVar, cast
from lancedb.background_loop import LOOP
from . import _lancedb
T = TypeVar("T")
class AsyncJob:
class AsyncJob(Generic[T]):
"""A handle to an operation that may still be running.
The operation may already be complete when the handle is created.
The operation may already be complete when the handle is created. ``T``
is the endpoint's terminal result type; unit-result jobs resolve to
``None``.
"""
def __init__(self, inner: Optional["_lancedb.Job"]):
def __init__(
self,
inner: Optional[Any],
result_decoder: Optional[Callable[[Any], T]] = None,
):
self._inner = inner
self._result_decoder = result_decoder
@property
def id(self) -> Optional[str]:
@@ -44,18 +53,24 @@ class AsyncJob:
return "finished"
return await self._inner.status()
async def wait(self, timeout: Optional[timedelta] = None):
async def wait(self, timeout: Optional[timedelta] = None) -> T:
"""Wait until the operation reaches a terminal state.
Returns the endpoint's typed result, or ``None`` for a unit-result
job.
Raises `JobFailedError` if the operation failed, `JobCancelledError`
if it was cancelled, and `TimeoutError` if `timeout` elapses first.
"""
if self._inner is None:
return
return cast(T, None)
if timeout is None:
await self._inner.wait()
result = await self._inner.wait()
else:
await asyncio.wait_for(self._inner.wait(), timeout.total_seconds())
result = await asyncio.wait_for(self._inner.wait(), timeout.total_seconds())
if self._result_decoder is not None:
return self._result_decoder(result)
return cast(T, result)
async def cancel(self):
"""Request cancellation. Cancelling a finished operation is a no-op."""
@@ -64,10 +79,10 @@ class AsyncJob:
await self._inner.cancel()
class Job:
"""Synchronous counterpart of `AsyncJob`."""
class Job(Generic[T]):
"""Synchronous counterpart of `AsyncJob` with the same result type."""
def __init__(self, inner: Optional[AsyncJob]):
def __init__(self, inner: Optional[AsyncJob[T]]):
self._inner = inner
@property
@@ -88,18 +103,28 @@ class Job:
return "finished"
return LOOP.run(self._inner.status())
def wait(self, timeout: Optional[timedelta] = None):
def wait(self, timeout: Optional[timedelta] = None) -> T:
"""Block until the operation reaches a terminal state.
Returns the endpoint's typed result, or ``None`` for a unit-result
job.
Raises `JobFailedError` if the operation failed, `JobCancelledError`
if it was cancelled, and `TimeoutError` if `timeout` elapses first.
"""
if self._inner is None:
return
LOOP.run(self._inner.wait(timeout))
return cast(T, None)
return LOOP.run(self._inner.wait(timeout))
def cancel(self):
"""Request cancellation. Cancelling a finished operation is a no-op."""
if self._inner is None:
return
LOOP.run(self._inner.cancel())
def _typed_job(
inner: "_lancedb.Job", result_decoder: Callable[[str], T]
) -> AsyncJob[T]:
"""Bind an internal JSON-producing job to its public result model."""
return AsyncJob(inner, result_decoder)
python/python/lancedb/materialized_view.py
+178
View File
@@ -0,0 +1,178 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Materialized views: tables defined by a query over a source table and
maintained by refresh. See ``DBConnection.create_materialized_view``."""
from __future__ import annotations
import json
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
from .background_loop import LOOP
if TYPE_CHECKING:
import pyarrow as pa
from ._lancedb import RefreshMaterializedViewResult
from .table import AsyncTable, LanceTable
DEFINITION_META_KEY = b"mv.definition"
SelectArg = Union[
str,
Sequence[Union[str, Tuple[str, str]]],
Dict[str, str],
None,
]
@dataclass
class MaterializedViewDefinition:
"""The query that defines a materialized view."""
source_table: str
"""Name of the source table, in the same database as the view."""
projections: List[Tuple[str, str]]
"""``(output column, SQL expression)`` pairs, in view schema order."""
filter: Optional[str] = None
"""SQL predicate selecting the source rows the view holds."""
limit: Optional[int] = None
"""Cap on the number of rows the view holds."""
inputs: List[str] = field(default_factory=list)
"""Source columns the projections and filter read."""
def _definition_from_schema(
schema: "pa.Schema", name: str
) -> MaterializedViewDefinition:
metadata = schema.metadata or {}
raw = metadata.get(DEFINITION_META_KEY)
if raw is None:
raise ValueError(f"Table '{name}' is not a materialized view")
value = json.loads(raw)
kind = value.get("kind")
if kind != "select":
raise NotImplementedError(
f"materialized view '{name}' is defined by '{kind}', which this "
"version of lancedb cannot refresh"
)
return MaterializedViewDefinition(
source_table=value["source_table"],
projections=[
(p["output"], p["expression"]) for p in value.get("projections", [])
],
filter=value.get("filter"),
limit=value.get("limit"),
inputs=value.get("inputs", []),
)
def _quote_identifier(name: str) -> str:
"""Quote a column name as a Lance SQL identifier (backticks)."""
escaped = name.replace("`", "``")
return f"`{escaped}`"
def normalize_select(select: SelectArg) -> Optional[List[Tuple[str, str]]]:
"""``select`` items may be a column name, an ``(alias, expression)`` pair,
or a dict of the same. A bare name projects itself and is quoted, so any
valid column name works; dict and pair entries are kept verbatim because
their right side is an expression.
A lone string is one column, not a sequence of its characters."""
if select is None:
return None
if isinstance(select, str):
select = [select]
if isinstance(select, dict):
return list(select.items())
normalized = []
for item in select:
if isinstance(item, str):
normalized.append((item, _quote_identifier(item)))
else:
alias, expression = item
normalized.append((alias, expression))
return normalized
class AsyncMaterializedView:
"""A handle on a materialized view: its table plus its definition.
Obtained from ``AsyncConnection.create_materialized_view`` or
``AsyncConnection.open_materialized_view``.
"""
def __init__(self, table: "AsyncTable"):
self._table = table
def __repr__(self) -> str:
return f"AsyncMaterializedView(name={self.name!r})"
@property
def name(self) -> str:
return self._table.name
@property
def table(self) -> "AsyncTable":
"""The view, as the table it is. Queries, indexes and search all
apply; writes are not blocked, but a rebuild replaces them."""
return self._table
async def definition(self) -> MaterializedViewDefinition:
"""The query that defines the view, read from its stored schema."""
return _definition_from_schema(await self._table.schema(), self.name)
async def refresh(
self, *, full: bool = False, source_version: Optional[int] = None
) -> "RefreshMaterializedViewResult":
"""Recompute the view from its source.
The refresh is incremental when the source's changes can be
reconciled into the view -- rows added, changed or removed since the
last one -- and otherwise rebuilds. ``full=True`` forces a rebuild;
``source_version`` refreshes to that source version instead of the
latest.
Concurrent refreshes of one view do not duplicate its rows. Two that
plan the same source rows conflict on commit, and the loser raises
rather than writing them a second time.
"""
return await self._table._inner.refresh_materialized_view(
full=full, source_version=source_version
)
class MaterializedView:
"""Synchronous variant of
[AsyncMaterializedView][lancedb.materialized_view.AsyncMaterializedView]."""
def __init__(self, table: "LanceTable"):
self._table = table
self._async = AsyncMaterializedView(table._table)
def __repr__(self) -> str:
return f"MaterializedView(name={self.name!r})"
@property
def name(self) -> str:
return self._table.name
@property
def table(self) -> "LanceTable":
"""The view, as the table it is."""
return self._table
@property
def definition(self) -> MaterializedViewDefinition:
"""The query that defines the view, read from its stored schema."""
return _definition_from_schema(self._table.schema, self.name)
def refresh(
self, *, full: bool = False, source_version: Optional[int] = None
) -> "RefreshMaterializedViewResult":
"""Recompute the view from its source. See
[AsyncMaterializedView.refresh][lancedb.materialized_view.AsyncMaterializedView.refresh]."""
return LOOP.run(self._async.refresh(full=full, source_version=source_version))
python/python/lancedb/namespace.py
+89
View File
@@ -49,6 +49,7 @@ from lancedb._lancedb import (
)
from lancedb.background_loop import LOOP
from lancedb.db import AsyncConnection, DBConnection
from lancedb.job import AsyncJob, Job
from lance_namespace import (
LanceNamespace,
connect as namespace_connect,
@@ -60,6 +61,11 @@ from lance_namespace import (
NamespaceExistsRequest,
TableExistsRequest,
)
from lancedb.materialized_view import (
AsyncMaterializedView,
MaterializedView,
SelectArg,
)
from lancedb.table import AsyncTable, LanceTable, Table
from lancedb.util import validate_table_name
from lancedb.common import DATA
@@ -618,12 +624,60 @@ class LanceNamespaceDBConnection(DBConnection):
tbl.checkout(version)
return tbl
@override
def create_materialized_view(
self,
name: str,
source: str,
*,
select: "SelectArg" = None,
where: Optional[str] = None,
limit: Optional[int] = None,
) -> "MaterializedView":
"""Define a materialized view over a table in the root namespace.
See
[DBConnection.create_materialized_view][lancedb.DBConnection.create_materialized_view].
"""
return MaterializedView(
self.open_table(
LOOP.run(
self._inner.create_materialized_view(
name, source, select=select, where=where, limit=limit
)
).name
)
)
@override
def open_materialized_view(self, name: str) -> "MaterializedView":
"""Open the materialized view named ``name``."""
view = MaterializedView(self.open_table(name))
view.definition
return view
@override
def list_materialized_views(self) -> List[str]:
"""The names of the materialized views in the root namespace."""
return LOOP.run(self._inner.list_materialized_views())
@override
def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
if namespace_path is None:
namespace_path = []
LOOP.run(self._inner.drop_table(name, namespace_path=namespace_path))
@override
def drop_table_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Job:
"""Start dropping a table and return its cleanup job."""
if namespace_path is None:
namespace_path = []
job = LOOP.run(
self._inner.drop_table_async(name, namespace_path=namespace_path)
)
return Job(job if isinstance(job, AsyncJob) else AsyncJob(job))
@override
def rename_table(
self,
@@ -1128,12 +1182,47 @@ class AsyncLanceNamespaceDBConnection:
route_pushdown_to_rust=self._route_pushdown_to_rust,
)
async def create_materialized_view(
self,
name: str,
source: str,
*,
select: "SelectArg" = None,
where: Optional[str] = None,
limit: Optional[int] = None,
) -> "AsyncMaterializedView":
"""Define a materialized view over a table in the root namespace."""
view = await self._inner.create_materialized_view(
name, source, select=select, where=where, limit=limit
)
# Reopen through the namespace so the view's table carries the
# namespace client and pushdown configuration a bare inner table lacks.
return AsyncMaterializedView(await self.open_table(view.name))
async def open_materialized_view(self, name: str) -> "AsyncMaterializedView":
"""Open the materialized view named ``name``."""
view = AsyncMaterializedView(await self.open_table(name))
await view.definition()
return view
async def list_materialized_views(self) -> List[str]:
"""The names of the materialized views in the root namespace."""
return await self._inner.list_materialized_views()
async def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
"""Drop a table from the namespace."""
if namespace_path is None:
namespace_path = []
await self._inner.drop_table(name, namespace_path=namespace_path)
async def drop_table_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> AsyncJob:
"""Start dropping a table and return its cleanup job."""
if namespace_path is None:
namespace_path = []
return await self._inner.drop_table_async(name, namespace_path=namespace_path)
async def rename_table(
self,
cur_name: str,
python/python/lancedb/permutation.py
+28 -18
View File
@@ -41,21 +41,15 @@ class PermutationBuilder:
The permutation is stored in memory and will be lost when the program exits.
"""
def __init__(self, table: LanceTable):
def __init__(self, table: Table):
"""
Creates a new permutation builder for the given table.
By default, the permutation builder will create a single split that contains all
rows in the same order as the base table.
Tables with an LSM write spec are rejected: unflushed rows have no row id.
"""
if not hasattr(table, "_inner"):
raise TypeError(
f"PermutationBuilder requires a local LanceTable, "
f"got {type(table).__name__}. "
"The permutation API is not supported on remote tables. "
"Remote tables connect to LanceDB Cloud or Enterprise and do not have "
"direct access to the underlying Lance dataset needed for permutations."
)
self._async = async_permutation_builder(table)
def split_random(
@@ -226,12 +220,12 @@ class PermutationBuilder:
async def do_execute():
inner_tbl = await self._async.execute()
return LanceTable.from_inner(inner_tbl)
return await LanceTable.from_inner(inner_tbl)
return LOOP.run(do_execute())
def permutation_builder(table: LanceTable) -> PermutationBuilder:
def permutation_builder(table: Table) -> PermutationBuilder:
return PermutationBuilder(table)
@@ -248,7 +242,7 @@ class Permutations:
Attributes
----------
base_table: LanceTable
base_table: Table
The base table that the permutations are based on.
permutation_table: LanceTable
The permutation table that defines the splits.
@@ -282,7 +276,7 @@ class Permutations:
{'train': 0, 'test': 1}
"""
def __init__(self, base_table: LanceTable, permutation_table: LanceTable):
def __init__(self, base_table: Table, permutation_table: LanceTable):
self.base_table = base_table
self.permutation_table = permutation_table
@@ -397,6 +391,15 @@ def _table_to_pickle_state(table: Table) -> dict[str, Any]:
}
def _drop_base_version(permutation_data: pa.Table) -> pa.Table:
"""Strip the recorded base version so the reader leaves the base table unpinned."""
metadata = dict(permutation_data.schema.metadata or {})
if metadata.pop(b"base_version", None) is None:
return permutation_data
metadata.pop(b"base_branch", None)
return permutation_data.replace_schema_metadata(metadata)
def _table_from_pickle_state(state: dict[str, Any]) -> Table:
from . import connect
@@ -685,11 +688,15 @@ class Permutation:
from . import connect
connection_factory = state["connection_factory"]
rebuilt_base = False
if connection_factory is not None:
base_table = connection_factory(state["base_table_name"])
elif "base_table_state" in state:
base_table = _table_from_pickle_state(state["base_table_state"])
base_state = state["base_table_state"]
rebuilt_base = base_state["kind"] == "memory"
base_table = _table_from_pickle_state(base_state)
elif "base_table_data" in state:
rebuilt_base = True
# In-memory base table inlined into the pickle; rebuild the same
# way we rebuild the in-memory permutation table.
mem_db = connect("memory://")
@@ -707,11 +714,14 @@ class Permutation:
)
permutation_table: Optional[Table] = None
if state["permutation_data"] is not None:
permutation_data = state["permutation_data"]
if permutation_data is not None:
if rebuilt_base:
# The base table was materialized from Arrow, so it is a fresh
# single-version dataset and the recorded pin cannot resolve on it.
permutation_data = _drop_base_version(permutation_data)
mem_db = connect("memory://")
permutation_table = mem_db.create_table(
"permutation", state["permutation_data"]
)
permutation_table = mem_db.create_table("permutation", permutation_data)
self.base_table = base_table
self.permutation_table = permutation_table
python/python/lancedb/py.typed
+1
View File
@@ -0,0 +1 @@
python/python/lancedb/pydantic.py
+19 -91
View File
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Pydantic (v1 / v2) adapter for LanceDB"""
"""Pydantic adapter for LanceDB."""
from __future__ import annotations
@@ -14,9 +14,6 @@ from enum import Enum
from typing import (
TYPE_CHECKING,
Any,
Callable,
Dict,
Generator,
List,
Type,
Union,
@@ -24,17 +21,9 @@ from typing import (
GenericAlias,
)
import numpy as np
import pyarrow as pa
import pydantic
from packaging.version import Version
PYDANTIC_VERSION = Version(pydantic.__version__)
try:
from pydantic_core import CoreSchema, core_schema
except ImportError:
if PYDANTIC_VERSION.major >= 2:
raise
from pydantic_core import CoreSchema, core_schema
if TYPE_CHECKING:
from pydantic.fields import FieldInfo
@@ -131,28 +120,18 @@ def Vector(
),
)
@classmethod
def __get_validators__(cls) -> Generator[Callable, None, None]:
yield cls.validate
# For pydantic v1
@classmethod
def validate(cls, v):
if not isinstance(v, (list, range, np.ndarray)) or len(v) != dim:
raise TypeError("A list of numbers or numpy.ndarray is needed")
return cls(v)
if PYDANTIC_VERSION.major < 2:
@classmethod
def __modify_schema__(cls, field_schema: Dict[str, Any]):
field_schema["items"] = {"type": "number"}
field_schema["maxItems"] = dim
field_schema["minItems"] = dim
return FixedSizeList
def _raise_bare_vector_error(*_args):
raise TypeError("Vector must be parameterized with a dimension, e.g. Vector(128).")
# Pydantic otherwise inspects the bare factory as a field type and produces
# misleading errors about its internal annotations.
setattr(Vector, "__get_pydantic_core_schema__", _raise_bare_vector_error)
def MultiVector(
dim: int, value_type: pa.DataType = pa.float32(), nullable: bool = True
) -> Type:
@@ -223,31 +202,6 @@ def MultiVector(
),
)
@classmethod
def __get_validators__(cls) -> Generator[Callable, None, None]:
yield cls.validate
# For pydantic v1
@classmethod
def validate(cls, v):
if not isinstance(v, (list, range)):
raise TypeError("A list of vectors is needed")
for vec in v:
if not isinstance(vec, (list, range, np.ndarray)) or len(vec) != dim:
raise TypeError(f"Each vector must be a list of {dim} numbers")
return cls(v)
if PYDANTIC_VERSION.major < 2:
@classmethod
def __modify_schema__(cls, field_schema: Dict[str, Any]):
field_schema["items"] = {
"type": "array",
"items": {"type": "number"},
"minItems": dim,
"maxItems": dim,
}
return MultiVectorList
@@ -293,20 +247,10 @@ def _py_type_to_arrow_type(py_type: Type[Any], field: FieldInfo) -> pa.DataType:
)
if PYDANTIC_VERSION.major < 2:
def _pydantic_model_to_fields(model: pydantic.BaseModel) -> List[pa.Field]:
return [
_pydantic_to_field(name, field) for name, field in model.__fields__.items()
]
else:
def _pydantic_model_to_fields(model: pydantic.BaseModel) -> List[pa.Field]:
return [
_pydantic_to_field(name, field)
for name, field in model.model_fields.items()
]
def _pydantic_model_to_fields(model: pydantic.BaseModel) -> List[pa.Field]:
return [
_pydantic_to_field(name, field) for name, field in model.model_fields.items()
]
def _pydantic_type_to_arrow_type(tp: Any, field: FieldInfo) -> pa.DataType:
@@ -499,8 +443,6 @@ class LanceModel(pydantic.BaseModel):
@classmethod
def safe_get_fields(cls):
if PYDANTIC_VERSION.major < 2:
return cls.__fields__
return cls.model_fields
@classmethod
@@ -538,23 +480,9 @@ def get_extras(field_info: FieldInfo, key: str) -> Any:
"""
Get the extra metadata from a Pydantic FieldInfo.
"""
if PYDANTIC_VERSION.major >= 2:
return (field_info.json_schema_extra or {}).get(key)
return (field_info.field_info.extra or {}).get("json_schema_extra", {}).get(key)
return (field_info.json_schema_extra or {}).get(key)
if PYDANTIC_VERSION.major < 2:
def model_to_dict(model: pydantic.BaseModel) -> Dict[str, Any]:
"""
Convert a Pydantic model to a dictionary.
"""
return model.dict()
else:
def model_to_dict(model: pydantic.BaseModel) -> Dict[str, Any]:
"""
Convert a Pydantic model to a dictionary.
"""
return model.model_dump()
def model_to_dict(model: pydantic.BaseModel) -> dict[str, Any]:
"""Convert a Pydantic model to a dictionary."""
return model.model_dump()
python/python/lancedb/query.py
+11 -17
View File
@@ -32,8 +32,6 @@ from typing_extensions import Annotated
from lancedb._lancedb import fts_query_to_json
from lancedb.background_loop import LOOP
from lancedb.pydantic import PYDANTIC_VERSION
from . import __version__
from .arrow import AsyncRecordBatchReader
from .dependencies import pandas as pd
@@ -827,12 +825,7 @@ class Query(pydantic.BaseModel):
# This tells pydantic to allow custom types (needed for the `vector` query since
# pa.Array wouln't be allowed otherwise)
if PYDANTIC_VERSION.major < 2: # Pydantic 1.x compat
class Config:
arbitrary_types_allowed = True
else:
model_config = {"arbitrary_types_allowed": True}
model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
class LanceQueryBuilder(ABC):
@@ -2235,6 +2228,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
reranker=self._reranker,
limit=self._limit,
with_row_ids=True,
offset=self._offset,
)
return self._finish_hybrid_results(results)
@@ -2256,6 +2250,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
reranker,
limit: int,
with_row_ids: bool,
offset: Optional[int] = None,
) -> pa.Table:
if norm == "rank":
vector_results = LanceHybridQueryBuilder._rank(vector_results, "_distance")
@@ -2332,7 +2327,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
score_i = results.column_names.index("_score")
results = results.set_column(score_i, "_score", original_scores)
results = results.slice(length=limit)
results = results.slice(offset=offset or 0, length=limit)
if not with_row_ids:
results = results.drop(["_rowid"])
@@ -2679,8 +2674,12 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
# Apply common configurations
if self._limit:
self._vector_query.limit(self._limit)
self._fts_query.limit(self._limit)
# The final offset/limit window is sliced out of the combined,
# reranked results, so each sub-query must fetch enough rows to
# cover the skipped prefix as well as the window itself.
sub_query_limit = self._limit + (self._offset or 0)
self._vector_query.limit(sub_query_limit)
self._fts_query.limit(sub_query_limit)
if self._columns:
self._vector_query.select(self._columns)
self._fts_query.select(self._columns)
@@ -3245,12 +3244,7 @@ class AsyncStandardQuery(AsyncQueryBase):
if ordering is None:
self._inner.order_by(None)
else:
self._inner.order_by(
[
o.model_dump() if hasattr(o, "model_dump") else o.dict()
for o in ordering
]
)
self._inner.order_by([o.model_dump() for o in ordering])
return self
def fast_search(self) -> Self:
python/python/lancedb/remote/db.py
+47 -1
View File
@@ -23,7 +23,9 @@ import pyarrow as pa
from ..common import DATA
from ..db import DBConnection, LOOP
from ..job import Job
from ..functions import FunctionVersion, UdfDefinition
from ..job import AsyncJob, Job
from ..materialized_view import MaterializedView, SelectArg
if TYPE_CHECKING:
from .._lancedb import JobDescription, JobInfo
@@ -647,6 +649,32 @@ class RemoteDBConnection(DBConnection):
namespace_path=namespace_path,
)
@override
def create_materialized_view(
self,
name: str,
source: str,
*,
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
) -> MaterializedView:
raise NotImplementedError(
"materialized views are supported only on local databases"
)
@override
def open_materialized_view(self, name: str) -> MaterializedView:
raise NotImplementedError(
"materialized views are supported only on local databases"
)
@override
def list_materialized_views(self) -> List[str]:
raise NotImplementedError(
"materialized views are supported only on local databases"
)
@override
def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
"""Drop a table from the database.
@@ -663,6 +691,16 @@ class RemoteDBConnection(DBConnection):
namespace_path = []
LOOP.run(self._conn.drop_table(name, namespace_path=namespace_path))
@override
def drop_table_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Job:
"""Start dropping a table and return its cleanup job."""
if namespace_path is None:
namespace_path = []
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 rename_table(
self,
@@ -703,6 +741,14 @@ class RemoteDBConnection(DBConnection):
"""
return Job(self._conn.job(job_id))
@override
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
return Job(LOOP.run(self._conn.create_function_async(definition)))
@override
def get_function(self, name: str, *, version: str) -> FunctionVersion:
return LOOP.run(self._conn.get_function(name, version=version))
@override
def list_jobs(self) -> List["JobInfo"]:
"""List server-side jobs across the database's tables."""
python/python/lancedb/remote/table.py
+44 -4
View File
@@ -49,6 +49,7 @@ from lancedb.index import (
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
import pyarrow as pa
@@ -609,6 +610,7 @@ class RemoteTable(Table):
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add more data to the [Table][lancedb.table.Table].
@@ -641,6 +643,8 @@ class RemoteTable(Table):
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Not supported on LanceDB Cloud. Setting this raises.
Returns
-------
@@ -657,6 +661,7 @@ class RemoteTable(Table):
fill_value=fill_value,
progress=progress,
write_parallelism=write_parallelism,
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
)
)
finally:
@@ -958,8 +963,21 @@ class RemoteTable(Table):
def count_rows(self, filter: Optional[str] = None) -> int:
return LOOP.run(self._table.count_rows(filter))
def add_columns(self, transforms: Dict[str, str]) -> AddColumnsResult:
return LOOP.run(self._table.add_columns(transforms))
def add_columns(
self,
transforms: Dict[str, str | FunctionApplication]
| FunctionApplication
| None = None,
*,
computed: Dict[str, str] | None = None,
) -> AddColumnsResult:
return LOOP.run(self._table.add_columns(transforms, computed=computed))
def refresh_column(self, column: str):
return LOOP.run(self._table.refresh_column(column))
def refresh_column_async(self, column: str) -> Job[RefreshColumnResult]:
return Job(LOOP.run(self._table.refresh_column_async(column)))
def alter_columns(
self, *alterations: Iterable[Dict[str, str]]
@@ -979,17 +997,39 @@ class RemoteTable(Table):
return LOOP.run(self._table.set_unenforced_primary_key(columns))
def set_lsm_write_spec(self, spec: "LsmWriteSpec") -> None:
"""Not supported on LanceDB Cloud."""
"""Install an LsmWriteSpec."""
return LOOP.run(self._table.set_lsm_write_spec(spec))
def unset_lsm_write_spec(self) -> None:
"""Not supported on LanceDB Cloud."""
"""Remove the LsmWriteSpec."""
return LOOP.run(self._table.unset_lsm_write_spec())
def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
"""Read the installed LsmWriteSpec, or ``None``."""
return LOOP.run(self._table.get_lsm_write_spec())
def checkpoint_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""
return LOOP.run(self._table.checkpoint_lsm())
def flush_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.flush_lsm`][lancedb.AsyncTable.flush_lsm]."""
return LOOP.run(self._table.flush_lsm())
def compact_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.compact_lsm`][lancedb.AsyncTable.compact_lsm]."""
return LOOP.run(self._table.compact_lsm())
def get_lsm_stats(self, *, include_generation_rows: bool = False) -> Optional[dict]:
"""Synchronous version of
[`AsyncTable.get_lsm_stats`][lancedb.AsyncTable.get_lsm_stats]."""
return LOOP.run(
self._table.get_lsm_stats(include_generation_rows=include_generation_rows)
)
def close_lsm_writers(self) -> None:
"""No-op on LanceDB Cloud (no local shard writers)."""
return LOOP.run(self._table.close_lsm_writers())
python/python/lancedb/streaming.py
+1298 -64
View File
File diff suppressed because it is too large Load Diff
python/python/lancedb/table.py
+447 -30
View File
@@ -40,7 +40,7 @@ from ._blob import (
from .types import BlobMode
from lancedb.arrow import peek_reader
from lancedb.background_loop import LOOP, embedding_executor
from lancedb.job import AsyncJob, Job
from lancedb.job import AsyncJob, Job, _typed_job
from .dependencies import (
_check_for_hugging_face,
_check_for_lance,
@@ -72,6 +72,10 @@ from .index import (
FTS,
)
from .expr import Expr
from .functions import (
FunctionApplication,
RefreshColumnResult as RefreshColumnJobResult,
)
from .merge import LanceMergeInsertBuilder
from .pydantic import LanceModel, model_to_dict
from .query import (
@@ -108,6 +112,11 @@ def _should_push_down_query_table(
return namespace_client is not None and "QueryTable" in pushdown_operations
def _polars_predicate_pushdown_barrier(frame: Any) -> Any:
"""Return a Polars frame unchanged while blocking predicate pushdown."""
return frame
_MODEL_BACKED_TOKENIZER_PREFIXES = ("jieba", "lindera")
_MODEL_BACKED_TOKENIZER_ERRORS = (
"unknown base tokenizer",
@@ -171,6 +180,7 @@ if TYPE_CHECKING:
CompactionStats,
Tag,
AddColumnsResult,
RefreshColumnResult,
AddResult,
AlterColumnsResult,
UpdateFieldMetadataResult,
@@ -427,6 +437,20 @@ def _cast_to_target_schema(
return pa.RecordBatchReader.from_batches(reordered_schema, gen())
def _field_extension_name(field: pa.Field) -> Optional[str]:
extension_name = getattr(field.type, "extension_name", None)
if extension_name is not None:
return extension_name
metadata = field.metadata or {}
extension_name = metadata.get(b"ARROW:extension:name") or metadata.get(
"ARROW:extension:name"
)
if isinstance(extension_name, bytes):
return extension_name.decode()
return extension_name
def _align_field_types(
fields: List[pa.Field],
target_fields: List[pa.Field],
@@ -439,6 +463,16 @@ def _align_field_types(
target_field = next((f for f in target_fields if f.name == field.name), None)
if target_field is None:
raise ValueError(f"Field '{field.name}' not found in target schema")
# Preserve arrow.json input until it reaches Lance. LanceDB exposes stored
# JSON columns as lance.json (JSONB-backed LargeBinary), but casting the
# input to that storage type here merely relabels the raw JSON bytes as
# JSONB. Lance must see arrow.json so it can perform the JSONB encoding.
if (
_field_extension_name(field) == "arrow.json"
and _field_extension_name(target_field) == "lance.json"
):
new_fields.append(field)
continue
if pa.types.is_struct(target_field.type):
if pa.types.is_struct(field.type):
new_type = pa.struct(
@@ -864,12 +898,18 @@ class Table(ABC):
"""
raise NotImplementedError
def to_polars(self, **kwargs) -> "pl.DataFrame":
"""Return the table as a polars.DataFrame.
def to_polars(self, **kwargs) -> "pl.LazyFrame":
"""Return the table as a Polars LazyFrame.
Note
----
The Polars streaming engine is not supported because it does not currently
implement Python PyArrow dataset scans. Use the default engine when collecting
this LazyFrame.
Returns
-------
polars.DataFrame
polars.LazyFrame
"""
raise NotImplementedError
@@ -1229,6 +1269,7 @@ class Table(ABC):
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add more data to the [Table][lancedb.table.Table].
@@ -1280,6 +1321,10 @@ class Table(ABC):
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Store blob URIs that sit outside registered blob bases. The row
keeps a reference, so the object has to stay readable. Local
tables only.
Returns
-------
@@ -1905,14 +1950,23 @@ class Table(ABC):
@abstractmethod
def add_columns(
self, transforms: Dict[str, str] | pa.Field | List[pa.Field] | pa.Schema
self,
transforms: Dict[str, str | FunctionApplication]
| FunctionApplication
| pa.Field
| List[pa.Field]
| pa.Schema
| None = None,
*,
computed: Dict[str, str] | None = None,
):
"""
Add new columns with defined values.
Parameters
----------
transforms: Dict[str, str], pa.Field, List[pa.Field], pa.Schema
transforms: Dict[str, str | FunctionApplication], FunctionApplication,
pa.Field, List[pa.Field], pa.Schema
A map of column name to a SQL expression to use to calculate the
value of the new column. These expressions will be evaluated for
each row in the table, and can reference existing columns.
@@ -1920,10 +1974,109 @@ class Table(ABC):
new columns with the specified data types. The new columns will
be initialized with null values.
A mapping with one ``FunctionApplication`` value keeps its scalar
or named-struct result in the named table column. A bare
named-struct application expands its ordered result fields as one
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only on LanceDB Cloud and
Enterprise.
computed: Dict[str, str], optional
A map of column name to a SQL expression defining the column. The
column's type and inputs are derived from the expression, so no
data type is supplied.
Unlike ``transforms``, the expression is stored rather than
evaluated now: the column is committed with no values, and rows get
them from [`refresh_column`][lancedb.table.Table.refresh_column].
Declaring one therefore costs the same on a large table as on an
empty one.
A refresh does not revisit rows it has already filled, so mutating
an input leaves the value computed at fill time; recomputing means
dropping the column and declaring it again. While a declaration
reads a column, that column cannot be renamed, retyped or dropped.
On LanceDB Cloud and Enterprise the expression is planned by the
server, and the refresh runs as a server job -- see
[`refresh_column_async`][lancedb.table.Table.refresh_column_async].
Cannot be combined with ``transforms``.
Returns
-------
AddColumnsResult
version: the new version number of the table after adding columns.
Examples
--------
>>> import lancedb
>>> db = lancedb.connect("./.lancedb")
>>> table = db.create_table("computed_demo", [{"x": 1}, {"x": 2}])
>>> table.add_columns(computed={"doubled": "x * 2"})
AddColumnsResult(version=2)
>>> table.refresh_column("doubled")
RefreshColumnResult(rows_filled=2, version=3)
>>> table.to_arrow().sort_by("x").to_pandas()
x doubled
0 1 2
1 2 4
"""
@abstractmethod
def refresh_column(self, column: str) -> "RefreshColumnResult":
"""
Fill the rows of a computed column that hold no value yet.
Declared with ``add_columns(computed=...)``, a column starts empty and
gets its values here. Rows appended since the last refresh are filled
by the next one; rows already filled are left as they are, so the call
is idempotent and does not observe a mutated input.
Local tables only: a remote refresh runs as a server job, through
[`refresh_column_async`][lancedb.table.Table.refresh_column_async].
Parameters
----------
column: str
The name of the computed column to fill.
Returns
-------
RefreshColumnResult
rows_filled: the number of rows given a value.
version: the new version number of the table.
"""
@abstractmethod
def refresh_column_async(self, column: str) -> Job[RefreshColumnJobResult]:
"""
Like :meth:`refresh_column`, but returns a handle to the refresh job
instead of blocking until it completes.
The job may already be complete when returned; callers must not assume
the column is filled until :meth:`Job.wait` returns. Invalid input --
an unknown column, or one that is not computed -- raises here rather
than failing the job. On local tables the job runs in-process; on
LanceDB Cloud and Enterprise it is the server's backfill job.
Returns
-------
Job[RefreshColumnResult]
A job whose successful ``wait`` returns row counts plus the source
and published table versions.
Examples
--------
>>> import lancedb
>>> db = lancedb.connect("./.lancedb")
>>> table = db.create_table("computed_job_demo", [{"x": 1}, {"x": 2}])
>>> table.add_columns(computed={"doubled": "x * 2"})
AddColumnsResult(version=2)
>>> job = table.refresh_column_async("doubled")
>>> result = job.wait()
>>> result.rows_assigned
2
>>> job.status()
'finished'
"""
@abstractmethod
@@ -2182,11 +2335,15 @@ class LanceTable(Table):
return self.name
@classmethod
def from_inner(cls, tbl: LanceDBTable):
from .db import LanceDBConnection
async def from_inner(cls, tbl: LanceDBTable):
from .db import AsyncConnection, LanceDBConnection
async_tbl = AsyncTable(tbl)
conn = LanceDBConnection.from_inner(tbl.database())
inner_conn = tbl.database()
read_consistency_interval = await AsyncConnection(
inner_conn
).get_read_consistency_interval()
conn = LanceDBConnection.from_inner(inner_conn, read_consistency_interval)
return cls(
conn,
async_tbl.name,
@@ -2565,6 +2722,9 @@ class LanceTable(Table):
2. Currently we've disabled push-down of the filters from polars
because polars pushdown into pyarrow uses pyarrow compute
expressions rather than SQl strings (which LanceDB supports)
3. The Polars streaming engine is not supported because it does not
currently implement Python PyArrow dataset scans. Use the default
engine when collecting this LazyFrame.
Returns
-------
@@ -2573,8 +2733,12 @@ class LanceTable(Table):
from lancedb.integrations.pyarrow import PyarrowDatasetAdapter
dataset = PyarrowDatasetAdapter(self)
return pl.scan_pyarrow_dataset(
dataset, allow_pyarrow_filter=False, batch_size=batch_size
# Polars 1.32's non-PyArrow callback path passes batch_size twice. Keep
# the compatible PyArrow path, but block predicates because this adapter
# cannot translate PyArrow expressions into LanceDB filters.
return pl.scan_pyarrow_dataset(dataset, batch_size=batch_size).map_batches(
_polars_predicate_pushdown_barrier,
predicate_pushdown=False,
)
# New unified API overload
@@ -3250,6 +3414,7 @@ class LanceTable(Table):
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add data to the table.
If vector columns are missing and the table
@@ -3277,6 +3442,9 @@ class LanceTable(Table):
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Allow blob URIs outside registered bases. See :meth:`Table.add`.
Local tables only.
Returns
-------
@@ -3293,6 +3461,7 @@ class LanceTable(Table):
fill_value=fill_value,
progress=progress,
write_parallelism=write_parallelism,
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
)
)
finally:
@@ -3917,9 +4086,29 @@ class LanceTable(Table):
return LOOP.run(self._table.index_stats(index_name))
def add_columns(
self, transforms: Dict[str, str] | pa.field | List[pa.field] | pa.Schema
self,
transforms: Dict[str, str | FunctionApplication]
| FunctionApplication
| pa.Field
| List[pa.Field]
| pa.Schema
| None = None,
*,
computed: Dict[str, str] | None = None,
) -> AddColumnsResult:
return LOOP.run(self._table.add_columns(transforms))
return LOOP.run(self._table.add_columns(transforms, computed=computed))
def refresh_column(self, column: str) -> "RefreshColumnResult":
"""Fill a computed column's unfilled rows. See
[`AsyncTable.refresh_column`][lancedb.AsyncTable.refresh_column]."""
return LOOP.run(self._table.refresh_column(column))
def refresh_column_async(self, column: str) -> Job[RefreshColumnJobResult]:
"""Fill a computed column's unfilled rows, returning a handle to the
refresh job. See
[`Table.refresh_column_async`][lancedb.table.Table.refresh_column_async].
"""
return Job(LOOP.run(self._table.refresh_column_async(column)))
def alter_columns(
self, *alterations: Iterable[Dict[str, str]]
@@ -3954,6 +4143,28 @@ class LanceTable(Table):
[`AsyncTable.get_lsm_write_spec`][lancedb.AsyncTable.get_lsm_write_spec]."""
return LOOP.run(self._table.get_lsm_write_spec())
def checkpoint_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""
return LOOP.run(self._table.checkpoint_lsm())
def flush_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.flush_lsm`][lancedb.AsyncTable.flush_lsm]."""
return LOOP.run(self._table.flush_lsm())
def compact_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.compact_lsm`][lancedb.AsyncTable.compact_lsm]."""
return LOOP.run(self._table.compact_lsm())
def get_lsm_stats(self, *, include_generation_rows: bool = False) -> Optional[dict]:
"""Synchronous version of
[`AsyncTable.get_lsm_stats`][lancedb.AsyncTable.get_lsm_stats]."""
return LOOP.run(
self._table.get_lsm_stats(include_generation_rows=include_generation_rows)
)
def close_lsm_writers(self) -> None:
"""Close cached MemWAL shard writers. See
[`AsyncTable.close_lsm_writers`][lancedb.AsyncTable.close_lsm_writers]."""
@@ -4632,6 +4843,13 @@ class AsyncTable:
via [`set_unenforced_primary_key`]; bucket sharding additionally
requires it to be the single column being bucketed.
By default the MemWAL maintains every index on the table, resolved
here — a snapshot, so an index created afterwards needs the spec unset
and set again. This fails if one cannot be maintained; name the set
with ``with_maintained_indexes`` to install anyway. That pins an exact
set (a still-building index is rejected, not omitted); ``[]`` maintains
none.
Parameters
----------
spec : LsmWriteSpec
@@ -4639,7 +4857,7 @@ class AsyncTable:
Examples
--------
>>> from lancedb._lancedb import LsmWriteSpec
>>> from lancedb import LsmWriteSpec
>>> # table.set_unenforced_primary_key("id")
>>> # table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 16))
"""
@@ -4658,12 +4876,73 @@ class AsyncTable:
Returns ``None`` when the MemWAL LSM write path is not enabled (no
spec has been set, or it was removed with `unset_lsm_write_spec`).
The returned spec — including its ``maintained_indexes`` and
``writer_config_defaults`` — mirrors what was passed to
`set_lsm_write_spec`.
The returned spec mirrors what was passed to `set_lsm_write_spec`,
except that ``maintained_indexes`` always reports the concrete list
resolved when the spec was set — ``None`` never round-trips.
"""
return await self._inner.get_lsm_write_spec()
async def checkpoint_lsm(self) -> None:
"""Converge this table's LSM write path into its base table.
One flush, sealing every memtable into L0, then compaction triggers
until every generation that existed at that moment has reached base.
The loop runs client-side, reading progress from ``get_lsm_stats``.
Best-effort: generations created *while* it runs are deliberately not
waited on, which is what lets it terminate on a table taking writes.
Idempotent and safe on a cadence.
There is no deadline, and the caller owns that. It returns when the
target generations are gone, raises on a terminal server fault, and
otherwise waits however long the server takes. A slow table and a
stuck one are the same picture from the client: the compactor pool is
shared across every table on the node, so a checkpoint queued behind
unrelated work looks exactly like one that is merging. Wrap this in
``asyncio.wait_for`` for a wall-clock bound; abandoning it partway
costs nothing.
"""
await self._inner.checkpoint_lsm()
async def flush_lsm(self) -> None:
"""Seal every bucket's active memtable into L0.
Does not touch the base table — moving L0 into base is
`compact_lsm`. On a node that has not claimed this table, this claims
it and replays its WAL log first.
"""
await self._inner.flush_lsm()
async def compact_lsm(self) -> None:
"""Trigger a background L0 to base compaction pass per bucket.
Returns once the passes are dispatched, not once they finish: watch
``get_lsm_stats`` for progress, or use ``checkpoint_lsm`` to loop
until the current L0 has reached base.
"""
await self._inner.compact_lsm()
async def get_lsm_stats(
self, *, include_generation_rows: bool = False
) -> Optional[dict]:
"""Read live per-bucket LSM state.
Answers "how far behind is my fresh tier", "which bucket is hot", and
"why is my fresh-tier vector search brute-force". Mutates no table
state, though on a node that has not claimed this table it claims it,
exactly as a read would.
Returns ``None`` only when the LSM write path is not enabled.
Parameters
----------
include_generation_rows
Report a row count per L0 generation. Off by default: each count
opens an uncached Lance dataset, and ``checkpoint_lsm`` polls this
needing only generation numbers.
"""
return await self._inner.get_lsm_stats(include_generation_rows)
async def close_lsm_writers(self) -> None:
"""Drain and close any cached MemWAL shard writers for this table.
@@ -5096,6 +5375,7 @@ class AsyncTable:
fill_value: Optional[float] = None,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add more data to the [AsyncTable][lancedb.table.AsyncTable].
@@ -5126,6 +5406,9 @@ class AsyncTable:
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Allow blob URIs outside registered bases. See :meth:`Table.add`.
Local tables only.
"""
schema = await self.schema()
@@ -5162,6 +5445,7 @@ class AsyncTable:
mode or "append",
progress=progress,
write_parallelism=write_parallelism,
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
)
except RuntimeError as e:
if "Cast error" in str(e):
@@ -5744,37 +6028,166 @@ class AsyncTable:
return await self._inner.update(updates_sql, where)
async def add_columns(
self, transforms: dict[str, str] | pa.field | List[pa.field] | pa.Schema
self,
transforms: dict[str, str | FunctionApplication]
| FunctionApplication
| pa.Field
| List[pa.Field]
| pa.Schema
| None = None,
*,
computed: dict[str, str] | None = None,
) -> AddColumnsResult:
"""
Add new columns with defined values.
Parameters
----------
transforms: Dict[str, str]
transforms: Dict[str, str | FunctionApplication] or FunctionApplication
A map of column name to a SQL expression to use to calculate the
value of the new column. These expressions will be evaluated for
each row in the table, and can reference existing columns.
Alternatively, you can pass a pyarrow field or schema to add
new columns with NULLs.
A mapping with one ``FunctionApplication`` value keeps its scalar
or named-struct result in the named table column. A bare
named-struct application expands its ordered result fields as one
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only on LanceDB Cloud and
Enterprise.
computed: Dict[str, str], optional
A map of column name to a SQL expression defining the column. The
column's type and inputs are derived from the expression.
Unlike ``transforms``, the expression is stored rather than
evaluated now: the column is committed with no values, and rows get
them from
[`refresh_column`][lancedb.table.AsyncTable.refresh_column].
A refresh does not revisit rows it has already filled, so mutating
an input leaves the value computed at fill time. While a
declaration reads a column, that column cannot be renamed, retyped
or dropped.
On LanceDB Cloud and Enterprise the expression is planned by
the server. Cannot be combined with ``transforms``.
Returns
-------
AddColumnsResult
version: the new version number of the table after adding columns.
"""
function_application = None
function_output_name = None
if isinstance(transforms, FunctionApplication):
function_application = transforms
elif isinstance(transforms, dict) and any(
isinstance(value, FunctionApplication) for value in transforms.values()
):
if len(transforms) != 1 or not all(
isinstance(value, FunctionApplication) for value in transforms.values()
):
raise ValueError(
"one add_columns call declares exactly one Function binding"
)
function_output_name, function_application = next(iter(transforms.items()))
if function_application is not None:
if computed:
raise ValueError(
"add_columns cannot mix a Function application with SQL "
"computed columns"
)
function_application._ensure_declarable()
return await self._inner.add_function_columns(
function_application.to_canonical_json(), function_output_name
)
if isinstance(transforms, pa.Field):
transforms = [transforms]
if isinstance(transforms, list) and all(
{isinstance(f, pa.Field) for f in transforms}
):
transforms = pa.schema(transforms)
if computed:
if transforms:
raise ValueError(
"add_columns cannot take both transforms and computed columns"
)
return await self._inner.add_computed_columns(list(computed.items()))
if transforms is None:
raise ValueError("add_columns requires transforms or computed columns")
if isinstance(transforms, pa.Schema):
return await self._inner.add_columns_with_schema(transforms)
else:
return await self._inner.add_columns(list(transforms.items()))
async def refresh_column(self, column: str) -> RefreshColumnResult:
"""
Fill the rows of a computed column that hold no value yet.
Declared with ``add_columns(computed=...)``, a column starts empty and
gets its values here. Rows appended since the last refresh are filled
by the next one; rows already filled are left as they are, so the call
is idempotent and does not observe a mutated input.
Local tables only: a remote refresh runs as a server job, through
[`refresh_column_async`][lancedb.table.Table.refresh_column_async].
Parameters
----------
column: str
The name of the computed column to fill.
Returns
-------
RefreshColumnResult
The number of rows filled and the new version of the table.
"""
return await self._inner.refresh_column(column)
async def refresh_column_async(
self, column: str
) -> AsyncJob[RefreshColumnJobResult]:
"""
Like :meth:`refresh_column`, but returns a handle to the refresh job
instead of blocking until it completes.
The job may already be complete when returned; callers must not assume
the column is filled until :meth:`AsyncJob.wait` resolves. Invalid
input -- an unknown column, or one that is not computed -- raises here
rather than failing the job. On local tables the job runs
in-process; on LanceDB Cloud and Enterprise it is the server's
backfill job.
Returns
-------
AsyncJob[RefreshColumnResult]
A job whose successful ``wait`` returns row counts plus the source
and published table versions.
Examples
--------
>>> import asyncio
>>> import lancedb
>>> async def refresh_in_background():
... db = await lancedb.connect_async("./.lancedb")
... table = await db.create_table("computed_job_async_demo", [{"x": 1}])
... await table.add_columns(computed={"doubled": "x * 2"})
... job = await table.refresh_column_async("doubled")
... result = await job.wait()
... assert result.rows_assigned == 1
... return await job.status()
>>> asyncio.run(refresh_in_background())
'finished'
"""
return _typed_job(
await self._inner.refresh_column_async(column),
RefreshColumnJobResult.from_json,
)
async def alter_columns(
self, *alterations: Iterable[dict[str, Any]]
) -> AlterColumnsResult:
@@ -6229,7 +6642,9 @@ class TableStatistics:
Attributes
----------
total_bytes: int
The total number of bytes in the table.
The total size, in bytes, of the table's data files, index files, and
overlay files. Read from the manifest, so this excludes deletion files
and manifests.
num_rows: int
The total number of rows in the table.
num_indices: int
@@ -6424,21 +6839,21 @@ class Branches:
"""Diff a branch against main."""
return LOOP.run(self._table.branches.diff(from_branch))
def merge(self, from_branch: str, dry_run: bool = False) -> Dict[str, Any]:
"""Merge a branch into main, or dry-run.
def cherry_pick(self, from_branch: str, dry_run: bool = False) -> Dict[str, Any]:
"""Cherry-pick a branch onto main, or dry-run.
Parameters
----------
from_branch: str
Branch to merge from.
Branch to cherry-pick from.
dry_run: bool, default False
When True, only preview. When False, attempt the merge.
When True, only preview. When False, attempt the cherry-pick.
Notes
-----
A rejected merge returns ``status="rejected"`` instead of raising.
A failed cherry-pick returns ``status="failed"`` instead of raising.
"""
return LOOP.run(self._table.branches.merge(from_branch, dry_run))
return LOOP.run(self._table.branches.cherry_pick(from_branch, dry_run))
def _wrap(
self, async_table: "AsyncTable", version: Optional[int] = None
@@ -6574,9 +6989,11 @@ class AsyncBranches:
"""Diff a branch against main."""
return await self._table.branches.diff(from_branch)
async def merge(self, from_branch: str, dry_run: bool = False) -> Dict[str, Any]:
"""Merge a branch into main, or dry-run.
async def cherry_pick(
self, from_branch: str, dry_run: bool = False
) -> Dict[str, Any]:
"""Cherry-pick a branch onto main, or dry-run.
A rejected merge returns ``status="rejected"`` instead of raising.
A failed cherry-pick returns ``status="failed"`` instead of raising.
"""
return await self._table.branches.merge(from_branch, dry_run)
return await self._table.branches.cherry_pick(from_branch, dry_run)
python/python/lancedb/util.py
+5
View File
@@ -395,6 +395,11 @@ def _(value: dict):
)
@value_to_sql.register(pa.Scalar)
def _(value: pa.Scalar):
return value_to_sql(value.as_py())
@value_to_sql.register(np.ndarray)
def _(value: np.ndarray):
return value_to_sql(value.tolist())
python/python/tests/test_blob.py
+68
View File
@@ -617,3 +617,71 @@ def test_fetch_blobs_nested_path_survives_sort_after_query():
def _identifiable_payload(size: int) -> bytes:
block = 256
return b"".join(bytes([i % 256]) * block for i in range(size // block))
def _external_uri_blob_array(uris):
blob_type = lancedb.blob("image").type
storage_type = blob_type.storage_type
child_names = [field.name for field in storage_type]
assert "uri" in child_names, "blob layout no longer has a uri child"
children = [
pa.array(uris if field.name == "uri" else [None] * len(uris), type=field.type)
for field in storage_type
]
storage = pa.StructArray.from_arrays(children, fields=list(storage_type))
return pa.ExtensionArray.from_storage(blob_type, storage)
def _external_uri_table_and_rows(name, uris):
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table(name, schema=schema)
rows = pa.Table.from_arrays(
[
pa.array(range(len(uris)), type=pa.int64()),
_external_uri_blob_array(uris),
],
schema=schema,
)
return table, rows
def test_add_external_uri_struct_round_trips_with_flag(tmp_path):
payload = b"external-uri-bytes"
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(payload)
table, rows = _external_uri_table_and_rows("external_struct", [blob_path.as_uri()])
table.add(rows, allow_external_blob_outside_bases=True)
hits = table.search().to_arrow()
blobs = table.fetch_blobs("image", hits)
assert blobs[0].as_py() == payload
def test_add_external_uri_without_flag_raises(tmp_path):
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(b"unreachable")
table, rows = _external_uri_table_and_rows("external_no_flag", [blob_path.as_uri()])
with pytest.raises(ValueError, match="allow_external_blob_outside_bases"):
table.add(rows)
assert table.count_rows() == 0
def test_add_external_uri_string_round_trips_with_flag(tmp_path):
payload = b"external-uri-bytes"
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(payload)
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("external_string", schema=schema)
table.add(
[{"id": 1, "image": blob_path.as_uri()}],
allow_external_blob_outside_bases=True,
)
hits = table.search().to_arrow()
blobs = table.fetch_blobs("image", hits)
assert blobs[0].as_py() == payload
python/python/tests/test_db.py
+47 -5
View File
@@ -2,9 +2,11 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import inspect
import re
import sys
from datetime import timedelta
from importlib import resources
import os
from types import SimpleNamespace
@@ -17,6 +19,10 @@ from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
from lancedb.pydantic import LanceModel, Vector
def test_package_includes_pep_561_marker():
assert resources.files(lancedb).joinpath("py.typed").is_file()
def test_basic(tmp_path):
db = lancedb.connect(tmp_path)
@@ -62,21 +68,44 @@ def test_basic(tmp_path):
assert db.open_table("test").name == db["test"].name
def test_sync_repr_does_not_use_background_loop(tmp_path, monkeypatch):
def test_sync_debugger_inspection_does_not_use_background_loop(tmp_path, monkeypatch):
from lancedb.background_loop import LOOP
db = lancedb.connect(tmp_path)
table = db.create_table("test", data=[{"id": 1}])
def fail_run(*args, **kwargs):
raise AssertionError("repr should not use the Python background loop")
raise AssertionError("debugger inspection should not use the background loop")
monkeypatch.setattr(LOOP, "run", fail_run)
# Debuggers enumerate and evaluate every exposed attribute when expanding a
# variable. This must remain safe while their breakpoint suspends LOOP's thread.
members = dict(inspect.getmembers(db))
assert members["uri"] == str(tmp_path)
assert members["read_consistency_interval"] is None
assert repr(db) == f"LanceDBConnection(uri={str(tmp_path)!r})"
assert repr(table) == f"LanceTable(name='test', _conn={db!r})"
def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkeypatch):
from lancedb.background_loop import LOOP
from lancedb.db import LanceDBConnection
consistency_interval = timedelta(seconds=5)
db = lancedb.connect(tmp_path, read_consistency_interval=consistency_interval)
db_from_inner = LanceDBConnection.from_inner(db._inner, consistency_interval)
def fail_run(*args, **kwargs):
raise AssertionError("properties should not use the Python background loop")
monkeypatch.setattr(LOOP, "run", fail_run)
assert db.read_consistency_interval == consistency_interval
assert db_from_inner.read_consistency_interval == consistency_interval
def test_ingest_pd(tmp_path):
db = lancedb.connect(tmp_path)
@@ -726,8 +755,7 @@ def test_delete_table(tmp_db: lancedb.DBConnection):
assert tmp_db.table_names() == []
@pytest.mark.asyncio
async def test_delete_table_async(tmp_db: lancedb.DBConnection):
def test_drop_table_async(tmp_db: lancedb.DBConnection):
data = pd.DataFrame(
{
"vector": [[3.1, 4.1], [5.9, 26.5]],
@@ -743,7 +771,10 @@ async def test_delete_table_async(tmp_db: lancedb.DBConnection):
assert tmp_db.table_names() == ["test"]
tmp_db.drop_table("test")
job = tmp_db.drop_table_async("test")
assert job.id is None
assert job.status() == "finished"
assert job.wait() is None
assert tmp_db.table_names() == []
tmp_db.create_table("test", data=data)
@@ -752,6 +783,17 @@ async def test_delete_table_async(tmp_db: lancedb.DBConnection):
tmp_db.drop_table("does_not_exist", ignore_missing=True)
@pytest.mark.asyncio
async def test_drop_table_async_connection(tmp_db_async: lancedb.AsyncConnection):
await tmp_db_async.create_table("test", data=pa.table({"id": [1, 2]}))
job = await tmp_db_async.drop_table_async("test")
assert job.id is None
assert await job.status() == "finished"
assert await job.wait() is None
assert await tmp_db_async.table_names() == []
def test_drop_database(tmp_db: lancedb.DBConnection):
data = pd.DataFrame(
{
python/python/tests/test_elastic_dataloader.py
+1481 -2
View File
File diff suppressed because it is too large Load Diff
python/python/tests/test_embeddings.py
+51 -27
View File
@@ -64,6 +64,23 @@ def test_embedding_function(tmp_path):
assert np.allclose(actual, expected)
def test_instructor_ndims_uses_instruction():
instructor = get_registry().get("instructor").create()
model = MagicMock()
model.encode.return_value = np.zeros((1, 384))
with patch.object(type(instructor), "get_model", return_value=model):
assert instructor.ndims() == 384
model.encode.assert_called_once_with(
[[instructor.source_instruction, "foo"]],
batch_size=instructor.batch_size,
show_progress_bar=instructor.show_progress_bar,
normalize_embeddings=instructor.normalize_embeddings,
device=instructor.device,
)
def test_embedding_function_variables():
@register("variable-testing")
class VariableTestingFunction(TextEmbeddingFunction):
@@ -115,34 +132,16 @@ def test_embedding_function_variables():
assert func.safe_model_dump()["secret_key"] == "$var:secret"
def test_parse_functions_with_variables():
@register("variable-parsing-test")
class VariableParsingFunction(TextEmbeddingFunction):
api_key: str
base_url: Optional[str] = None
@staticmethod
def sensitive_keys():
return ["api_key"]
def ndims(self):
return 10
def generate_embeddings(self, texts):
# Mock implementation that just returns random embeddings
# In real usage, this would use the api_key to call an API
return [np.random.rand(self.ndims()).tolist() for _ in texts]
def test_openai_variables_survive_metadata_round_trip():
registry = EmbeddingFunctionRegistry.get_instance()
registry.set_var("test_api_key", "sk-test-key-12345")
registry.set_var("test_base_url", "https://api.example.com")
conf = EmbeddingFunctionConfig(
source_column="text",
vector_column="vector",
function=registry.get("variable-parsing-test").create(
api_key="$var:test_api_key", base_url="$var:test_base_url"
function=registry.get("openai").create(
api_key="$var:test_api_key", base_url="https://api.example.com"
),
)
@@ -150,7 +149,10 @@ def test_parse_functions_with_variables():
# Create a mock arrow table with the metadata
schema = pa.schema(
[pa.field("text", pa.string()), pa.field("vector", pa.list_(pa.float32(), 10))]
[
pa.field("text", pa.string()),
pa.field("vector", pa.list_(pa.float32(), 1536)),
]
)
table = pa.table({"text": [], "vector": []}, schema=schema)
table = table.replace_schema_metadata(metadata)
@@ -164,13 +166,15 @@ def test_parse_functions_with_variables():
assert parsed_func.api_key == "sk-test-key-12345"
assert parsed_func.base_url == "https://api.example.com"
embeddings = parsed_func.generate_embeddings(["test text"])
assert len(embeddings) == 1
assert len(embeddings[0]) == 10
assert parsed_func.safe_model_dump()["api_key"] == "$var:test_api_key"
with patch("lancedb.embeddings.openai.attempt_import_or_raise") as import_openai:
parsed_func._openai_client
import_openai.return_value.OpenAI.assert_called_once_with(
api_key="sk-test-key-12345", base_url="https://api.example.com"
)
def test_embedding_with_bad_results(tmp_path):
@register("null-embedding")
@@ -627,3 +631,23 @@ def test_url_retrieve_downloads_image():
image_bytes = url_retrieve(image_url)
img = Image.open(io.BytesIO(image_bytes))
assert img.size[0] > 0 and img.size[1] > 0
def test_jina_generate_image_input_dict_local_path(tmp_path):
"""
JinaEmbeddings._generate_image_input_dict must accept a local image path
(str or Path), not just bytes. Previously it crashed with
`AttributeError: 'function' object has no attribute 'urlparse'` on any
str/Path input because it called `urlparse.urlparse(image)` instead of
`urlparse(image)` (urlparse was imported as a function, not a module).
"""
Image = pytest.importorskip("PIL.Image")
from lancedb.embeddings.jinaai import JinaEmbeddings
image_path = tmp_path / "test.png"
Image.new("RGB", (4, 4), color="red").save(image_path, format="PNG")
for image in (str(image_path), image_path):
image_dict = JinaEmbeddings._generate_image_input_dict(image)
assert "image" in image_dict
assert isinstance(image_dict["image"], str) and len(image_dict["image"]) > 0
python/python/tests/test_expr.py
+98
View File
@@ -632,3 +632,101 @@ class TestExprBytesIntegration:
.to_arrow()
)
assert result.num_rows == 2
# ── datetime / timezone integration for lit() (issue #3262) ──────────────────
class TestExprDatetimeTimezoneIntegration:
"""Integration coverage for lit(datetime) against table timestamp columns.
PyArrow stores naive timestamps as UTC wall-clock microseconds. Python's
datetime.timestamp() treats naive values as *local* time, which used to
shift lit(naive) by the host UTC offset and break equality filters on
non-UTC machines. These cases lock the expected semantics.
"""
def test_both_naive_match(self, tmp_path):
"""Table naive + lit naive with the same wall clock must match."""
db = lancedb.connect(str(tmp_path / "naive"))
ts = datetime(2024, 7, 1, 10, 0, 0)
table = db.create_table(
"t", [{"id": 1, "ts": ts}, {"id": 2, "ts": datetime(2024, 7, 2, 10, 0, 0)}]
)
result = table.search().where(col("ts") == lit(ts)).to_list()
assert len(result) == 1
assert result[0]["id"] == 1
def test_both_same_timezone_match(self, tmp_path):
"""Table UTC + lit UTC for the same instant must match."""
db = lancedb.connect(str(tmp_path / "utc"))
ts = datetime(2024, 7, 1, 10, 0, 0, tzinfo=timezone.utc)
table = db.create_table(
"t",
pa.table(
{
"id": [1, 2],
"ts": pa.array(
[ts, datetime(2024, 7, 2, 10, 0, 0, tzinfo=timezone.utc)],
type=pa.timestamp("us", tz="UTC"),
),
}
),
)
result = table.search().where(col("ts") == lit(ts)).to_list()
assert len(result) == 1
assert result[0]["id"] == 1
def test_different_timezones_same_instant(self, tmp_path):
"""UTC table row equals lit of the same instant in a different zone."""
db = lancedb.connect(str(tmp_path / "diff_tz"))
ts_utc = datetime(2024, 7, 1, 10, 0, 0, tzinfo=timezone.utc)
# Same instant as 06:00 in UTC-4
ts_est = datetime(2024, 7, 1, 6, 0, 0, tzinfo=timezone(timedelta(hours=-4)))
table = db.create_table(
"t",
pa.table(
{
"id": [1],
"ts": pa.array([ts_utc], type=pa.timestamp("us", tz="UTC")),
}
),
)
result = table.search().where(col("ts") == lit(ts_est)).to_list()
assert len(result) == 1
assert result[0]["id"] == 1
def test_table_tz_literal_naive(self, tmp_path):
"""UTC table + naive lit uses wall-clock equality (10:00 == 10:00 UTC)."""
db = lancedb.connect(str(tmp_path / "tz_naive"))
ts_utc = datetime(2024, 7, 1, 10, 0, 0, tzinfo=timezone.utc)
ts_naive = datetime(2024, 7, 1, 10, 0, 0)
table = db.create_table(
"t",
pa.table(
{
"id": [1],
"ts": pa.array([ts_utc], type=pa.timestamp("us", tz="UTC")),
}
),
)
result = table.search().where(col("ts") == lit(ts_naive)).to_list()
assert len(result) == 1
assert result[0]["id"] == 1
def test_table_naive_literal_aware(self, tmp_path):
"""Naive table + UTC lit with the same wall clock must match."""
db = lancedb.connect(str(tmp_path / "naive_aware"))
ts_naive = datetime(2024, 7, 1, 10, 0, 0)
ts_utc = datetime(2024, 7, 1, 10, 0, 0, tzinfo=timezone.utc)
table = db.create_table("t", [{"id": 1, "ts": ts_naive}])
result = table.search().where(col("ts") == lit(ts_utc)).to_list()
assert len(result) == 1
assert result[0]["id"] == 1
def test_naive_lit_sql_is_wall_clock_not_local_shifted(self):
"""Regression: naive lit must not apply the host local UTC offset."""
ts = datetime(2024, 7, 1, 10, 0, 0)
sql = lit(ts).to_sql()
# Must encode 10:00 wall clock, not 10:00+local_offset.
assert "2024-07-01 10:00:00" in sql
python/python/tests/test_first_class_function_slice1.py
+358
View File
@@ -0,0 +1,358 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import json
from pathlib import Path
import pytest
from lancedb import col
import lancedb.functions as functions
from lancedb.functions import (
FunctionApplication,
FunctionBinding,
FunctionVersion,
PythonRuntimeSpec,
RefreshColumnResult,
)
from lancedb.table import AsyncTable
FIXTURES = (
Path(__file__).parents[3]
/ "rust"
/ "lancedb"
/ "tests"
/ "fixtures"
/ "first_class_functions"
/ "v1"
)
def fixture(name: str) -> str:
return (FIXTURES / name).read_text()
def job_result(name: str) -> dict:
return json.loads(fixture(name))["result"]
def test_public_function_values_are_in_api_reference():
docs = Path(__file__).parents[3] / "docs" / "src" / "python" / "python.md"
rendered = docs.read_text()
for name in functions.__all__:
assert f"::: lancedb.functions.{name}" in rendered
@pytest.mark.parametrize(
("fixture_name", "canonical_name", "model", "nested_result"),
[
(
"remote_function_job.json",
"remote_function_version.canonical.json",
FunctionVersion,
True,
),
(
"remote_function_application.json",
"remote_function_application.canonical.json",
FunctionApplication,
False,
),
(
"remote_function_binding.json",
"remote_function_binding.canonical.json",
FunctionBinding,
False,
),
(
"remote_refresh_job.json",
"remote_refresh_result.canonical.json",
RefreshColumnResult,
True,
),
(
"remote_refresh_result_without_published_version.json",
"remote_refresh_result_without_published_version.canonical.json",
RefreshColumnResult,
False,
),
],
)
def test_python_and_rust_share_remote_canonical_goldens(
fixture_name, canonical_name, model, nested_result
):
value = json.loads(fixture(fixture_name))
if nested_result:
value = value["result"]
decoded = model.from_json(json.dumps(value))
assert decoded.to_canonical_json() == fixture(canonical_name).strip()
def test_function_version_identity_is_immutable_and_exact():
value = job_result("remote_function_job.json")
version = FunctionVersion.from_json(json.dumps(value))
assert version.name == "embed"
assert version.version == "fv_01K3EXACT"
with pytest.raises((TypeError, ValueError)):
version.version = "fv_changed"
with pytest.raises(TypeError, match="immutable"):
version.runtime.env["TOKENIZERS_PARALLELISM"] = "true"
changed = dict(value)
changed["version"] = "fv_changed"
assert FunctionVersion(**changed) != version
def test_function_version_binds_named_columns_as_one_immutable_application():
version = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
)
application = version(text=col("documents.body"))
assert application.function.name == version.name
assert application.function.version == version.version
assert application.output is version.signature.output
assert [
(value.parameter, value.kind, value.value["path"])
for value in application.inputs
] == [("text", "column", "documents.body")]
def test_function_version_binding_validates_names_and_direct_columns():
version = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
)
with pytest.raises(TypeError, match=r"missing inputs: \['text'\]"):
version()
with pytest.raises(TypeError, match=r"unknown inputs: \['body'\]"):
version(text=col("text"), body=col("body"))
with pytest.raises(TypeError, match="direct col"):
version(text=col("text").lower())
def test_function_version_keeps_named_struct_outputs_in_one_application():
value = job_result("remote_function_job.json")
value["name"] = "text_features"
value["version"] = "fv_multi_output"
value["signature"] = {
"inputs": [
{"name": "title", "arrow_type": "utf8", "nullable": True},
{"name": "body", "arrow_type": "utf8", "nullable": True},
],
"output": {
"kind": "named_struct",
"fields": [
{
"name": "normalized_text",
"arrow_type": "utf8",
"nullable": False,
},
{
"name": "token_count",
"arrow_type": "int64",
"nullable": False,
},
],
},
}
version = FunctionVersion(**value)
application = version(body=col("body"), title=col("title")).rename(
columns={
"normalized_text": "search_text",
"token_count": "search_token_count",
}
)
assert [value.parameter for value in application.inputs] == ["title", "body"]
assert [field.name for field in application.output.fields] == [
"normalized_text",
"token_count",
]
assert dict(application.columns) == {
"normalized_text": "search_text",
"token_count": "search_token_count",
}
def test_unknown_fields_and_discriminators_are_forward_decodable():
value = job_result("remote_function_job.json")
value["future_version_metadata"] = {"retention_class": "catalog"}
value["runtime"] = {"kind": "wasm", "module_digest": "sha256:wasm"}
value["signature"]["output"]["kind"] = "future_output_shape"
version = FunctionVersion.from_json(json.dumps(value))
assert version.runtime.kind == "wasm"
assert version.runtime.python_version is None
assert version.runtime.environment is None
assert json.loads(version.to_canonical_json())["runtime"] == {"kind": "wasm"}
assert version.signature.output.kind == "future_output_shape"
def test_function_application_uses_rename_columns_only():
application = FunctionApplication.from_json(
fixture("remote_function_application.json")
)
renamed = application.rename(columns={"normalized_text": "body_normalized"})
assert application.columns["normalized_text"] == "search_text"
assert renamed.columns["normalized_text"] == "body_normalized"
assert renamed.function == application.function
assert not hasattr(application, "rename_outputs")
with pytest.raises(TypeError, match="immutable"):
renamed.columns["normalized_text"] = "changed"
with pytest.raises(TypeError, match="immutable"):
application.inputs[0].value["path"] = "changed"
with pytest.raises(ValueError, match="unknown Function result fields"):
application.rename(columns={"missing": "search_text"})
with pytest.raises(ValueError, match="destinations must be unique"):
application.rename(columns={"normalized_text": "same", "token_count": "same"})
bare_value = json.loads(fixture("remote_function_application.json"))
bare_value.pop("columns")
bare = FunctionApplication(**bare_value)
with pytest.raises(ValueError, match="destinations must be unique"):
bare.rename(columns={"normalized_text": "token_count"})
def test_binding_and_refresh_result_keep_stable_remote_fields():
binding = FunctionBinding.from_json(fixture("remote_function_binding.json"))
assert binding.function.version == "fv_01K3TEXT"
assert [output.output_ordinal for output in binding.outputs] == [0, 1]
assert binding.input_schema is not None
assert binding.output_schema is not None
result = RefreshColumnResult.from_json(
json.dumps(job_result("remote_refresh_job.json"))
)
assert result.rows_filled == result.rows_assigned
assert result.version == result.published_version
result = RefreshColumnResult.from_json(
fixture("remote_refresh_result_without_published_version.json")
)
assert result.published_version is None
assert RefreshColumnResult.from_json(result.to_canonical_json()) == result
def test_function_literal_numeric_domain_matches_rust():
with pytest.raises(ValueError, match="floating-point Function literals"):
FunctionApplication.from_json(fixture("remote_function_application_float.json"))
value = json.loads(fixture("remote_function_application_float.json"))
value["inputs"][0]["value"] = 2**64
with pytest.raises(ValueError, match="outside the canonical JSON range"):
FunctionApplication.from_json(json.dumps(value))
def test_empty_default_maps_have_stable_canonical_bytes():
runtime = PythonRuntimeSpec(
kind="python", python_version="3.12", environment={"kind": "pip"}
)
assert runtime.to_canonical_json() == (
'{"environment":{"kind":"pip"},"kind":"python","python_version":"3.12"}'
)
value = json.loads(fixture("remote_function_application.json"))
value.pop("columns")
application = FunctionApplication.from_json(json.dumps(value))
assert "columns" not in json.loads(application.to_canonical_json())
@pytest.mark.parametrize("field", ["rows_assigned", "source_version"])
def test_refresh_result_rejects_non_u64_values(field):
value = job_result("remote_refresh_job.json")
value[field] = -1
with pytest.raises(ValueError):
RefreshColumnResult.from_json(json.dumps(value))
value[field] = "1"
with pytest.raises(ValueError):
RefreshColumnResult.from_json(json.dumps(value))
class _FunctionDeclarationInner:
def __init__(self):
self.calls = []
async def add_function_columns(self, application_json, output_name):
self.calls.append((json.loads(application_json), output_name))
return "declared"
def known_application() -> FunctionApplication:
value = json.loads(fixture("remote_function_application.json"))
value.pop("future_application")
return FunctionApplication(**value)
@pytest.mark.asyncio
async def test_add_columns_routes_struct_as_one_and_multi_output_binding_atomically():
inner = _FunctionDeclarationInner()
table = AsyncTable(inner)
application = known_application()
result = await table.add_columns(
{"features": application._copy(update={"columns": {}})}
)
assert result == "declared"
assert inner.calls[-1][1] == "features"
bare = application._copy(update={"columns": {}}).rename(
columns={"normalized_text": "search_text"}
)
result = await table.add_columns(bare)
assert result == "declared"
assert inner.calls[-1][1] is None
assert inner.calls[-1][0]["columns"] == {"normalized_text": "search_text"}
@pytest.mark.asyncio
async def test_add_columns_rejects_multiple_bindings_and_unknown_newer_application():
inner = _FunctionDeclarationInner()
table = AsyncTable(inner)
application = known_application()
with pytest.raises(ValueError, match="exactly one Function binding"):
await table.add_columns({"a": application, "b": application})
future = json.loads(fixture("remote_function_application.json"))
application = FunctionApplication(**future)
with pytest.raises(ValueError, match="newer contract"):
await table.add_columns(application)
future.pop("future_application")
future["output"]["assignment"] = "cell_flag"
application = FunctionApplication(**future)
assert "assignment" not in json.loads(application.to_canonical_json())["output"]
with pytest.raises(ValueError, match="output.assignment"):
await table.add_columns(application)
assert inner.calls == []
def test_rename_requires_named_struct_and_keeps_partial_mapping_immutable():
scalar = FunctionApplication.from_json(
json.dumps(
{
"function": {"name": "embed", "version": "fv_exact"},
"inputs": [],
"output": {
"kind": "scalar",
"arrow_type": "list<float32>",
"nullable": False,
},
}
)
)
with pytest.raises(ValueError, match="named-struct"):
scalar.rename(columns={"value": "embedding"})
application = known_application()._copy(update={"columns": {}})
renamed = application.rename(columns={"normalized_text": "search_text"})
assert dict(application.columns) == {}
assert dict(renamed.columns) == {"normalized_text": "search_text"}
python/python/tests/test_first_class_function_slice2.py
+620
View File
@@ -0,0 +1,620 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from __future__ import annotations
import base64
import contextlib
import functools
import importlib.util
import types
from datetime import date
import http.server
import json
from pathlib import Path
import threading
from typing import Optional
import pyarrow as pa
import pytest
import lancedb
from lancedb.functions import UdfDefinition, udf
THRESHOLD = 20
_CACHE = None
FIXTURES = (
Path(__file__).parents[3]
/ "rust"
/ "lancedb"
/ "tests"
/ "fixtures"
/ "first_class_functions"
/ "v1"
)
@udf(
pip=["numpy>=2"],
env={"MODE": "test"},
python_version="3.12",
)
def normalize_score(value: float) -> float:
return value / 100.0
def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
assert isinstance(normalize_score, UdfDefinition)
assert normalize_score(25.0) == 0.25
assert (
normalize_score.registration_request.to_canonical_json()
== (FIXTURES / "remote_function_registration_request.canonical.json")
.read_text()
.strip()
)
request = json.loads(normalize_score.registration_request.to_canonical_json())
assert request["artifact"]["adapter"] == {
"kind": "scalar_to_arrow_batch",
"version": 1,
}
def _run_packaged(definition, *args):
"""Execute the shipped artifact in a fresh namespace, as a worker would."""
source = base64.b64decode(definition.registration_request.artifact.content.data)
namespace: dict = {}
exec(compile(source, "<udf>", "exec"), namespace)
return namespace[definition.registration_request.artifact.entrypoint](*args)
def test_udf_packages_attribute_access_and_body_imports():
@udf
def word_norm(body: str) -> float:
import numpy as np
try:
words = body.split()
except AttributeError as error:
raise ValueError(str(error)) from error
return float(np.linalg.norm([len(w) for w in words]))
assert _run_packaged(word_norm, "aa bb") == pytest.approx(8**0.5)
def test_udf_packages_module_globals_and_global_caches():
@udf
def label(value: int) -> str:
return "big" if value >= THRESHOLD else "small"
assert _run_packaged(label, 21) == "big"
@udf
def cached(value: int) -> int:
global _CACHE
if _CACHE is None:
_CACHE = 40
return _CACHE + value
assert _run_packaged(cached, 2) == 42
def test_udf_annotations_are_not_runtime_names():
@udf
def identity(value: date) -> date:
return value
assert _run_packaged(identity, date(2026, 8, 25)) == date(2026, 8, 25)
def test_udf_nested_scopes_resolve_lexically():
@udf
def score(value: int) -> int:
offset = 2
def add_offset() -> int:
return value + offset
return add_offset() + sum(v for v in [0])
assert _run_packaged(score, 3) == 5
def test_udf_resolves_module_globals_before_builtins(tmp_path):
module_path = tmp_path / "shadowing_udfs.py"
module_path.write_text(
"max = 7\n"
"len = lambda _: 99\n"
"\n"
"def uses_literal_shadow(value: int) -> int:\n"
" def nested() -> int:\n"
" return max\n"
" return nested() + value\n"
"\n"
"def uses_callable_shadow(value: int) -> int:\n"
" def nested() -> int:\n"
" return len([1])\n"
" return nested() + value\n"
)
spec = importlib.util.spec_from_file_location("shadowing_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
# The module's `max = 7` is what the interpreter would use, so it ships.
assert _run_packaged(udf(module.uses_literal_shadow), 1) == 8
# A callable global cannot ship; it must not be silently swapped for the builtin.
with pytest.raises(TypeError, match="unsupported global value of type function"):
udf(module.uses_callable_shadow)
def test_canonical_arrow_type_is_exactly_the_grammar():
from lancedb.functions import _GRAMMAR_PRIMITIVES, _canonical_arrow_type
golden = json.loads(
(
Path(__file__).parents[3]
/ "rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json"
).read_text()
)
primitives = [
case["arrow_type"] for case in golden["valid"] if "<" not in case["arrow_type"]
]
assert [name for _, name in _GRAMMAR_PRIMITIVES] == primitives
for outside in [
pa.timestamp("us"),
pa.decimal128(10, 2),
pa.large_string(),
pa.large_binary(),
pa.binary(4),
pa.duration("s"),
pa.struct([pa.field("a", pa.int32())]),
pa.list_(pa.float32(), 0),
pa.list_(pa.timestamp("us")),
]:
with pytest.raises(TypeError, match="unsupported Arrow type"):
_canonical_arrow_type(outside)
def test_udf_nested_annotations_are_postponed_in_the_artifact():
@udf
def score(value: int) -> int:
def identity(item: date) -> date:
return item
identity(date(2026, 8, 25))
return value
assert _run_packaged(score, 3) == 3
def test_udf_ships_globals_the_body_deletes():
@udf
def clear(value: int) -> int:
global _CACHE
del _CACHE
return value
assert _run_packaged(clear, 3) == 3
def test_udf_rejects_a_module_global_that_does_not_import_as_itself(tmp_path):
module_path = tmp_path / "fake_module_udfs.py"
module_path.write_text(
"import types\n"
"np = types.ModuleType('numpy')\n"
"np.sqrt = lambda x: 0\n"
"\n"
"def score(value: int) -> int:\n"
" return int(np.sqrt(value))\n"
)
spec = importlib.util.spec_from_file_location("fake_module_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
with pytest.raises(TypeError, match="does not import as 'numpy'"):
udf(module.score)
def test_udf_rejects_a_module_level_namespace_alias(tmp_path):
module_path = tmp_path / "aliasing_udfs.py"
module_path.write_text(
"import builtins as b\n"
"THRESHOLD = 5\n"
"\n"
"def score(value: int) -> int:\n"
" return value + b.vars(b.__import__('aliasing_udfs'))['THRESHOLD']\n"
)
spec = importlib.util.spec_from_file_location("aliasing_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
with pytest.raises(ValueError, match="dynamic namespace access"):
udf(module.score)
@pytest.mark.parametrize(
"access",
[
"globals()['THRESHOLD']",
"eval('THRESHOLD')",
"(lambda g: g()['THRESHOLD'])(globals)",
"__import__('sys').modules[__name__].THRESHOLD",
"sys.modules[__name__].THRESHOLD",
],
)
def test_udf_rejects_dynamic_namespace_access(access):
namespace: dict = {}
exec(
f"def score(value: int) -> int:\n return value + {access}\n",
{"THRESHOLD": 5},
namespace,
)
with pytest.raises(ValueError, match="dynamic namespace access"):
_package_from_text(
"def score(value: int) -> int:\n"
" import sys\n"
f" return value + {access}\n"
)
def _package_from_text(source: str, module_globals: dict | None = None):
"""Load `source` as a real module file so the packager can inspect it."""
import tempfile
directory = tempfile.mkdtemp()
path = Path(directory) / "generated_udf_module.py"
path.write_text(source)
spec = importlib.util.spec_from_file_location(f"generated_udf_{id(source)}", path)
module = importlib.util.module_from_spec(spec)
if module_globals:
module.__dict__.update(module_globals)
spec.loader.exec_module(module)
functions = [
value
for value in vars(module).values()
if callable(value) and getattr(value, "__module__", None) == module.__name__
]
return udf(functions[0])
def test_udf_rejects_a_non_standard_builtins_environment():
def score(value: int) -> int:
return len([1]) + value
score.__globals__ # noqa: B018 -- real function, real globals
import builtins
patched = types.FunctionType(
score.__code__,
{"__builtins__": {**vars(builtins), "len": lambda _: 99}},
"score",
)
patched.__annotations__ = score.__annotations__
assert patched(3) == 102
with pytest.raises(ValueError, match="non-standard builtins environment"):
udf(patched)
class ReportingDict(dict): # reports standard entries, resolves differently
def __missing__(self, key):
return vars(builtins)[key]
disguised = types.FunctionType(
score.__code__, {"__builtins__": ReportingDict(len=lambda _: 99)}, "score"
)
disguised.__annotations__ = score.__annotations__
assert disguised(3) == 102
with pytest.raises(ValueError, match="non-standard builtins environment"):
udf(disguised)
hooked = types.FunctionType(
score.__code__,
{"__builtins__": {**vars(builtins), "__import__": lambda *a, **k: None}},
"score",
)
hooked.__annotations__ = score.__annotations__
with pytest.raises(ValueError, match="non-standard builtins environment"):
udf(hooked)
def test_udf_recursion_versus_a_rebound_module_name(tmp_path):
module_path = tmp_path / "rebound_udfs.py"
module_path.write_text(
"def fact(value: int) -> int:\n"
" return 1 if value <= 1 else value * fact(value - 1)\n"
"\n"
"def score(value: int) -> int:\n"
" return score + value\n"
)
spec = importlib.util.spec_from_file_location("rebound_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
assert _run_packaged(udf(module.fact), 5) == 120
raw = module.score
module.score = 10
with pytest.raises(ValueError, match="binds that name to another value"):
udf(raw)
# A wrapper that merely exposes __wrapped__ is not the function.
module.score = functools.wraps(raw)(lambda value: 41)
with pytest.raises(ValueError, match="binds that name to another value"):
udf(raw)
# The decorator's own result is; a subclass of it is not.
module.fact = udf(module.fact)
assert _run_packaged(module.fact, 4) == 24
class Twisted(UdfDefinition):
def __call__(self, *args, **kwargs):
return 41
raw_fact = module.fact._function
module.fact = Twisted(
raw_fact,
name=None,
input_schema=None,
output_schema=None,
pip=(),
env={},
python_version=None,
)
with pytest.raises(ValueError, match="binds that name to another value"):
udf(raw_fact)
def test_canonical_arrow_type_rejects_unrepresentable_list_children():
from lancedb.functions import _canonical_arrow_type
for outside in [
pa.list_(pa.float32()), # pyarrow default: nullable child
pa.list_(pa.field("custom", pa.float32(), nullable=False)),
pa.list_(pa.field("item", pa.float32(), nullable=False, metadata={"k": "v"})),
pa.list_(pa.field("item", pa.float32(), nullable=False), 0),
]:
with pytest.raises(TypeError, match="unsupported Arrow type"):
_canonical_arrow_type(outside)
assert (
_canonical_arrow_type(
pa.list_(pa.field("item", pa.float32(), nullable=False), 3)
)
== "fixed_size_list<float32, 3>"
)
def _calls_missing(value: int) -> int:
return missing(value) # noqa: F821
def _shadows_missing_in_a_comprehension(value: int) -> int:
return missing(value) + sum(missing for missing in ()) # noqa: F821
def _shadows_missing_in_a_lambda(value: int) -> int:
return (lambda missing: missing)(value) + missing # noqa: F821
@pytest.mark.parametrize(
"function",
[_calls_missing, _shadows_missing_in_a_comprehension, _shadows_missing_in_a_lambda],
)
def test_udf_rejects_a_truly_unresolved_global(function):
with pytest.raises(ValueError, match=r"unresolved global names: \['missing'\]"):
udf(function)
def _arrow_type_from_golden(spec: dict) -> pa.DataType:
kind = spec["type"]
if kind in ("list", "large_list", "fixed_size_list"):
item = _arrow_type_from_golden(spec["fields"][0]["type"])
field = pa.field("item", item, nullable=False)
if kind == "list":
return pa.list_(field)
if kind == "large_list":
return pa.large_list(field)
return pa.list_(field, spec["length"])
return {
"null": pa.null(),
"bool": pa.bool_(),
"utf8": pa.string(),
"binary": pa.binary(),
"float16": pa.float16(),
"float32": pa.float32(),
"float64": pa.float64(),
"date32": pa.date32(),
"date64": pa.date64(),
}.get(kind) or getattr(pa, kind)()
def test_arrow_type_grammar_matches_the_shared_golden():
golden = json.loads(
(
Path(__file__).parents[3]
/ "rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json"
).read_text()
)
from lancedb.functions import _canonical_arrow_type
emitted = {
case["arrow_type"]: _canonical_arrow_type(_arrow_type_from_golden(case["json"]))
for case in golden["valid"]
}
assert emitted == {
case["arrow_type"]: case["arrow_type"] for case in golden["valid"]
}
assert not set(emitted) & set(golden["invalid"])
for case in golden["server_only"]:
with pytest.raises(TypeError, match="unsupported Arrow type"):
_canonical_arrow_type(_arrow_type_from_golden(case["json"]))
def test_explicit_arrow_schema_is_deterministic():
input_schema = pa.schema([pa.field("value", pa.float32(), nullable=True)])
output_schema = pa.field(
"embedding",
pa.list_(pa.field("item", pa.float32(), nullable=False), 3),
nullable=False,
)
@udf(input_schema=input_schema, output_schema=output_schema)
def explicit(value):
return [value, value, value]
signature = explicit.registration_request.signature
assert signature.inputs[0].arrow_type == "float32"
assert signature.inputs[0].nullable is True
assert signature.output.arrow_type == "fixed_size_list<float32, 3>"
assert signature.output.nullable is False
def test_annotation_and_explicit_schema_validation_fail_closed():
with pytest.raises(TypeError, match="missing Function annotations"):
@udf
def missing(value):
return value
with pytest.raises(TypeError, match="unsupported Function annotation"):
@udf
def unsupported(value: set[str]) -> str:
return ""
with pytest.raises(ValueError, match="output must be non-nullable"):
@udf
def nullable_output(value: int) -> Optional[int]:
return value
with pytest.raises(ValueError, match="provided together"):
@udf(input_schema=pa.schema([pa.field("value", pa.int64())]))
def partial_schema(value):
return value
with pytest.raises(ValueError, match="exactly match callable parameters"):
@udf(
input_schema=pa.schema([pa.field("other", pa.int64())]),
output_schema=pa.int64(),
)
def wrong_name(value):
return value
with pytest.raises(ValueError, match="output must be non-nullable"):
@udf(
input_schema=pa.schema([pa.field("value", pa.int64())]),
output_schema=pa.field("result", pa.int64(), nullable=True),
)
def nullable_explicit(value):
return value
def test_local_function_catalog_operations_are_not_supported(tmp_path):
db = lancedb.connect(tmp_path)
message = "Function catalog operations are not supported by this database"
with pytest.raises(NotImplementedError, match=message):
db.create_function(normalize_score)
with pytest.raises(NotImplementedError, match=message):
db.create_function_async(normalize_score)
with pytest.raises(NotImplementedError, match=message):
db.get_function("normalize_score", version="fv_exact")
@contextlib.contextmanager
def _mock_remote_function_catalog():
state = {"requests": [], "version": None}
class Handler(http.server.BaseHTTPRequestHandler):
def log_message(self, *args):
pass
def do_POST(self):
length = int(self.headers.get("Content-Length", "0"))
body = json.loads(self.rfile.read(length) or b"{}")
state["requests"].append((self.path, body))
status = 200
if self.path == "/v1/functions/create":
state["version"] = {
"name": body["name"],
"version": "fv_exact",
"artifact": {
key: body["artifact"][key]
for key in ("kind", "digest", "entrypoint")
},
"signature": body["signature"],
"runtime": body["runtime"],
"runtime_digest": "sha256:runtime",
"environment_digest": "sha256:environment",
"created_at": "2026-08-21T00:00:00Z",
}
response = {"job_id": "job-register"}
status = 202
elif self.path == "/v1/jobs/describe":
assert body == {"job_id": "job-register"}
response = {
"job_id": "job-register",
"job_type": "create_function",
"job_state": "DONE",
"result": state["version"],
}
elif self.path == "/v1/functions/describe":
assert body == {
"name": "normalize_score",
"version": "fv_exact",
}
response = state["version"]
else:
status = 404
response = {"error": "not found"}
encoded = json.dumps(response).encode()
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(encoded)))
self.end_headers()
self.wfile.write(encoded)
with http.server.HTTPServer(("localhost", 0), Handler) as server:
thread = threading.Thread(target=server.serve_forever)
thread.start()
try:
yield f"http://localhost:{server.server_address[1]}", state
finally:
server.shutdown()
thread.join()
def test_remote_registration_job_and_exact_version_reopen_round_trip():
with _mock_remote_function_catalog() as (host, state):
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
registration = db.create_function_async(normalize_score)
assert registration.id == "job-register"
created = registration.wait()
reopened = db.get_function("normalize_score", version=created.version)
assert created == reopened
assert reopened.name == "normalize_score"
assert reopened.version == "fv_exact"
create_request = state["requests"][0][1]
assert create_request == json.loads(
normalize_score.registration_request.to_canonical_json()
)
def test_blocking_remote_registration_returns_function_version():
with _mock_remote_function_catalog() as (host, state):
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
created = db.create_function(normalize_score)
assert created.name == "normalize_score"
assert created.version == "fv_exact"
assert [path for path, _ in state["requests"]] == [
"/v1/functions/create",
"/v1/jobs/describe",
]
python/python/tests/test_fts.py
+8
View File
@@ -245,6 +245,14 @@ def test_create_inverted_index_rejects_invalid_block_size(table):
table.create_index("text", config=FTS(block_size=129))
def test_create_inverted_index_respects_build_memory_limit(table):
with pytest.raises(ValueError, match="exceeds worker memory limit"):
table.create_index(
"text",
config=FTS(memory_limit=0, num_workers=1),
)
def test_custom_stop_words_list(table):
table.create_index(
"text",
python/python/tests/test_hybrid_query.py
+106 -1
View File
@@ -12,7 +12,7 @@ import pyarrow.compute as pc
import pytest
import pytest_asyncio
from lancedb.index import FTS
from lancedb.index import BTree, FTS, IvfPq
from lancedb.table import AsyncTable, Table
@@ -99,6 +99,86 @@ async def test_async_hybrid_query_filters(table: AsyncTable):
assert result["text"].to_pylist() == ["cat", "b"]
@pytest.mark.asyncio
async def test_hybrid_query_with_stale_fixed_size_binary_prefilter(
tmpdir_factory,
):
tmp_path = str(tmpdir_factory.mktemp("stale_scalar_prefilter"))
db = await lancedb.connect_async(tmp_path)
def fixed_size_binary(value: int) -> bytes:
return value.to_bytes(16, byteorder="big")
num_rows = 1000
data = pa.table(
{
"space_id": pa.array(
[fixed_size_binary(i) for i in range(num_rows)],
type=pa.binary(16),
),
"text": ["book"] * num_rows,
"vector": pa.array(
[[float(i), float(i)] for i in range(num_rows)],
type=pa.list_(pa.float32(), 2),
),
}
)
table = await db.create_table("test", data)
await table.create_index(
"vector", config=IvfPq(num_partitions=4, num_sub_vectors=2)
)
await table.create_index("space_id", config=BTree())
await table.create_index("text", config=FTS(with_position=False))
# Advance the search indices without advancing the scalar index. This is the
# state that previously let hybrid search use an incomplete scalar prefilter.
await table.add(data)
lance_dataset = await table.to_lance()
lance_dataset.optimize.optimize_indices(index_names=["vector_idx", "text_idx"])
await table.checkout_latest()
scalar_stats = await table.index_stats("space_id_idx")
assert scalar_stats is not None
assert scalar_stats.num_indexed_rows == num_rows
assert scalar_stats.num_unindexed_rows == num_rows
for index_name in ["vector_idx", "text_idx"]:
search_stats = await table.index_stats(index_name)
assert search_stats is not None
assert search_stats.num_indexed_rows == num_rows * 2
assert search_stats.num_unindexed_rows == 0
matching_ids = [5, 10, 15, 20, 25, 30]
literals = [
f"arrow_cast(0x{fixed_size_binary(i).hex()}, 'FixedSizeBinary(16)')"
for i in matching_ids
]
predicate = f"space_id IN ({', '.join(literals)})"
expected_ids = sorted(fixed_size_binary(i) for i in matching_ids for _ in range(2))
vector_query = (
table.query().where(predicate).nearest_to([5.0, 5.0]).limit(num_rows * 2)
)
vector_results = await vector_query.to_arrow()
assert sorted(vector_results["space_id"].to_pylist()) == expected_ids
fts_query = (
table.query().where(predicate).nearest_to_text("book").limit(num_rows * 2)
)
fts_results = await fts_query.to_arrow()
assert sorted(fts_results["space_id"].to_pylist()) == expected_ids
hybrid_results = await (
table.query()
.where(predicate)
.nearest_to([5.0, 5.0])
.nearest_to_text("book")
.limit(num_rows * 2)
.to_arrow()
)
assert sorted(hybrid_results["space_id"].to_pylist()) == expected_ids
@pytest.mark.asyncio
async def test_async_hybrid_query_default_limit(table: AsyncTable):
# add 10 new rows
@@ -123,6 +203,31 @@ async def test_async_hybrid_query_default_limit(table: AsyncTable):
assert texts.count("a") == 1
def test_hybrid_query_offset(sync_table: Table):
# The offset window of a hybrid query must be a suffix of the same query
# run without an offset -- it must not be silently ignored.
full = (
sync_table.search(query_type="hybrid")
.vector([0.0, 0.4])
.text("dog")
.limit(4)
.with_row_id(True)
.to_arrow()
)
assert len(full) == 4
offset_result = (
sync_table.search(query_type="hybrid")
.vector([0.0, 0.4])
.text("dog")
.offset(2)
.limit(2)
.with_row_id(True)
.to_arrow()
)
assert offset_result["_rowid"].to_pylist() == full["_rowid"].to_pylist()[2:]
def test_hybrid_query_minimum_nprobes_zero_raises(sync_table: Table):
# minimum_nprobes(0) must raise the same validation error a plain vector
# query raises, not silently no-op because 0 is falsy.
python/python/tests/test_import.py
+33
View File
@@ -0,0 +1,33 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import re
import shutil
import subprocess
import sys
import lancedb._lancedb as _lancedb
import pytest
@pytest.mark.skipif(sys.platform != "linux", reason="ldd is Linux-specific")
def test_native_extension_does_not_link_openssl():
"""OpenSSL-linked wheels abort when imported on RHEL hosts in FIPS mode."""
ldd = shutil.which("ldd")
if ldd is None:
pytest.skip("ldd is not installed")
result = subprocess.run(
[ldd, _lancedb.__file__],
check=True,
capture_output=True,
text=True,
)
openssl_libraries = re.findall(
r"^\s*(lib(?:crypto|ssl)\S*)\s+=>", result.stdout, flags=re.MULTILINE
)
assert not openssl_libraries, (
"the LanceDB native extension must use rustls instead of linking OpenSSL: "
f"{openssl_libraries}"
)
python/python/tests/test_index.py
+26 -1
View File
@@ -88,7 +88,7 @@ async def binary_table(db_async):
async def test_create_index_async_returns_done_job(some_table: AsyncTable):
job = await some_table.create_index_async("id", config=BTree())
assert job.id is None
await job.wait()
assert await job.wait() is None
assert len(await some_table.list_indices()) == 1
await job.cancel()
@@ -372,6 +372,31 @@ async def test_create_vector_index(some_table: AsyncTable):
assert stats.num_indices == 1
@pytest.mark.asyncio
async def test_create_ivf_index_reports_unsplittable_partitions(db_async):
dim = 8
num_partitions = 300 # More than 256 selects hierarchical k-means.
base_vectors = [[float(row == column) for column in range(dim)] for row in range(5)]
vectors = pa.array(base_vectors * 200, pa.list_(pa.float32(), dim))
table = await db_async.create_table(
"unsplittable_partitions",
pa.table({"vector": vectors}),
)
error_pattern = (
rf"Cannot create {num_partitions} IVF partitions: k-means could only form"
)
with pytest.raises(RuntimeError, match=error_pattern):
await table.create_index(
"vector",
config=IvfFlat(
distance_type="dot",
num_partitions=num_partitions,
max_iterations=10,
),
)
@pytest.mark.asyncio
async def test_create_4bit_ivfpq_index(some_table: AsyncTable):
# Can create
python/python/tests/test_lsm_write_spec.py
+16 -4
View File
@@ -21,6 +21,11 @@ SCHEMA = pa.schema(
)
def test_lsm_write_spec_module_metadata():
assert lancedb.LsmWriteSpec is LsmWriteSpec
assert LsmWriteSpec.__module__ == "lancedb._lancedb"
def _batch(ids, vs):
return pa.RecordBatch.from_arrays(
[pa.array(ids, type=pa.utf8()), pa.array(vs, type=pa.int32())],
@@ -83,7 +88,9 @@ def test_lsm_write_spec_repr():
assert s.spec_type == "bucket"
assert s.column == "id"
assert s.num_buckets == 4
assert s.maintained_indexes == []
# A fresh spec defers its maintained set to install time.
assert s.maintained_indexes is None
assert s.with_maintained_indexes([]).maintained_indexes == []
assert "bucket" in repr(s)
assert "id" in repr(s)
assert "4" in repr(s)
@@ -169,18 +176,23 @@ def test_get_lsm_write_spec(tmp_path):
table.unset_lsm_write_spec()
assert table.get_lsm_write_spec() is None
# Identity round-trips (column recovered from the schema).
# Identity round-trips (column recovered from the schema). Leaving the
# maintained set to be inferred picks up the index on the table, so the
# spec reads back naming it rather than as "infer".
table.set_lsm_write_spec(LsmWriteSpec.identity("id"))
spec = table.get_lsm_write_spec()
assert spec.spec_type == "identity"
assert spec.column == "id"
assert spec.maintained_indexes == [idx_name]
table.unset_lsm_write_spec()
# Unsharded round-trips (no routing column).
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
# Unsharded round-trips (no routing column). Opting out is distinct from
# the inferred default.
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
spec = table.get_lsm_write_spec()
assert spec.spec_type == "unsharded"
assert spec.column is None
assert spec.maintained_indexes == []
@pytest.mark.asyncio
python/python/tests/test_materialized_views.py
+268
View File
@@ -0,0 +1,268 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import lancedb
import pytest
from lancedb.materialized_view import MaterializedViewDefinition
STABLE_ROW_IDS = {"new_table_enable_stable_row_ids": "true"}
def make_db(tmp_path):
db = lancedb.connect(tmp_path, storage_options=STABLE_ROW_IDS)
db.create_table(
"people",
[
{"name": "ada", "age": 36},
{"name": "kid", "age": 7},
{"name": "grace", "age": 85},
],
)
return db
def test_create_refresh_and_query(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view(
"adults",
"people",
select=["name", ("shout", "upper(name)")],
where="age >= 18",
)
assert view.name == "adults"
assert view.table.count_rows() == 0
result = view.refresh()
assert result.mode == "rebuild"
assert result.rows_written == 2
rows = view.table.search().to_list()
assert sorted(row["shout"] for row in rows) == ["ADA", "GRACE"]
def test_definition_round_trips(tmp_path):
db = make_db(tmp_path)
db.create_materialized_view("adults", "people", where="age >= 18")
view = db.open_materialized_view("adults")
assert view.definition == MaterializedViewDefinition(
source_table="people",
projections=[("name", "`name`"), ("age", "`age`")],
filter="age >= 18",
inputs=["age", "name"],
)
def test_incremental_refresh_after_append(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view("copy", "people")
view.refresh()
db.open_table("people").add([{"name": "alan", "age": 41}])
result = view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
assert view.table.count_rows() == 4
assert view.refresh().mode == "no_op"
def test_incremental_refresh_after_update(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view("copy", "people")
view.refresh()
db.open_table("people").update(where="name = 'kid'", values={"age": 8})
result = view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
rows = view.table.search().to_list()
assert sorted(row["age"] for row in rows) == [8, 36, 85]
def test_legacy_storage_source_update_rebuilds(tmp_path):
db = lancedb.connect(
tmp_path,
storage_options={**STABLE_ROW_IDS, "new_table_data_storage_version": "legacy"},
)
db.create_table("people", [{"name": "ada", "age": 36}, {"name": "kid", "age": 7}])
view = db.create_materialized_view("copy", "people")
view.refresh()
db.open_table("people").update(where="name = 'kid'", values={"age": 8})
result = view.refresh()
assert result.mode == "rebuild"
rows = view.table.search().to_list()
assert sorted(row["age"] for row in rows) == [8, 36]
def test_list_and_not_a_view(tmp_path):
db = make_db(tmp_path)
db.create_materialized_view("adults", "people", where="age >= 18")
assert db.list_materialized_views() == ["adults"]
with pytest.raises(ValueError, match="not a materialized view"):
db.open_materialized_view("people")
def test_invalid_expression_fails_at_create(tmp_path):
db = make_db(tmp_path)
with pytest.raises(Exception, match="missing"):
db.create_materialized_view("bad", "people", select=[("x", "missing + 1")])
assert "bad" not in db.list_tables().tables
@pytest.mark.asyncio
async def test_async_create_refresh_and_open(tmp_path):
db = await lancedb.connect_async(tmp_path, storage_options=STABLE_ROW_IDS)
await db.create_table("people", [{"name": "ada", "age": 36}])
view = await db.create_materialized_view(
"shouts", "people", select=[("shout", "upper(name)")]
)
result = await view.refresh()
assert result.mode == "rebuild"
assert result.rows_written == 1
reopened = await db.open_materialized_view("shouts")
definition = await reopened.definition()
assert definition.projections == [("shout", "upper(name)")]
assert await db.list_materialized_views() == ["shouts"]
@pytest.mark.asyncio
async def test_async_incremental(tmp_path):
db = await lancedb.connect_async(tmp_path, storage_options=STABLE_ROW_IDS)
await db.create_table("people", [{"name": "ada", "age": 36}])
view = await db.create_materialized_view("copy", "people")
await view.refresh()
table = await db.open_table("people")
await table.add([{"name": "alan", "age": 41}])
result = await view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
def test_source_requires_stable_row_ids(tmp_path):
db = lancedb.connect(tmp_path)
db.create_table("plain", [{"x": 1}])
with pytest.raises(Exception, match="stable row ids"):
db.create_materialized_view("v", "plain")
def test_bare_select_names_are_quoted(tmp_path):
db = lancedb.connect(tmp_path, storage_options=STABLE_ROW_IDS)
db.create_table("odd_names", [{"order item": "widget", "select": 2}])
view = db.create_materialized_view(
"quoted", "odd_names", select=["order item", "select"]
)
result = view.refresh()
assert result.rows_written == 1
rows = view.table.search().to_list()
assert rows[0]["order item"] == "widget"
assert rows[0]["select"] == 2
@pytest.mark.asyncio
async def test_async_remote_is_refused_without_network():
db = await lancedb.connect_async(
"db://nowhere", api_key="sk_test", region="us-east-1"
)
with pytest.raises(NotImplementedError, match="local"):
await db.create_materialized_view("v", "src")
with pytest.raises(NotImplementedError, match="local"):
await db.open_materialized_view("v")
with pytest.raises(NotImplementedError, match="local"):
await db.list_materialized_views()
def test_scalar_select_is_one_column(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view("just_name", "people", select="name")
view.refresh()
rows = view.table.search().to_list()
assert set(rows[0]) - {"__source_row_id"} == {"name"}
assert sorted(row["name"] for row in rows) == ["ada", "grace", "kid"]
@pytest.mark.asyncio
async def test_async_scalar_select_is_one_column(tmp_path):
db = await lancedb.connect_async(tmp_path, storage_options=STABLE_ROW_IDS)
await db.create_table("people", [{"name": "ada", "age": 36}])
view = await db.create_materialized_view("just_name", "people", select="name")
await view.refresh()
rows = await view.table.query().to_list()
assert set(rows[0]) - {"__source_row_id"} == {"name"}
def test_limit_above_i64_max_is_refused(tmp_path):
db = make_db(tmp_path)
with pytest.raises(ValueError, match="exceeds the maximum"):
db.create_materialized_view("too_big", "people", limit=2**63)
# The boundary is fine, and zero still means an empty view.
db.create_materialized_view("at_max", "people", limit=2**63 - 1)
empty = db.create_materialized_view("none", "people", limit=0)
empty.refresh()
assert empty.table.count_rows() == 0
def _namespace_db(tmp_path):
return lancedb.connect_namespace(
"dir",
{"root": str(tmp_path)},
storage_options=STABLE_ROW_IDS,
)
def test_namespace_connection_materialized_views(tmp_path):
db = _namespace_db(tmp_path)
db.create_table(
"people",
[{"name": "ada", "age": 36}, {"name": "kid", "age": 7}],
storage_options=STABLE_ROW_IDS,
)
view = db.create_materialized_view("adults", "people", where="age >= 18")
view.refresh()
assert view.table.count_rows() == 1
assert db.list_materialized_views() == ["adults"]
reopened = db.open_materialized_view("adults")
assert reopened.definition.source_table == "people"
with pytest.raises(ValueError, match="not a materialized view"):
db.open_materialized_view("people")
@pytest.mark.asyncio
async def test_async_namespace_connection_materialized_views(tmp_path):
db = lancedb.connect_namespace_async(
"dir",
{"root": str(tmp_path)},
storage_options=STABLE_ROW_IDS,
)
await db.create_table(
"people",
[{"name": "ada", "age": 36}, {"name": "kid", "age": 7}],
storage_options=STABLE_ROW_IDS,
)
view = await db.create_materialized_view("adults", "people", where="age >= 18")
await view.refresh()
assert await view.table.count_rows() == 1
assert await db.list_materialized_views() == ["adults"]
reopened = await db.open_materialized_view("adults")
assert (await reopened.definition()).source_table == "people"
# The view's table came through the namespace, not straight from the
# inner connection: a bare inner table carries no namespace context, so
# its pushdown routing differs from a table the namespace opened.
through_namespace = await db.open_table("adults")
for handle in (view.table, reopened.table):
assert (
handle._route_pushdown_to_rust == through_namespace._route_pushdown_to_rust
)
assert handle._namespace_path == through_namespace._namespace_path
python/python/tests/test_merge_insert_lsm.py
+2 -2
View File
@@ -544,7 +544,7 @@ def test_lsm_read_fts_unmaintained_index_errors(tmp_path):
table.create_index("text", config=FTS())
# No maintained indexes: the active memtable FTS arm cannot serve un-compacted
# docs, so the search would silently omit them — reject instead.
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
with pytest.raises(Exception, match="maintained"):
table.search("fox", query_type="fts", fts_columns="text").to_arrow()
@@ -631,7 +631,7 @@ def test_lsm_read_vector_unmaintained_index_errors(tmp_path):
)
# Spec with NO maintained indexes: the base vector index's catch-up is untracked,
# so the scanner rejects rather than risk dropping compacted-but-unindexed rows.
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
with pytest.raises(Exception, match="maintained"):
table.search([1.0] * VECTOR_DIM).to_arrow()
python/python/tests/test_package_metadata.py
+42
View File
@@ -0,0 +1,42 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import importlib
import re
import sys
from pathlib import Path
import pytest
def test_pyo3_abi_matches_minimum_supported_python():
project_dir = Path(__file__).parents[2]
pyproject = (project_dir / "pyproject.toml").read_text()
cargo_manifest = (project_dir / "Cargo.toml").read_text()
minimum_python = re.search(
r'^requires-python\s*=\s*">=(\d+)\.(\d+)"$', pyproject, re.MULTILINE
)
assert minimum_python is not None
major, minor = minimum_python.groups()
expected_abi = f"abi3-py{major}{minor}"
configured_abis = re.findall(r'"(abi3-py\d+)"', cargo_manifest)
assert configured_abis == [expected_abi, expected_abi], (
"the pyo3 runtime and build ABI features must both match requires-python"
)
@pytest.mark.skipif(sys.platform != "win32", reason="Windows wheel regression test")
def test_windows_wheel_tag_and_native_import():
project_dir = Path(__file__).parents[2]
wheels = list((project_dir.parent / "target" / "wheels").glob("lancedb-*.whl"))
if not wheels:
pytest.skip("no wheel artifact is available in this development environment")
assert len(wheels) == 1
assert wheels[0].name.endswith("-cp310-abi3-win_amd64.whl")
native_module = importlib.import_module("lancedb._lancedb")
assert Path(native_module.__file__).suffix == ".pyd"
python/python/tests/test_permutation.py
+104
View File
@@ -6,7 +6,13 @@ import math
import pytest
from lancedb import DBConnection, Table, connect
from lancedb.background_loop import LOOP
from lancedb.permutation import Permutation, Permutations, permutation_builder
from utils import (
MockPermutationServer,
assert_server_safe_row_id_requests,
mock_remote_table,
)
def test_split_random_ratios(mem_db):
@@ -31,6 +37,50 @@ def test_split_random_ratios(mem_db):
assert 65 <= split_1_count <= 75 # ~70% ± tolerance
def test_execute_does_not_reenter_background_loop(tmp_path, monkeypatch):
import threading
db = connect(tmp_path)
tbl = db.create_table("test_table", pa.table({"x": range(10)}))
original_run = LOOP.run
def fail_on_reentry(future):
assert threading.current_thread() is not LOOP.thread
return original_run(future)
monkeypatch.setattr(LOOP, "run", fail_on_reentry)
permutation_tbl = permutation_builder(tbl).execute()
assert permutation_tbl.count_rows() == 10
assert permutation_tbl._conn.read_consistency_interval is None
def test_pickled_permutation_reads_pinned_version(tmp_path):
"""An unpickled copy must still read the pinned version, which also covers the
version surviving the ``to_arrow()`` round trip in ``__getstate__``."""
import pickle
db = connect(tmp_path)
tbl = db.create_table("base", pa.table({"idx": range(20)}))
permutation_tbl = permutation_builder(tbl).execute()
perm = Permutation.from_tables(tbl, permutation_tbl)
payload = pickle.dumps(perm)
# Compact so the stored row addresses no longer describe these rows at latest.
tbl.delete("true")
tbl.optimize()
assert tbl.count_rows() == 0
# Unpickle after the mutation: __setstate__ reopens at latest, so this only
# passes if the recorded version is applied on reopen.
restored = pickle.loads(payload)
assert len(restored) == 20
rows = restored.__getitems__(list(range(20)))
assert sorted(row["idx"] for row in rows) == list(range(20))
def test_split_random_counts(mem_db):
"""Test random splitting with absolute counts."""
tbl = mem_db.create_table(
@@ -1194,3 +1244,57 @@ def test_remove_rowid_after_select(some_permutation: Permutation):
perm_without_rowid = perm_with_rowid.remove_columns(["_rowid"])
assert "_rowid" not in perm_without_rowid.column_names
assert perm_without_rowid.column_names == ["id"]
def test_permutation_is_stable_when_remote_scan_order_varies():
"""Splits are assigned by scan position, and every rank builds its own
permutation, so two ranks seeing different scan orders must still agree."""
server = MockPermutationServer(num_rows=16, vary_scan_order=True)
def split_of_each_row(permutation_tbl):
# Sequential splits are assigned by position, so a reversed scan would put
# the last rows in split 0. Compare the mapping rather than the table order,
# which the split-id sort does not pin down.
rows = permutation_tbl.search(None).to_arrow().to_pydict()
return dict(zip(rows["row_id"], rows["split_id"]))
with mock_remote_table(server) as table:
first = split_of_each_row(
permutation_builder(table).split_sequential(fixed=2).execute()
)
second = split_of_each_row(
permutation_builder(table).split_sequential(fixed=2).execute()
)
assert server.scan_calls == 2, "both builds must have scanned"
assert first == second
assert first[0] == 0 and first[server.num_rows - 1] == 1, first
def test_permutation_over_remote_table():
"""The permutation API accepts a remote table, addressing rows by `_rowid` just
as `take_row_ids` does. Also pins the request shapes sent to the server.
"""
server = MockPermutationServer()
with mock_remote_table(server) as table:
permutation_tbl = permutation_builder(table).split_sequential(fixed=2).execute()
assert permutation_tbl.count_rows() == server.num_rows
permutation = Permutation.from_tables(table, permutation_tbl, 0)
assert permutation.num_rows == server.num_rows // 2
# Compare against the permutation's own order; the split-id sort is not stable.
rows = permutation_tbl.search(None).to_arrow().to_pydict()
split0 = [
row_id
for row_id, split in zip(rows["row_id"], rows["split_id"])
if not split
]
# The mock table's `id` equals its `_rowid`.
assert permutation.take_offsets([2, 0]) == [
{"id": split0[2]},
{"id": split0[0]},
]
assert_server_safe_row_id_requests(server)
python/python/tests/test_pydantic.py
+17 -16
View File
@@ -10,11 +10,10 @@ import pyarrow as pa
import pydantic
import pytest
from lancedb.pydantic import (
PYDANTIC_VERSION,
LanceModel,
MultiVector,
Vector,
pydantic_to_schema,
MultiVector,
)
from pydantic import BaseModel
from pydantic import Field
@@ -415,22 +414,27 @@ def test_nullable_vector():
assert schema == pa.schema([pa.field("vec", pa.list_(pa.float32(), 16), True)])
def test_bare_vector_raises_clear_error():
namespace = {
"__name__": "test_model_without_pyarrow",
"LanceModel": LanceModel,
"Vector": Vector,
}
with pytest.raises(TypeError, match=r"Vector must be parameterized.*Vector\(128\)"):
exec("class TestModel(LanceModel):\n vector: Vector", namespace)
def test_fixed_size_list_field():
class TestModel(pydantic.BaseModel):
vec: Vector(16)
li: List[int]
data = TestModel(vec=list(range(16)), li=[1, 2, 3])
if PYDANTIC_VERSION.major >= 2:
assert json.loads(data.model_dump_json()) == {
"vec": list(range(16)),
"li": [1, 2, 3],
}
else:
assert data.dict() == {
"vec": list(range(16)),
"li": [1, 2, 3],
}
assert json.loads(data.model_dump_json()) == {
"vec": list(range(16)),
"li": [1, 2, 3],
}
schema = pydantic_to_schema(TestModel)
assert schema == pa.schema(
@@ -440,10 +444,7 @@ def test_fixed_size_list_field():
]
)
if PYDANTIC_VERSION.major >= 2:
json_schema = TestModel.model_json_schema()
else:
json_schema = TestModel.schema()
json_schema = TestModel.model_json_schema()
assert json_schema == {
"properties": {
python/python/tests/test_query.py
+9
View File
@@ -570,6 +570,15 @@ def test_query_builder(table):
assert all(np.array(rs[0]["vector"]) == [1, 2])
def test_query_multiple_vectors(table):
results = table.search([np.array([1, 2]), np.array([4, 5])]).limit(1).to_list()
assert len(results) == 2
results_by_query = {result["query_index"]: result for result in results}
assert results_by_query[0]["id"] == 1
assert results_by_query[1]["id"] == 2
def test_with_row_id(table: lancedb.table.Table):
rs = table.search().with_row_id(True).to_arrow()
assert "_rowid" in rs.column_names
python/python/tests/test_remote_db.py
+214 -10
View File
@@ -35,6 +35,12 @@ def make_mock_http_handler(handler):
return MockLanceDBHandler
@pytest.mark.parametrize("db_name", ["a" * 64, "invalid..database"])
def test_connect_rejects_invalid_cloud_dns_hostname(db_name):
with pytest.raises(ValueError, match="DNS labels must contain 1 to 63 bytes"):
lancedb.connect(f"db://{db_name}", api_key="fake")
@contextlib.contextmanager
def mock_lancedb_connection(handler):
with http.server.HTTPServer(
@@ -236,8 +242,8 @@ def test_remote_table_branches_sync():
table.branches.delete("exp")
def test_remote_table_branch_merge_defaults_to_execute():
merge_bodies = []
def test_remote_table_cherry_pick_defaults_to_execute():
cherry_pick_bodies = []
diff = {
"fromBranch": "exp",
"parentVersion": 1,
@@ -259,8 +265,7 @@ def test_remote_table_branch_merge_defaults_to_execute():
"changedColumns": [],
"addedIndexes": [],
"removedIndexes": [],
"mergeable": True,
"mergeBlockers": [],
"errors": [],
}
def handler(request):
@@ -270,11 +275,11 @@ def test_remote_table_branch_merge_defaults_to_execute():
else:
content_len = int(request.headers.get("Content-Length"))
request_body = json.loads(request.rfile.read(content_len))
merge_bodies.append(request_body)
cherry_pick_bodies.append(request_body)
dry_run = request_body["dry_run"]
status = 200 if dry_run else 409
body = {
"status": "ready" if dry_run else "rejected",
"status": "ready" if dry_run else "failed",
"diff": diff,
"preview": {"promotedColumns": []},
}
@@ -286,10 +291,10 @@ def test_remote_table_branch_merge_defaults_to_execute():
with mock_lancedb_connection(handler) as db:
branches = db.open_table("test").branches
assert branches.merge("exp")["status"] == "rejected"
assert branches.merge("exp", dry_run=True)["status"] == "ready"
assert branches.cherry_pick("exp")["status"] == "failed"
assert branches.cherry_pick("exp", dry_run=True)["status"] == "ready"
assert merge_bodies == [
assert cherry_pick_bodies == [
{"from_branch": "exp", "dry_run": False},
{"from_branch": "exp", "dry_run": True},
]
@@ -870,11 +875,85 @@ def test_remote_create_index_async_returns_job():
table = db.create_table("test", [{"id": 1}])
job = table.create_index_async("id", config=BTree())
assert job.id == "job-1"
job.wait(timeout=timedelta(seconds=30))
assert job.wait(timeout=timedelta(seconds=30)) is None
assert len(describe_calls) == 2
job.cancel()
def test_remote_refresh_async_returns_typed_terminal_result():
terminal_result = {
"rows_assigned": 12,
"rows_failed": 0,
"rows_remaining": 0,
"source_version": 7,
"published_version": 8,
}
def handler(request):
content_len = int(request.headers.get("Content-Length", 0))
body = request.rfile.read(content_len) if content_len > 0 else b""
if request.path == "/v1/table/test/backfill_column":
assert json.loads(body)["column"] == "derived"
request.send_response(202)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"job_id": "refresh-1"}')
elif request.path == "/v1/jobs/describe":
assert json.loads(body)["job_id"] == "refresh-1"
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(
json.dumps(
{
"job_id": "refresh-1",
"job_type": "function_refresh",
"job_state": "DONE",
"result": terminal_result,
}
).encode()
)
elif request.path == "/v1/table/test/create/?mode=create":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b"{}")
elif request.path == "/v1/table/test/describe/":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(
json.dumps(
{
"version": 1,
"schema": {
"fields": [
{
"name": "id",
"type": {"type": "int64"},
"nullable": False,
}
]
},
}
).encode()
)
else:
request.send_response(404)
request.end_headers()
with mock_lancedb_connection(handler) as db:
table = db.create_table("test", [{"id": 1}])
job = table.refresh_column_async("derived")
assert job.id == "refresh-1"
result = job.wait(timeout=timedelta(seconds=30))
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.model_dump() == terminal_result
assert result.rows_filled == 12
assert result.version == 8
def test_remote_job_wait_raises_on_failure():
from lancedb.exceptions import JobFailedError
from lancedb.index import BTree
@@ -1127,6 +1206,131 @@ def test_stats():
assert res == stats
@contextlib.contextmanager
def lsm_test_table(lsm_handler):
"""A remote table whose LSM routes are served by ``lsm_handler``.
``lsm_handler(request, route)`` is called for ``/v1/table/test/<route>/``
where route is one of flush_lsm, compact_lsm, get_lsm_stats, and is
responsible for writing the response.
"""
routes = ("flush_lsm", "compact_lsm", "get_lsm_stats")
def handler(request):
match = re.fullmatch(r"/v1/table/test/(\w+)/", request.path)
route = match.group(1) if match else None
if route in routes:
lsm_handler(request, route)
elif route == "describe":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
else:
request.send_response(404)
request.end_headers()
with mock_lancedb_connection(handler) as db:
yield db.open_table("test")
def read_json_body(request):
content_len = int(request.headers.get("Content-Length"))
return json.loads(request.rfile.read(content_len))
def send_json(request, payload, status=200):
request.send_response(status)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(json.dumps(payload).encode())
def test_get_lsm_stats_sync():
"""The sync wrapper round-trips the server payload into a dict."""
bucket = {
"shard_id": "b0",
"status": "Active",
"writer_epoch": 3,
"manifest_version": 12,
"current_generation": 6,
"replay_after_wal_entry_position": 40,
"wal_entry_position_last_seen": 42,
"generations": [{"generation": 5, "bytes": 1024, "rows": 7}],
"compacting": False,
"memtables": [
{
"generation": 6,
"rows": 2,
"bytes": 64,
"batches": 1,
"indexes": ["vec_idx"],
}
],
}
seen_bodies = []
def lsm_handler(request, route):
assert route == "get_lsm_stats"
seen_bodies.append(read_json_body(request))
send_json(request, {"lsm_stats": {"buckets": [bucket]}})
with lsm_test_table(lsm_handler) as table:
assert table.get_lsm_stats() == {"buckets": [bucket]}
# Off by default, and forwarded when asked for.
assert seen_bodies == [{"include_generation_rows": False}]
table.get_lsm_stats(include_generation_rows=True)
assert seen_bodies[-1] == {"include_generation_rows": True}
def test_get_lsm_stats_sync_returns_none_when_lsm_disabled():
"""A null envelope means the LSM write path is not enabled, not an error."""
def lsm_handler(request, route):
send_json(request, {"lsm_stats": None})
with lsm_test_table(lsm_handler) as table:
assert table.get_lsm_stats() is None
def test_flush_and_compact_lsm_sync():
"""Both are one-shot POSTs answered 202 with no body."""
called = []
def lsm_handler(request, route):
called.append(route)
request.send_response(202)
request.end_headers()
with lsm_test_table(lsm_handler) as table:
assert table.flush_lsm() is None
assert table.compact_lsm() is None
assert called == ["flush_lsm", "compact_lsm"]
def test_checkpoint_lsm_sync():
"""Seal, read the watermark, and return once L0 holds nothing.
The convergence loop itself is covered in Rust; this pins the sync
binding to the endpoints it drives.
"""
called = []
def lsm_handler(request, route):
called.append(route)
if route == "get_lsm_stats":
# An empty L0 yields no target watermark, so the loop is done
# after the seal without ever polling compaction.
send_json(request, {"lsm_stats": {"buckets": []}})
else:
request.send_response(202)
request.end_headers()
with lsm_test_table(lsm_handler) as table:
assert table.checkpoint_lsm() is None
assert called == ["flush_lsm", "get_lsm_stats"]
@contextlib.contextmanager
def query_test_table(query_handler, *, server_version=Version("0.1.0")):
def handler(request):
python/python/tests/test_table.py
+430 -5
View File
@@ -2,10 +2,14 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import ctypes
import gc
import os
import sys
import threading
import warnings
import weakref
from concurrent.futures import ThreadPoolExecutor
from datetime import date, datetime, timedelta
from time import sleep
from typing import List
@@ -98,6 +102,30 @@ def test_basic(mem_db: DBConnection):
assert table.to_arrow() == expected_data
def test_search_preserves_nulls_from_sliced_arrow_table(mem_db: DBConnection):
data = pa.table(
{
"id": [0, 1, 2, 3, 4],
"score_cn": [None, 22, None, 5, 8],
"score_mt": [None, 42, None, 5, 8],
"vector": [
[20, 19, -1, -1],
[41, 38, 22, 42],
[10, 10, -1, -1],
[5, 5, 5, 5],
[8, 8, 8, 8],
],
}
).slice(1)
table = mem_db.create_table("sliced_nullable", data=data)
result = table.search([41, 38, 22, 42]).limit(1).to_arrow()
assert result["id"].to_pylist() == [1]
assert result["score_cn"].to_pylist() == [22]
assert result["score_mt"].to_pylist() == [42]
def test_table_to_pandas_default_matches_arrow(tmp_db: DBConnection):
pd = pytest.importorskip("pandas")
data = pa.table({"id": [1, 2], "text": ["one", "two"]})
@@ -434,6 +462,38 @@ def test_add(mem_db: DBConnection):
_add(table, schema)
def test_add_releases_arrow_buffers_without_gc(mem_db: DBConnection):
"""Regression test for https://github.com/lancedb/lancedb/issues/2512."""
schema = pa.schema([pa.field("x", pa.int64())])
table = mem_db.create_table("test_add_releases_arrow_buffers", schema=schema)
class BufferOwner:
def __init__(self, size: int):
self.memory = ctypes.create_string_buffer(size)
owner_refs = []
gc_was_enabled = gc.isenabled()
gc.disable()
try:
for _ in range(3):
size = 8 * 1024
owner = BufferOwner(size)
arrow_buffer = pa.foreign_buffer(
ctypes.addressof(owner.memory), size, owner
)
array = pa.Array.from_buffers(pa.int64(), 1024, [None, arrow_buffer])
batch = pa.RecordBatch.from_arrays([array], schema=schema)
owner_refs.append(weakref.ref(owner))
table.add(batch)
del batch, array, arrow_buffer, owner
assert all(owner_ref() is None for owner_ref in owner_refs)
finally:
if gc_was_enabled:
gc.enable()
def test_add_write_parallelism(mem_db: DBConnection):
schema = pa.schema([pa.field("id", pa.int64())])
table = mem_db.create_table("test", schema=schema)
@@ -869,6 +929,7 @@ def test_polars(mem_db: DBConnection):
# enter table to polars dataframe
result = table.to_polars()
assert isinstance(result, pl.LazyFrame)
assert np.allclose(result.collect()["vector"].to_list(), data["vector"])
# make sure filtering isn't broken
@@ -1406,7 +1467,7 @@ def test_create_index_async_returns_done_job(mem_db: DBConnection):
table = mem_db.create_table("job_test", [{"id": i} for i in range(10)])
job = table.create_index_async("id", config=BTree())
assert job.id is None
job.wait()
assert job.wait() is None
assert len(table.list_indices()) == 1
job.cancel()
@@ -1785,6 +1846,27 @@ def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection):
assert np.allclose(data["embedding"].to_pylist()[0], np.array([0.1] * 16))
def test_add_nullable_fixed_size_list_with_none(mem_db: DBConnection):
"""Regression test for issue #2340."""
table = mem_db.create_table(
"test_nullable_fixed_size_list",
schema=pa.schema(
[
pa.field("id", pa.string()),
pa.field("feature", pa.list_(pa.float32(), 256)),
pa.field("tags", pa.list_(pa.string())),
]
),
)
table.add([{"id": "1", "feature": None, "tags": ["tag1", "tag2"]}])
result = table.to_arrow()
assert result.to_pylist() == [
{"id": "1", "feature": None, "tags": ["tag1", "tag2"]}
]
def test_add_nullable_struct_with_none(mem_db: DBConnection):
"""Regression test for issue #2654: a nullable struct column whose
first batch contains only None values must not crash in
@@ -1824,6 +1906,33 @@ def test_add_nullable_struct_with_none(mem_db: DBConnection):
assert result.column("data").to_pylist() == [{"x": 1.0}, None]
def test_read_mostly_null_list_v2_2_page_boundary(tmp_path):
# Regression test for #3194. This row/value count crosses a v2.2 structural
# encoding page boundary where Lance 3.0.0 sliced repetition/definition
# levels by row offset and decoded child arrays at different lengths.
num_rows = 64_885
num_values = 217
list_type = pa.list_(pa.float32())
source = pa.table(
{
"id": np.arange(num_rows, dtype=np.int64),
"coords": pa.array(
[[1.0, 2.0, 3.0, 4.0]] * num_values + [None] * (num_rows - num_values),
type=list_type,
),
}
)
db = lancedb.connect(
tmp_path,
storage_options={"new_table_data_storage_version": "2.2"},
)
table = db.create_table("test_sparse_nullable_list", data=source)
result = table.search().select(["id", "coords"]).limit(num_rows).to_arrow()
assert result.equals(source)
def test_add_with_integer_embeddings_preserves_casting(mem_db: DBConnection):
class Schema(LanceModel):
text: str
@@ -2109,6 +2218,45 @@ def test_merge(tmp_db: DBConnection, tmp_path):
table.merge(other_dataset, left_on="id")
@pytest.mark.parametrize("storage_version", ["legacy", "stable"])
def test_search_after_merge(tmp_path, storage_version):
pytest.importorskip("lance")
pd = pytest.importorskip("pandas")
db = lancedb.connect(
tmp_path,
storage_options={"new_table_data_storage_version": storage_version},
)
rng = np.random.default_rng(42)
row_count = 512
vectors = rng.standard_normal((row_count, 8)).astype(np.float32)
table = db.create_table(
"search_after_merge",
data=pd.DataFrame(
{
"id": [str(i) for i in range(row_count)],
"vector": list(vectors),
}
),
)
table.create_index("vector", config=IvfPq(num_partitions=1, num_sub_vectors=2))
links = pd.DataFrame(
{
"id": [str(i) for i in range(row_count // 2)],
"link": [f"https://example.com/{i}" for i in range(row_count // 2)],
}
)
table.merge(links, left_on="id")
query = table.search(vectors[-1]).refine_factor(50).limit(10)
assert "ANN" in query.explain_plan(verbose=True)
result = query.to_arrow()
links_by_id = dict(zip(result["id"].to_pylist(), result["link"].to_pylist()))
assert links_by_id[str(row_count - 1)] is None
def test_delete(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
@@ -2124,6 +2272,27 @@ def test_delete(mem_db: DBConnection):
assert table.to_arrow()["id"].to_pylist() == [1]
def test_concurrent_deletes_are_thread_safe(mem_db: DBConnection):
num_workers = 8
table = mem_db.create_table(
"my_table", data=[{"id": row_id} for row_id in range(num_workers)]
)
barrier = threading.Barrier(num_workers)
def delete(row_id: int):
barrier.wait()
return table.delete(f"id = {row_id}")
with ThreadPoolExecutor(max_workers=num_workers) as pool:
results = list(pool.map(delete, range(num_workers)))
assert all(result.num_deleted_rows == 1 for result in results)
assert sorted(result.version for result in results) == list(
range(2, num_workers + 2)
)
assert table.count_rows() == 0
def test_delete_expr(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
@@ -2174,6 +2343,20 @@ def test_update(mem_db: DBConnection):
assert np.allclose(v, np.array([[1.2, 1.9], [1.1, 1.1]]))
def test_update_with_arrow_scalar(mem_db: DBConnection):
schema = pa.schema({"id": pa.int64(), "vector": pa.list_(pa.float32(), 4)})
table = mem_db.create_table("my_table", schema=schema)
table.add([{"id": 1, "vector": [1.0, 2.0, 3.0, 4.0]}])
value = table.search().select(["vector"]).limit(1).to_arrow()["vector"][0]
assert isinstance(value, pa.FixedSizeListScalar)
result = table.update(where="id == 1", values={"vector": value})
assert result.rows_updated == 1
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0, 3.0, 4.0]]
def test_update_types(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
@@ -2341,6 +2524,55 @@ def test_merge_insert(mem_db: DBConnection):
)
def test_merge_insert_nullable_pandas_into_pydantic_schema(mem_db: DBConnection):
# Regression test for https://github.com/lancedb/lancedb/issues/2366
pd = pytest.importorskip("pandas")
class Document(LanceModel):
id: int
title: str
content: str
table = mem_db.create_table("documents", schema=Document)
table.add(
pd.DataFrame(
{
"title": ["Old title", "Unchanged"],
"id": [2, 3],
"content": ["Old content", "Keep this"],
}
)
)
# Pandas produces nullable Arrow fields, in an order that differs from the
# non-nullable Pydantic schema. This is valid as long as the data has no nulls.
new_data = pd.DataFrame(
{
"title": ["Inserted", "Updated"],
"id": [1, 2],
"content": ["New row", "New content"],
}
)
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(new_data)
)
assert result.num_inserted_rows == 1
assert result.num_updated_rows == 1
expected = pa.Table.from_pylist(
[
{"id": 1, "title": "Inserted", "content": "New row"},
{"id": 2, "title": "Updated", "content": "New content"},
{"id": 3, "title": "Unchanged", "content": "Keep this"},
],
schema=Document.to_arrow_schema(),
)
assert table.to_arrow().sort_by("id") == expected
def test_merge_insert_by_source_delete_expr(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
@@ -2441,6 +2673,36 @@ def test_merge_insert_subschema(mem_db: DBConnection, data_format):
assert table.to_arrow().sort_by("id") == expected
def test_repeated_partial_merge_insert_with_scalar_index(mem_db: DBConnection):
def make_batch(start: int) -> pa.Table:
return pa.table(
{
"id": [f"id-{i:04}" for i in range(start, start + 100)],
"category": ["A"] * 100,
"value_a": [float(i) for i in range(start, start + 100)],
"value_b": [float(i) / 10 for i in range(100)],
}
)
table = mem_db.create_table("my_table", data=make_batch(0))
table.add(make_batch(100))
table.add(make_batch(200))
table.create_index("id", config=BTree())
ids = [f"id-{i:04}" for i in range(100, 200)]
for value in (999.0, 888.0):
result = (
table.merge_insert("id")
.when_matched_update_all()
.execute(pa.table({"id": ids, "value_a": [value] * 100}))
)
assert result.num_updated_rows == 100
actual = table.to_arrow().sort_by("id")
assert actual.num_rows == 300
assert actual["value_a"].to_pylist()[100:200] == [888.0] * 100
@pytest.mark.asyncio
async def test_merge_insert_async(mem_db_async: AsyncConnection):
data = pa.table({"a": [1, 2, 3], "b": ["a", "b", "c"]})
@@ -2510,6 +2772,56 @@ async def test_merge_insert_async(mem_db_async: AsyncConnection):
assert (await table.to_arrow()).sort_by("a") == expected
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
@pytest.mark.asyncio
async def test_merge_insert_encodes_json(mem_db_async: AsyncConnection):
json_type = pa.json_()
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", json_type)])
def json_table(rows):
json_values = pa.ExtensionArray.from_storage(
json_type,
pa.array([value for _, value in rows], type=json_type.storage_type),
)
return pa.Table.from_arrays(
[pa.array([row_id for row_id, _ in rows]), json_values], schema=schema
)
table = await mem_db_async.create_table("json_merge", schema=schema)
await table.add(json_table([("a", '{"k": 1}'), ("b", '{"k": 9}')]))
await (
table.merge_insert("id")
.when_matched_update_all()
.execute(json_table([("a", '{"k": 2}')]))
)
rows = sorted(await table.query().to_list(), key=lambda row: row["id"])
assert rows == [
{"id": "a", "j": '{"k":2}'},
{"id": "b", "j": '{"k":9}'},
]
filtered = await table.query().where("json_extract(j, '$.k') = '2'").to_list()
assert filtered == [{"id": "a", "j": '{"k":2}'}]
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
@pytest.mark.asyncio
async def test_add_sanitization_encodes_json(mem_db_async: AsyncConnection):
json_type = pa.json_()
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", json_type)])
json_values = pa.ExtensionArray.from_storage(
json_type, pa.array(['{"k": 3}'], type=json_type.storage_type)
)
data = pa.Table.from_arrays([pa.array(["c"]), json_values], schema=schema)
table = await mem_db_async.create_table("json_add", schema=schema)
await table.add(data, on_bad_vectors="fill")
rows = await table.query().where("json_extract(j, '$.k') = '3'").to_list()
assert rows == [{"id": "c", "j": '{"k":3}'}]
def test_create_with_embedding_function(mem_db: DBConnection):
class MyTable(LanceModel):
text: str
@@ -2537,15 +2849,40 @@ def test_create_with_embedding_function(mem_db: DBConnection):
assert actual == expected
def test_create_f16_table_from_arrow_data(mem_db: DBConnection):
dimension = 32
num_rows = 512
values = pa.array(
np.random.default_rng(42)
.standard_normal(num_rows * dimension)
.astype(np.float16)
)
df = pa.table(
{
"text": [f"s-{i}" for i in range(num_rows)],
"vector": pa.FixedSizeListArray.from_arrays(values, dimension),
}
)
table = mem_db.create_table("f16_tbl", data=df)
assert table.schema.field("vector").type == pa.list_(pa.float16(), dimension)
table.create_index(num_partitions=2, num_sub_vectors=2)
query = df["vector"][2].as_py()
expected = table.search(query).limit(2).to_arrow()
assert "s-2" in expected["text"].to_pylist()
def test_create_f16_table(mem_db: DBConnection):
class MyTable(LanceModel):
text: str
vector: Vector(32, value_type=pa.float16())
rng = np.random.default_rng(42)
df = pa.table(
{
"text": [f"s-{i}" for i in range(512)],
"vector": [np.random.randn(32).astype(np.float16) for _ in range(512)],
"vector": [rng.standard_normal(32).astype(np.float16) for _ in range(512)],
}
)
table = mem_db.create_table(
@@ -3426,7 +3763,8 @@ def test_stats(mem_db: DBConnection):
stats = table.stats()
print(f"{stats=}")
assert stats == {
"total_bytes": 60,
# Full on-disk size of the data file, footer and metadata included.
"total_bytes": 633,
"num_rows": 2,
"num_indices": 0,
"fragment_stats": {
@@ -3444,6 +3782,13 @@ def test_stats(mem_db: DBConnection):
},
}
# Index files count toward total_bytes too (only deletion files and
# manifests are excluded).
table.create_index("id", config=BTree())
stats_with_index = table.stats()
assert stats_with_index["num_indices"] == 1
assert stats_with_index["total_bytes"] > stats["total_bytes"]
def test_create_table_empty_list_with_schema(mem_db: DBConnection):
"""Test creating table with empty list data and schema
@@ -3467,8 +3812,8 @@ def test_create_table_empty_list_no_schema_error(mem_db: DBConnection):
mem_db.create_table("test_empty_no_schema", data=[])
def test_add_table_with_empty_embeddings(tmp_path):
"""Test exact scenario from issue #1968
def test_create_table_without_data_with_vector_schema(tmp_path):
"""Test exact scenario from issue #1968.
Regression test for issue #1968:
https://github.com/lancedb/lancedb/issues/1968
@@ -3480,6 +3825,9 @@ def test_add_table_with_empty_embeddings(tmp_path):
embedding: Vector(16)
table = db.create_table("test", schema=MySchema)
assert table.count_rows() == 0
assert table.schema == MySchema.to_arrow_schema()
table.add(
[{"text": "bar", "embedding": [0.1] * 16}],
on_bad_vectors="drop",
@@ -3556,3 +3904,80 @@ async def test_async_search_runs_embedding_on_dedicated_executor(
assert all(name.startswith("lancedb-embedding") for name in captured_threads), (
f"embedding ran off the dedicated executor: {captured_threads}"
)
def test_computed_column_declare_and_refresh(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("computed", [{"x": 1}, {"x": 2}])
table.add_columns(computed={"doubled": "x * 2"})
assert table.to_arrow()["doubled"].to_pylist() == [None, None]
result = table.refresh_column("doubled")
assert result.rows_filled == 2
assert sorted(table.to_arrow()["doubled"].to_pylist()) == [2, 4]
table.add([{"x": 5}])
assert table.refresh_column("doubled").rows_filled == 1
assert sorted(table.to_arrow()["doubled"].to_pylist()) == [2, 4, 10]
def test_computed_column_rejects_transforms_and_computed_together(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("computed_mixed", [{"x": 1}])
with pytest.raises(ValueError):
table.add_columns({"a": "x + 1"}, computed={"b": "x * 2"})
@pytest.mark.asyncio
async def test_computed_column_async(tmp_path):
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("computed_async", [{"x": 3}])
await table.add_columns(computed={"tripled": "x * 3"})
await table.refresh_column("tripled")
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
def test_refresh_column_async_returns_job(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("computed_job", [{"x": 1}, {"x": 2}])
table.add_columns(computed={"doubled": "x * 2"})
job = table.refresh_column_async("doubled")
assert job.id is None # in-process jobs have no server id
result = job.wait()
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.rows_assigned == 2
assert result.rows_failed == 0
assert result.rows_remaining == 0
assert result.source_version == 2
assert result.published_version == 3
assert job.status() == "finished"
assert sorted(table.to_arrow()["doubled"].to_pylist()) == [2, 4]
no_op = table.refresh_column_async("doubled").wait()
assert no_op.rows_assigned == 0
assert no_op.source_version == 3
assert no_op.published_version is None
# Bad input raises at the call, not through the job.
with pytest.raises(Exception, match="not a computed column"):
table.refresh_column_async("x")
@pytest.mark.asyncio
async def test_refresh_column_async_job_async_table(tmp_path):
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("computed_job_async", [{"x": 3}])
await table.add_columns(computed={"tripled": "x * 3"})
job = await table.refresh_column_async("tripled")
result = await job.wait()
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.rows_assigned == 1
assert result.source_version == 2
assert result.published_version == 3
assert await job.status() == "finished"
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
python/python/tests/test_voyageai_embeddings.py
+16
View File
@@ -75,6 +75,22 @@ class TestVoyageAIModelRegistration:
with pytest.raises(ValueError, match="not supported"):
func.ndims()
def test_voyage3_source_embeddings_use_text_api(self, mock_voyageai_client):
"""Regression test for text table data being sent to the multimodal API."""
mock_voyageai_client.tokenize.return_value = [["hello", "world"]]
mock_voyageai_client.embed.return_value.embeddings = [[0.1] * 1024]
registry = get_registry()
func = registry.get("voyageai").create(name="voyage-3")
embeddings = func.compute_source_embeddings("hello world")
assert embeddings == [[0.1] * 1024]
mock_voyageai_client.embed.assert_called_once_with(
texts=["hello world"], model="voyage-3", input_type="document"
)
mock_voyageai_client.multimodal_embed.assert_not_called()
@pytest.mark.parametrize(
"model_name",
[
python/python/tests/utils.py
+206
View File
@@ -1,7 +1,17 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import contextlib
import http.server
import json
import re
import threading
import lancedb
import pyarrow as pa
import pytest
ARROW_FILE_CONTENT_TYPE = "application/vnd.apache.arrow.file"
def exception_output(e_info: pytest.ExceptionInfo):
import traceback
@@ -9,3 +19,199 @@ def exception_output(e_info: pytest.ExceptionInfo):
# skip traceback part, since it's not worth checking in tests
lines = traceback.format_exception_only(e_info.type, e_info.value)
return "".join(lines).strip()
def parse_in_list(filter_sql: str) -> list[int]:
"""Pull the integers out of a `<col> IN (a, b, c)` predicate.
Scoped to the parenthesised list so a cast in the SQL adds no phantom values.
"""
match = re.search(r"\bIN\s*\(([^)]*)\)", filter_sql, re.IGNORECASE)
assert match is not None, f"expected an IN list, got: {filter_sql}"
return [int(m) for m in re.findall(r"-?\d+", match.group(1))]
def is_row_id_take(body) -> bool:
"""True when a query body fetches specific rows by row id."""
return "_rowid" in (body.get("filter") or "")
def arrow_file_bytes(table: pa.Table) -> bytes:
"""Serialize to the Arrow IPC *file* framing the /query/ route answers with."""
sink = pa.BufferOutputStream()
with pa.ipc.new_file(sink, table.schema) as writer:
writer.write_table(table)
return sink.getvalue().to_pybytes()
class MockPermutationServer:
"""A stand-in LanceDB server hosting one table whose ``id`` equals its ``_rowid``.
Records every ``/query/`` body so tests can assert on the request shapes sent to
the server, which is the part that has to stay compatible.
"""
def __init__(self, name="remote_data", num_rows=8, vary_scan_order=False):
self.name = name
self.num_rows = num_rows
self.query_bodies = []
# Stand in for a distributed scan that answers in no fixed order.
self.vary_scan_order = vary_scan_order
self.scan_calls = 0
def __call__(self, request):
path = request.path
if path == f"/v1/table/{self.name}/describe/":
return self._json(
request,
{
"version": 1,
"schema": {
"fields": [
{"name": "id", "type": {"type": "int64"}, "nullable": False}
]
},
},
)
if path == f"/v1/table/{self.name}/get_lsm_write_spec/":
self._read_body(request)
# Null spec: this table has no LSM write path.
return self._json(request, {"lsm_write_spec": None})
if path == f"/v1/table/{self.name}/count_rows/":
self._read_body(request)
return self._json(request, self.num_rows)
if path == f"/v1/table/{self.name}/query/":
return self._query(request, self._read_body(request))
# Drain first, so an unexpected route cannot desync a keep-alive connection.
self._read_body(request)
request.send_response(404)
request.end_headers()
@property
def scans(self):
"""Bodies of the permutation build scan: the row id column, nothing else."""
return [b for b in self.query_bodies if b.get("columns") == ["_rowid"]]
@property
def takes(self):
"""Bodies of the row-id takes the loader fetches batches with.
Keyed on `_rowid`, not "has a filter": the schema probe also has a predicate.
"""
return [b for b in self.query_bodies if is_row_id_take(b)]
@staticmethod
def _read_body(request):
content_len = int(request.headers.get("Content-Length") or 0)
return json.loads(request.rfile.read(content_len)) if content_len else {}
@staticmethod
def _json(request, payload):
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(json.dumps(payload).encode())
@staticmethod
def _arrow(request, table):
body = arrow_file_bytes(table)
request.send_response(200)
request.send_header("Content-Type", ARROW_FILE_CONTENT_TYPE)
request.send_header("Content-Length", str(len(body)))
request.end_headers()
request.wfile.write(body)
def _query(self, request, body):
self.query_bodies.append(body)
if is_row_id_take(body):
# A row-id take. Answer ascending, so tests prove the client reorders.
row_ids = sorted(parse_in_list(body["filter"]))
return self._arrow(
request,
pa.table(
{
"id": pa.array(row_ids, pa.int64()),
"_rowid": pa.array(row_ids, pa.uint64()),
}
),
)
if body.get("columns") == ["_rowid"]:
# The permutation build scan: row ids and nothing else.
row_ids = list(range(self.num_rows))
if self.vary_scan_order and self.scan_calls % 2:
row_ids.reverse()
self.scan_calls += 1
return self._arrow(
request,
pa.table({"_rowid": pa.array(row_ids, pa.uint64())}),
)
# The schema probe: filtered to nothing, so it carries schema and no rows.
return self._arrow(request, pa.table({"id": pa.array([], pa.int64())}))
def _make_handler(serve):
class MockLanceDBHandler(http.server.BaseHTTPRequestHandler):
def do_GET(self):
serve(self)
def do_POST(self):
serve(self)
def log_message(self, *args):
pass # keep pytest output readable
return MockLanceDBHandler
@contextlib.contextmanager
def mock_remote_table(server):
"""Run ``server`` on a local port and yield an open remote table against it.
Threading: the loader fans out fetch threads a single-threaded server would
serialize, hiding the prefetch overlap under test.
"""
with http.server.ThreadingHTTPServer(
("localhost", 0), _make_handler(server)
) as srv:
thread = threading.Thread(target=srv.serve_forever)
thread.start()
try:
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=f"http://localhost:{srv.server_address[1]}",
client_config={"timeout_config": {"connect_timeout": 5}},
)
yield db.open_table(server.name)
finally:
srv.shutdown()
thread.join()
def assert_server_safe_row_id_requests(server):
"""Assert the loader fetched rows by row id and bounded everything else.
`.get`, not `[...]`, so a dropped field reads as the assertion, not a KeyError.
"""
for body in server.takes:
# The fetch needs the row id back to restore the requested order.
assert body.get("with_row_id") is True, body
assert "_rowid" in body["filter"], body
# Only the one-off permutation scan may scan the whole table; the schema probe is
# built once per split per epoch. `k == 0` counts as unbounded: lance reads a zero
# limit as "no limit".
def is_unbounded(body):
if is_row_id_take(body):
return False
k = body.get("k")
return k is None or k == 0 or k > server.num_rows
unbounded = [b for b in server.query_bodies if is_unbounded(b)]
assert unbounded == server.scans, (
f"only the permutation scan may be unbounded, got {unbounded}"
)
python/python/type_tests/connect.py
+15
View File
@@ -0,0 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from typing import assert_type
import lancedb
from lancedb import AsyncConnection, DBConnection
def check_connect_type() -> None:
assert_type(lancedb.connect("memory://"), DBConnection)
async def check_connect_async_type() -> None:
assert_type(await lancedb.connect_async("memory://"), AsyncConnection)
python/src/connection.rs
+83
View File
@@ -333,6 +333,40 @@ impl Connection {
})
}
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None))]
pub fn create_materialized_view(
self_: PyRef<'_, Self>,
name: String,
source: String,
projections: Option<Vec<(String, String)>>,
filter: Option<String>,
limit: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.create_materialized_view(name, source);
if let Some(projections) = projections {
builder = builder.select(projections);
}
if let Some(filter) = filter {
builder = builder.only_if(filter);
}
if let Some(limit) = limit {
builder = builder.limit(limit);
}
let view = builder.execute().await.infer_error()?;
Ok(Table::new(view.table().clone()))
})
}
pub fn list_materialized_views(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let views = inner.list_materialized_views().await.infer_error()?;
Ok(views.into_iter().map(|view| view.name).collect::<Vec<_>>())
})
}
#[pyo3(signature = (name, namespace_path=None))]
pub fn drop_table(
self_: PyRef<'_, Self>,
@@ -346,6 +380,23 @@ impl Connection {
})
}
#[pyo3(signature = (name, namespace_path=None))]
pub fn drop_table_async(
self_: PyRef<'_, Self>,
name: String,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let ns_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner
.drop_table_async(name, &ns_path)
.await
.infer_error()
.map(crate::job::Job::new)
})
}
#[pyo3(signature = (namespace_path=None,))]
pub fn drop_all_tables(
self_: PyRef<'_, Self>,
@@ -546,6 +597,38 @@ impl Connection {
Ok(crate::job::Job::new(inner.job(job_id).infer_error()?))
}
pub fn create_function_async(
self_: PyRef<'_, Self>,
request_json: String,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let request = lancedb::function::FunctionRegistrationRequest::from_json(&request_json)
.infer_error()?;
future_into_py(self_.py(), async move {
inner
.create_function_async(request)
.await
.infer_error()
.map(crate::job::Job::new_typed)
})
}
pub fn get_function(
self_: PyRef<'_, Self>,
name: String,
version: String,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
inner
.get_function(name, version)
.await
.infer_error()?
.to_canonical_json()
.infer_error()
})
}
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
python/src/expr.rs
+20 -1
View File
@@ -191,8 +191,27 @@ pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult<PyExpr> {
}
// datetime.datetime is a subclass of datetime.date, so it must be checked first.
//
// Python's datetime.timestamp() treats *naive* datetimes as local wall time.
// PyArrow (and therefore Lance table storage) encodes naive timestamps as
// UTC wall-clock microseconds. Using .timestamp() for naive values therefore
// shifts the literal by the local UTC offset on non-UTC machines, so
// `col("ts") == lit(naive_dt)` fails against a table that holds the same
// naive value. Fix: treat naive datetimes as UTC wall clock (match Arrow);
// keep aware datetimes on the real .timestamp() path (correct epoch).
if let Ok(dt) = value.cast::<PyDateTime>() {
let ts: f64 = dt.call_method0("timestamp")?.extract()?;
let ts: f64 = if dt.getattr("tzinfo")?.is_none() {
// Force UTC interpretation of the naive wall clock.
let utc = pyo3::types::PyModule::import(value.py(), "datetime")?
.getattr("timezone")?
.getattr("utc")?;
let kwargs = pyo3::types::PyDict::new(value.py());
kwargs.set_item("tzinfo", utc)?;
let aware = dt.call_method("replace", (), Some(&kwargs))?;
aware.call_method0("timestamp")?.extract()?
} else {
dt.call_method0("timestamp")?.extract()?
};
let micros = (ts * 1_000_000.0).round() as i64;
return Ok(PyExpr(df_lit(ScalarValue::TimestampMicrosecond(
Some(micros),
python/src/index.rs
+58 -2
View File
@@ -42,7 +42,7 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
"Fm" => Ok(LanceDbIndex::Fm(FmIndexBuilder::default())),
"FTS" => {
let params = source.extract::<FtsParams>()?;
let inner_opts = FtsIndexBuilder::default()
let mut inner_opts = FtsIndexBuilder::default()
.base_tokenizer(params.base_tokenizer)
.language(&params.language)
.map_err(|_| {
@@ -61,6 +61,12 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
.ngram_max_length(params.ngram_max_length)
.ngram_prefix_only(params.prefix_only)
.custom_stop_words(params.custom_stop_words);
if let Some(memory_limit) = params.memory_limit {
inner_opts = inner_opts.memory_limit_mb(memory_limit);
}
if let Some(num_workers) = params.num_workers {
inner_opts = inner_opts.num_workers(num_workers);
}
let inner_opts = inner_opts
.block_size(params.block_size)
.map_err(|err| PyValueError::new_err(err.to_string()))?;
@@ -213,6 +219,8 @@ struct FtsParams {
ngram_max_length: u32,
prefix_only: bool,
block_size: usize,
memory_limit: Option<u64>,
num_workers: Option<usize>,
}
#[derive(FromPyObject)]
@@ -289,7 +297,7 @@ struct IvfHnswFlatParams {
target_partition_size: Option<u32>,
}
#[pyclass(get_all)]
#[pyclass(module = "lancedb._lancedb", get_all)]
/// A description of an index currently configured on a column
pub struct IndexConfig {
/// The type of the index
@@ -444,3 +452,51 @@ impl IndexConfig {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use pyo3::types::{PyDict, PyDictMethods};
use serde_json::json;
#[test]
fn fts_build_controls_are_forwarded() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(
c"class FTS:
with_position = True
base_tokenizer = 'simple'
language = 'English'
max_token_length = None
lower_case = True
stem = False
remove_stop_words = False
custom_stop_words = None
ascii_folding = False
ngram_min_length = 3
ngram_max_length = 3
prefix_only = False
block_size = 128
memory_limit = 2048
num_workers = 7
config = FTS()",
None,
Some(&locals),
)
.unwrap();
let config = locals.get_item("config").unwrap().unwrap();
let index = extract_index_params(&Some(config)).unwrap();
let LanceDbIndex::FTS(params) = index else {
panic!("expected FTS index parameters");
};
let training_json = params.to_training_json().unwrap();
assert_eq!(training_json.get("memory_limit"), Some(&json!(2048)));
assert_eq!(training_json.get("num_workers"), Some(&json!(7)));
});
}
}
python/src/job.rs
+20 -4
View File
@@ -5,18 +5,32 @@ use std::sync::Arc;
use crate::runtime::future_into_py;
use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods};
use serde::Serialize;
use crate::error::PythonErrorExt;
#[pyclass]
pub struct Job {
inner: Arc<lancedb::Job>,
inner: Arc<lancedb::Job<std::result::Result<Option<String>, String>>>,
}
impl Job {
pub(crate) fn new(inner: lancedb::Job) -> Self {
Self {
inner: Arc::new(inner),
inner: Arc::new(inner.map(|()| Ok(None))),
}
}
pub(crate) fn new_typed<T>(inner: lancedb::Job<T>) -> Self
where
T: Clone + Serialize + Send + Sync + 'static,
{
Self {
inner: Arc::new(inner.map(|result| {
serde_json::to_string(&result)
.map(Some)
.map_err(|error| format!("failed to serialize typed job result: {error}"))
})),
}
}
}
@@ -39,8 +53,10 @@ impl Job {
pub fn wait(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner.wait().await.infer_error()?;
Ok(())
let result = inner.wait().await.infer_error()?;
result
.map_err(|message| lancedb::Error::Runtime { message })
.infer_error()
})
}
python/src/lib.rs
+4 -1
View File
@@ -16,7 +16,8 @@ use query::{FTSQuery, HybridQuery, Query, VectorQuery};
use session::Session;
use table::{
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken,
LsmWriteSpec, MergeResult, PyBlobFile, Table, UpdateFieldMetadataResult, UpdateResult,
LsmWriteSpec, MergeResult, PyBlobFile, RefreshColumnResult, RefreshMaterializedViewResult,
Table, UpdateFieldMetadataResult, UpdateResult,
};
pub mod arrow;
@@ -57,6 +58,8 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<VectorQuery>()?;
m.add_class::<RecordBatchStream>()?;
m.add_class::<AddColumnsResult>()?;
m.add_class::<RefreshColumnResult>()?;
m.add_class::<RefreshMaterializedViewResult>()?;
m.add_class::<AlterColumnsResult>()?;
m.add_class::<UpdateFieldMetadataResult>()?;
m.add_class::<AddResult>()?;
python/src/permutation.rs
+3 -1
View File
@@ -268,7 +268,9 @@ impl PyPermutationReader {
.await
.infer_error()?
} else {
PermutationReader::identity(base_table).await
PermutationReader::identity(base_table)
.await
.infer_error()?
};
Ok(Self::from_reader(reader))
})
python/src/session.rs
+1 -1
View File
@@ -11,7 +11,7 @@ use pyo3::{PyResult, pyclass, pymethods};
/// Sessions allow you to configure cache sizes for index and metadata caches,
/// which can significantly impact memory use and performance. They can
/// also be re-used across multiple connections to share the same cache state.
#[pyclass(from_py_object)]
#[pyclass(module = "lancedb._lancedb", from_py_object)]
#[derive(Clone)]
pub struct Session {
pub(crate) inner: Arc<LanceSession>,
python/src/table.rs
+284 -22
View File
@@ -28,11 +28,72 @@ use pyo3::{
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
exceptions::{PyRuntimeError, PyValueError},
pyclass, pyfunction, pymethods,
types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods},
types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods, PyList, PyListMethods},
};
mod scannable;
/// Convert `LsmStats` to a Python dict, preserving the per-bucket list.
///
/// Deliberately not flattened to a table-level summary: a table is N
/// buckets on one node, and the per-bucket detail is the reason the
/// endpoint exists — flattening hides the single hot bucket someone opened
/// it to find.
fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult<Py<PyDict>> {
let out = PyDict::new(py);
let buckets = PyList::empty(py);
for b in &stats.buckets {
let e = PyDict::new(py);
e.set_item("shard_id", &b.shard_id)?;
e.set_item("status", &b.status)?;
e.set_item("writer_epoch", b.writer_epoch)?;
e.set_item("manifest_version", b.manifest_version)?;
e.set_item("current_generation", b.current_generation)?;
e.set_item(
"replay_after_wal_entry_position",
b.replay_after_wal_entry_position,
)?;
e.set_item(
"wal_entry_position_last_seen",
b.wal_entry_position_last_seen,
)?;
let generations = PyList::empty(py);
for g in &b.generations {
let ge = PyDict::new(py);
ge.set_item("generation", g.generation)?;
ge.set_item("bytes", g.bytes)?;
ge.set_item("rows", g.rows)?;
generations.append(ge)?;
}
e.set_item("generations", generations)?;
e.set_item("compacting", b.compacting)?;
e.set_item(
"memtables",
b.memtables
.as_ref()
.map(|ms| {
let l = PyList::empty(py);
for m in ms {
let d = PyDict::new(py);
d.set_item("generation", m.generation)?;
d.set_item("rows", m.rows)?;
d.set_item("bytes", m.bytes)?;
d.set_item("batches", m.batches)?;
d.set_item("indexes", m.indexes.clone())?;
l.append(d)?;
}
PyResult::Ok(l.unbind())
})
.transpose()?,
)?;
buckets.append(e)?;
}
out.set_item("buckets", buckets)?;
Ok(out.unbind())
}
#[derive(FromPyObject)]
enum PredicateArg {
Expr(PyExpr),
@@ -185,13 +246,23 @@ impl From<lancedb::table::MergeResult> for MergeResult {
}
}
/// Render for `__repr__`, so the default reads as Python's `None` rather than
/// Rust's `Some([..])`.
fn fmt_maintained(maintained: &Option<Vec<String>>) -> String {
match maintained {
Some(names) => format!("{:?}", names),
None => "None".to_string(),
}
}
/// Specification selecting Lance's MemWAL LSM-style write path for
/// `merge_insert`.
///
/// Constructed via the `bucket(...)`, `identity(...)`, or `unsharded()`
/// classmethods, then optionally chain `with_maintained_indexes(...)` and
/// `with_writer_config_defaults(...)`.
#[pyclass(from_py_object)]
/// `with_writer_config_defaults(...)`. A fresh spec maintains every index the
/// MemWAL supports, resolved on install.
#[pyclass(module = "lancedb._lancedb", from_py_object)]
#[derive(Clone, Debug)]
pub struct LsmWriteSpec {
inner: lancedb::table::LsmWriteSpec,
@@ -230,11 +301,11 @@ impl LsmWriteSpec {
}
}
/// Replace the list of indexes the MemWAL should keep up to date as
/// rows are appended. Each name must reference an index that
/// already exists on the table at the time `set_lsm_write_spec`
/// is called.
pub fn with_maintained_indexes(&self, indexes: Vec<String>) -> Self {
/// Set which indexes the MemWAL maintains. `None` (the default)
/// resolves every supported index on install; a list is verbatim,
/// and an empty list maintains nothing.
#[pyo3(signature = (indexes))]
pub fn with_maintained_indexes(&self, indexes: Option<Vec<String>>) -> Self {
Self {
inner: self.inner.clone().with_maintained_indexes(indexes),
}
@@ -256,23 +327,29 @@ impl LsmWriteSpec {
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={:?}, writer_config_defaults={:?})",
column, num_buckets, maintained_indexes, writer_config_defaults,
"LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={}, writer_config_defaults={:?})",
column,
num_buckets,
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
lancedb::table::LsmWriteSpec::Identity {
column,
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.identity(column={:?}, maintained_indexes={:?}, writer_config_defaults={:?})",
column, maintained_indexes, writer_config_defaults,
"LsmWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})",
column,
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
lancedb::table::LsmWriteSpec::Unsharded {
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.unsharded(maintained_indexes={:?}, writer_config_defaults={:?})",
maintained_indexes, writer_config_defaults,
"LsmWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})",
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
}
}
@@ -307,10 +384,10 @@ impl LsmWriteSpec {
}
}
/// Names of indexes the MemWAL should keep up to date during writes.
/// Indexes the MemWAL keeps up to date, or `None` for every supported one.
#[getter]
pub fn maintained_indexes(&self) -> Vec<String> {
self.inner.maintained_indexes().to_vec()
pub fn maintained_indexes(&self) -> Option<Vec<String>> {
self.inner.maintained_indexes().map(<[String]>::to_vec)
}
/// Default `ShardWriter` configuration recorded by this spec.
@@ -338,6 +415,67 @@ pub struct AddColumnsResult {
pub version: u64,
}
#[pyclass(get_all, from_py_object)]
#[derive(Clone, Debug)]
pub struct RefreshColumnResult {
pub rows_filled: u64,
pub version: u64,
}
#[pymethods]
impl RefreshColumnResult {
pub fn __repr__(&self) -> String {
format!(
"RefreshColumnResult(rows_filled={}, version={})",
self.rows_filled, self.version
)
}
}
impl From<lancedb::table::RefreshColumnResult> for RefreshColumnResult {
fn from(result: lancedb::table::RefreshColumnResult) -> Self {
Self {
rows_filled: result.rows_filled,
version: result.version,
}
}
}
#[pyclass(get_all, from_py_object)]
#[derive(Clone, Debug)]
pub struct RefreshMaterializedViewResult {
pub mode: String,
pub rows_written: u64,
pub source_version: u64,
pub version: u64,
}
#[pymethods]
impl RefreshMaterializedViewResult {
pub fn __repr__(&self) -> String {
format!(
"RefreshMaterializedViewResult(mode={}, rows_written={}, source_version={}, version={})",
self.mode, self.rows_written, self.source_version, self.version
)
}
}
impl From<lancedb::RefreshMaterializedViewResult> for RefreshMaterializedViewResult {
fn from(result: lancedb::RefreshMaterializedViewResult) -> Self {
let mode = match result.mode {
lancedb::RefreshMode::Rebuild => "rebuild",
lancedb::RefreshMode::Incremental => "incremental",
lancedb::RefreshMode::NoOp => "no_op",
};
Self {
mode: mode.to_string(),
rows_written: result.rows_written,
source_version: result.source_version,
version: result.version,
}
}
}
#[pymethods]
impl AddColumnsResult {
pub fn __repr__(&self) -> String {
@@ -502,7 +640,7 @@ impl PyBlobFile {
}
}
#[pyclass(get_all, from_py_object)]
#[pyclass(module = "lancedb._lancedb", get_all, from_py_object)]
#[derive(Clone, Debug)]
pub struct FtsToken {
pub text: String,
@@ -642,15 +780,19 @@ impl Table {
})
}
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None))]
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None, allow_external_blob_outside_bases=false))]
pub fn add<'a>(
self_: PyRef<'a, Self>,
data: PyScannable,
mode: String,
progress: Option<Py<PyAny>>,
write_parallelism: Option<usize>,
allow_external_blob_outside_bases: bool,
) -> PyResult<Bound<'a, PyAny>> {
let mut op = self_.inner_ref()?.add(data);
let mut op = self_
.inner_ref()?
.add(data)
.allow_external_blob_outside_bases(allow_external_blob_outside_bases);
if mode == "append" {
op = op.mode(AddDataMode::Append);
} else if mode == "overwrite" {
@@ -745,6 +887,9 @@ impl Table {
#[allow(private_interfaces)]
pub fn delete(self_: PyRef<'_, Self>, condition: PredicateArg) -> PyResult<Bound<'_, PyAny>> {
// Do not hold the Python borrow across the await. The cloned Rust table
// handle is thread-safe and allows deletes on the same Python table to
// run concurrently without PyO3 reporting "Already borrowed".
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let result = match &condition {
@@ -1336,6 +1481,51 @@ impl Table {
})
}
/// Converge the table's LSM write path into its base table.
///
/// Best-effort: with writes flowing, new rows may land after the last
/// pass. Errors if the table stops making progress.
pub fn checkpoint_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.checkpoint_lsm().await.infer_error()
})
}
/// Seal every bucket's active memtable into L0.
pub fn flush_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(
self_.py(),
async move { inner.flush_lsm().await.infer_error() },
)
}
/// Trigger a background L0 → base pass per bucket. Returns once the
/// passes are dispatched, not once they finish — watch `get_lsm_stats`.
pub fn compact_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.compact_lsm().await.infer_error()
})
}
/// Live LSM state, or `None` when the LSM write path is not enabled.
#[pyo3(signature = (include_generation_rows=false))]
pub fn get_lsm_stats(
self_: PyRef<'_, Self>,
include_generation_rows: bool,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let stats = inner
.get_lsm_stats(include_generation_rows)
.await
.infer_error()?;
Python::attach(|py| stats.map(|s| lsm_stats_to_py(py, &s)).transpose())
})
}
pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
@@ -1385,6 +1575,78 @@ impl Table {
})
}
pub fn add_computed_columns(
self_: PyRef<'_, Self>,
columns: Vec<(String, String)>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.add_columns();
for (name, expression) in columns {
builder = builder.computed(name, expression);
}
let result = builder.execute().await.infer_error()?;
Ok(AddColumnsResult::from(result))
})
}
pub fn add_function_columns(
self_: PyRef<'_, Self>,
application_json: String,
output_name: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let application =
lancedb::function::FunctionApplication::from_json(&application_json).infer_error()?;
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let builder = match output_name {
Some(name) => inner.add_columns().function_as(name, application),
None => inner.add_columns().function(application),
};
let result = builder.execute().await.infer_error()?;
Ok(AddColumnsResult::from(result))
})
}
pub fn refresh_column(self_: PyRef<'_, Self>, column: String) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let result = inner.refresh_column(column).await.infer_error()?;
Ok(RefreshColumnResult::from(result))
})
}
pub fn refresh_column_async(
self_: PyRef<'_, Self>,
column: String,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let job = inner.refresh_column_async(column).await.infer_error()?;
Ok(crate::job::Job::new_typed(job))
})
}
#[pyo3(signature = (full=false, source_version=None))]
pub fn refresh_materialized_view(
self_: PyRef<'_, Self>,
full: bool,
source_version: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let view = lancedb::MaterializedView::from_table(inner)
.await
.infer_error()?;
let mut builder = view.refresh().full(full);
if let Some(version) = source_version {
builder = builder.source_version(version);
}
let result = builder.execute().await.infer_error()?;
Ok(RefreshMaterializedViewResult::from(result))
})
}
pub fn add_columns_with_schema(
self_: PyRef<'_, Self>,
schema: PyArrowType<Schema>,
@@ -1682,7 +1944,7 @@ impl Branches {
}
#[pyo3(signature = (from_branch, dry_run=false))]
pub fn merge(
pub fn cherry_pick(
self_: PyRef<'_, Self>,
from_branch: String,
dry_run: bool,
@@ -1690,7 +1952,7 @@ impl Branches {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let result = inner
.merge_branch(&from_branch, dry_run)
.cherry_pick(&from_branch, dry_run)
.await
.infer_error()?;
Python::attach(|py| struct_to_wire_py(py, &result))
python/uv.lock
Generated
+2 -2
View File
@@ -1998,12 +1998,12 @@ requires-dist = [
{ name = "pillow", marker = "extra == 'clip'", specifier = ">=12.1.1" },
{ name = "pillow", marker = "extra == 'embeddings'", specifier = ">=12.1.1" },
{ name = "pillow", marker = "extra == 'siglip'", specifier = ">=12.1.1" },
{ name = "polars", marker = "extra == 'tests'", specifier = ">=0.19,<=1.3.0" },
{ name = "polars", marker = "extra == 'tests'", specifier = ">=0.19,<=1.32.3" },
{ name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.5.0" },
{ name = "pyarrow", specifier = ">=16" },
{ name = "pyarrow", marker = "extra == 'tests'", specifier = "<25" },
{ name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" },
{ name = "pydantic", specifier = ">=1.10" },
{ name = "pydantic", specifier = ">=2.7.4,<3" },
{ name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" },
{ name = "pylance", marker = "extra == 'tests'", specifier = "==9.0.0rc1" },
{ name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" },