rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-09-22 21:15:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Gatefixer
2026-09-17 06:49:46 +00:00
parent 9bfd475702 60a1b4c219
commit 64e4d42488
221 changed files with 34375 additions and 14194 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/Cargo.toml
+4 -2
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb-python"
version = "0.38.0-beta.12"
version = "0.40.0-beta.1"
publish = false
edition.workspace = true
description = "Python bindings for LanceDB"
@@ -15,6 +15,7 @@ name = "_lancedb"
crate-type = ["cdylib"]
[dependencies]
arc-swap = "1.9"
arrow = { workspace = true, features = ["pyarrow"] }
async-trait.workspace = true
bytes.workspace = true
@@ -28,7 +29,7 @@ env_logger.workspace = true
log.workspace = true
# 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"] }
pyo3 = { version = "0.28", features = ["abi3-py310", "chrono", "uuid"] }
chrono.workspace = true
pyo3-async-runtimes = { version = "0.28", features = [
"attributes",
@@ -40,6 +41,7 @@ serde.workspace = true
serde_json.workspace = true
snafu.workspace = true
tokio.workspace = true
uuid.workspace = true
libc = "0.2"
[build-dependencies]
python/pyproject.toml
+2 -1
View File
@@ -63,7 +63,7 @@ tests = [
"polars>=0.19, <=1.32.3",
"pyarrow<25",
"pyarrow-stubs>=16.0",
"pylance==9.0.0rc1",
"pylance==9.0.0",
"requests>=2.31.0",
"datafusion>=54,<55",
"opentelemetry-sdk>=1.30.0",
@@ -139,6 +139,7 @@ include = [
"python/lancedb/exceptions.py",
"python/lancedb/background_loop.py",
"python/lancedb/schema.py",
"python/lancedb/sql.py",
"python/lancedb/remote/__init__.py",
"python/lancedb/remote/errors.py",
"python/lancedb/embeddings/__init__.py",
python/python/lancedb/__init__.py
+33
View File
@@ -22,7 +22,11 @@ from .remote.db import RemoteDBConnection
from .expr import Expr, col, lit, func
from .schema import blob, vector
from .job import AsyncJob, Job
from .sql import AsyncQuery as AsyncSqlQuery
from .sql import Query as SqlQuery
from .sql import QueryDescription
from .functions import (
AssignmentMapping as AssignmentMapping,
FunctionArtifactRequest as FunctionArtifactRequest,
FunctionApplication as FunctionApplication,
FunctionBinding as FunctionBinding,
@@ -33,6 +37,8 @@ from .functions import (
UdfDefinition as UdfDefinition,
udf as udf,
)
from .secrets import EnvVarSecret as EnvVarSecret
from .secrets import SecretInfo as SecretInfo
from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
@@ -48,6 +54,14 @@ from .namespace import (
AsyncLanceNamespaceDBConnection,
)
from .catalog import (
AsyncCatalog,
Catalog,
ListDatabasesResponse,
connect_catalog,
connect_catalog_async,
)
if TYPE_CHECKING:
from lance.blob import BlobType as BlobType
@@ -101,6 +115,7 @@ def connect(
api_key: Optional[str] = None,
region: str = "us-east-1",
host_override: Optional[str] = None,
sql_host_override: Optional[str] = None,
read_consistency_interval: Optional[timedelta] = None,
request_thread_pool: Optional[Union[int, ThreadPoolExecutor]] = None,
client_config: Union[ClientConfig, Dict[str, Any], None] = None,
@@ -129,6 +144,9 @@ def connect(
The region to use for LanceDB Cloud.
host_override: str, optional
The override url for LanceDB Cloud.
sql_host_override: str, optional
The remote SQL service endpoint override. The client connects lazily when SQL
is first executed and retains that connection.
read_consistency_interval: timedelta, default None
The interval at which to check for updates to the table from other
processes. If None, then consistency is not checked. For performance
@@ -270,6 +288,7 @@ def connect(
api_key,
region,
host_override,
sql_host_override=sql_host_override,
# TODO: remove this (deprecation warning downstream)
request_thread_pool=request_thread_pool,
client_config=client_config,
@@ -412,6 +431,7 @@ def deserialize_conn(
parsed["api_key"],
parsed.get("region", "us-east-1"),
host_override=parsed.get("host_override"),
sql_host_override=parsed.get("sql_host_override"),
client_config=parsed.get("client_config"),
storage_options=storage_options,
)
@@ -425,6 +445,7 @@ async def connect_async(
api_key: Optional[str] = None,
region: str = "us-east-1",
host_override: Optional[str] = None,
sql_host_override: Optional[str] = None,
read_consistency_interval: Optional[timedelta] = None,
client_config: Optional[Union[ClientConfig, Dict[str, Any]]] = None,
storage_options: Optional[Dict[str, str]] = None,
@@ -447,6 +468,9 @@ async def connect_async(
The region to use for LanceDB Cloud.
host_override: str, optional
The override url for LanceDB Cloud.
sql_host_override: str, optional
The remote SQL service endpoint override. The client connects lazily when SQL
is first executed and retains that connection.
read_consistency_interval: timedelta, default None
The interval at which to check for updates to the table from other
processes. If None, then consistency is not checked. For performance
@@ -534,6 +558,7 @@ async def connect_async(
api_key,
region,
host_override,
sql_host_override,
read_consistency_interval_secs,
client_config,
storage_options,
@@ -546,6 +571,11 @@ async def connect_async(
__all__ = [
"Catalog",
"AsyncCatalog",
"ListDatabasesResponse",
"connect_catalog",
"connect_catalog_async",
"AsyncMaterializedView",
"MaterializedView",
"MaterializedViewDefinition",
@@ -556,6 +586,7 @@ __all__ = [
"connect_namespace_async",
"AsyncConnection",
"AsyncJob",
"AsyncSqlQuery",
"AsyncLanceNamespaceDBConnection",
"AsyncTable",
"CompactionOptions",
@@ -571,6 +602,8 @@ __all__ = [
"vector",
"DBConnection",
"Job",
"QueryDescription",
"SqlQuery",
"LanceDBConnection",
"LanceNamespaceDBConnection",
"LsmWriteSpec",
python/python/lancedb/_lancedb.pyi
+140 -6
View File
@@ -1,6 +1,7 @@
from datetime import date, datetime, timedelta
from decimal import Decimal
from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal
from uuid import UUID
import pyarrow as pa
@@ -147,15 +148,35 @@ class Connection(object):
start_after: Optional[str],
limit: Optional[int],
) -> list[str]: ... # Deprecated: Use list_tables instead
def job(self, job_id: str) -> Job: ...
async def open_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_functions(self) -> List[str]: ...
async def drop_function(self, name: str, version: str) -> bool: ...
async def create_secret(
self, name: str, value: str, namespace_path: Optional[List[str]] = None
) -> None: ...
async def alter_secret(
self, name: str, value: str, namespace_path: Optional[List[str]] = None
) -> None: ...
async def list_secrets(
self, namespace_path: Optional[List[str]] = None
) -> List[str]: ...
async def drop_secret(
self, name: str, namespace_path: Optional[List[str]] = None
) -> None: ...
async def describe_secret(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Tuple[str, int, int]: ...
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: ...
async def job_history(
self, job_id: Optional[str] = None
) -> List[pa.RecordBatch]: ...
async def execute_query_async(
self,
query: str,
*,
default_namespace_path: Optional[List[str]] = None,
) -> SqlQuery: ...
async def describe_query(self, query_id: UUID) -> QueryDescription: ...
async def create_table(
self,
name: str,
@@ -205,8 +226,24 @@ class Connection(object):
projections: Optional[List[Tuple[str, str]]] = None,
filter: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> Table: ...
async def create_materialized_view_async(
self,
name: str,
source: str,
projections: Optional[List[Tuple[str, str]]] = None,
filter: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> Job: ...
async def list_materialized_views(self) -> List[str]: ...
async def drop_materialized_view(
self, name: str, namespace_path: Optional[List[str]] = None
) -> None: ...
async def drop_materialized_view_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Job: ...
async def drop_table(
self, name: str, namespace_path: Optional[List[str]] = None
) -> None: ...
@@ -234,9 +271,20 @@ class BlobFile:
class Job:
@property
def id(self) -> Optional[str]: ...
@property
def _state(self) -> Optional[str]: ...
@property
def _description(self) -> Optional[JobDescription]: ...
async def status(self) -> str: ...
async def wait(self) -> Optional[str]: ...
async def cancel(self) -> None: ...
async def refresh(self) -> None: ...
async def events(
self,
*,
limit: Optional[int] = None,
filter: Optional[str] = None,
) -> pa.Table: ...
class JobInfo:
@property
@@ -250,6 +298,37 @@ class JobInfo:
@property
def created_at_millis(self) -> int: ...
class SessionStatus:
@property
def resource(self) -> Optional[str]: ...
@property
def audience(self) -> Optional[str]: ...
@property
def refreshable(self) -> bool: ...
@property
def issuer_url(self) -> str: ...
@property
def client_id(self) -> str: ...
@property
def scopes(self) -> List[str]: ...
@property
def flow(self) -> str: ...
@property
def obtained_at(self) -> Optional[int]: ...
def __repr__(self) -> str: ...
class SessionLogout:
@property
def removed(self) -> bool: ...
def __repr__(self) -> str: ...
class OAuthSession:
def __init__(self, config: Any) -> None: ...
async def login(self) -> SessionStatus: ...
async def status(self) -> SessionStatus: ...
async def logout(self) -> SessionLogout: ...
def __repr__(self) -> str: ...
class JobFailureInfo:
@property
def phase(self) -> Optional[str]: ...
@@ -268,10 +347,33 @@ class JobDescription:
@property
def creation_ms(self) -> int: ...
@property
def spec_json(self) -> Optional[str]: ...
def _spec_json(self) -> Optional[str]: ...
@property
def _result_json(self) -> Optional[str]: ...
@property
def spec(self) -> Optional[Any]: ...
@property
def result(self) -> Optional[Any]: ...
@property
def failure(self) -> Optional[JobFailureInfo]: ...
class SqlQuery:
@property
def id(self) -> UUID: ...
async def describe(self) -> QueryDescription: ...
async def reader(self) -> RecordBatchStream: ...
async def cancel(self) -> None: ...
class QueryDescription:
@property
def id(self) -> UUID: ...
@property
def status(self) -> str: ...
@property
def progress(self) -> Optional[float]: ...
@property
def expires_at(self) -> Optional[datetime]: ...
class Table:
def name(self) -> str: ...
def __repr__(self) -> str: ...
@@ -362,6 +464,10 @@ class Table:
async def refresh_materialized_view(
self, full: bool = False, source_version: Optional[int] = None
) -> RefreshMaterializedViewResult: ...
async def refresh_materialized_view_async(
self, full: bool = False, source_version: Optional[int] = None
) -> Job: ...
async def materialized_view_definition(self) -> str: ...
async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ...
async def alter_columns(
self, columns: list[dict[str, Any]]
@@ -451,6 +557,7 @@ async def connect(
api_key: Optional[str],
region: Optional[str],
host_override: Optional[str],
sql_host_override: Optional[str],
read_consistency_interval: Optional[float],
client_config: Optional[Union[ClientConfig, Dict[str, Any]]],
storage_options: Optional[Dict[str, str]],
@@ -607,6 +714,7 @@ class FullTextQuery:
class PyQueryRequest:
limit: Optional[int]
offset: Optional[int]
take_offsets: Optional[List[int]]
filter: Optional[Union[str, bytes]]
full_text_search: Optional[FullTextQuery]
select: Optional[Union[str, List[str]]]
@@ -714,6 +822,8 @@ class RefreshColumnResult:
version: int
class RefreshMaterializedViewResult:
@staticmethod
def from_json(value: str) -> RefreshMaterializedViewResult: ...
mode: str
rows_written: int
source_version: int
@@ -762,3 +872,27 @@ def fts_query_to_json(query: Any) -> str: ...
class PermutationReader:
def __init__(self, base_table: Table, permutation_table: Table): ...
class Catalog:
@property
def uri(self) -> str: ...
async def create_database(
self, name: str, *, exist_ok: bool = False
) -> Connection: ...
async def connect_database(self, name: str) -> Connection: ...
async def drop_database(
self, name: str, *, ignore_missing: bool = False
) -> None: ...
async def list_databases(
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
) -> tuple[list[str], Optional[str]]: ...
async def connect_catalog(
endpoint: str,
*,
api_key: Optional[str] = None,
client_config: Optional[Any] = None,
sql_host_override: Optional[str] = None,
read_consistency_interval: Optional[float] = None,
oauth_config: Optional[Any] = None,
) -> Catalog: ...
python/python/lancedb/catalog.py
+196
View File
@@ -0,0 +1,196 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Remote catalogs manage databases through a server's root namespace."""
from dataclasses import dataclass
from datetime import timedelta
from typing import Any, Optional, Union
from . import _lancedb
from .background_loop import LOOP
from .db import AsyncConnection, DBConnection
from .remote import ClientConfig, OAuthConfig
from .remote.db import RemoteDBConnection
@dataclass
class ListDatabasesResponse:
"""A page of database names and an optional continuation token."""
databases: list[str]
page_token: Optional[str] = None
class AsyncCatalog:
"""An asynchronous remote catalog returned by
[connect_catalog_async][lancedb.connect_catalog_async].
Create/connect return ordinary [AsyncConnection][lancedb.db.AsyncConnection]
instances. Drop uses restricted behavior: remove the database's tables first.
"""
def __init__(self, inner: _lancedb.Catalog):
self._inner = inner
@property
def uri(self) -> str:
"""The catalog's root namespace endpoint."""
return self._inner.uri
async def create_database(
self, name: str, *, exist_ok: bool = False
) -> AsyncConnection:
"""Create a database, or open an existing one when ``exist_ok=True``."""
return AsyncConnection(
await self._inner.create_database(name, exist_ok=exist_ok)
)
async def connect_database(self, name: str) -> AsyncConnection:
"""Connect to an existing database by its logical name."""
return AsyncConnection(await self._inner.connect_database(name))
async def list_databases(
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
) -> ListDatabasesResponse:
"""List a page of databases. Pass the returned token for the next page."""
names, token = await self._inner.list_databases(
limit=limit, page_token=page_token
)
return ListDatabasesResponse(names, token)
async def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
"""Drop an empty database. A nonempty database is an error."""
await self._inner.drop_database(name, ignore_missing=ignore_missing)
class Catalog:
"""A synchronous remote catalog returned by
[connect_catalog][lancedb.connect_catalog].
Examples
--------
```python
catalog = lancedb.connect_catalog("https://my-server.example", api_key="secret")
db = catalog.create_database("analytics", exist_ok=True)
page = catalog.list_databases(limit=20)
```
"""
def __init__(
self,
inner: AsyncCatalog,
*,
api_key=None,
client_config=None,
sql_host_override: Optional[str] = None,
oauth_config: Optional[OAuthConfig] = None,
):
self._inner = inner
self._api_key = api_key
self._client_config = client_config
self._sql_host_override = sql_host_override
self._oauth_config = oauth_config
@property
def uri(self) -> str:
"""The catalog's root namespace endpoint."""
return self._inner.uri
def create_database(self, name: str, *, exist_ok: bool = False) -> DBConnection:
"""Create a database, or open an existing one when ``exist_ok=True``."""
inner = LOOP.run(self._inner.create_database(name, exist_ok=exist_ok))
return self._wrap_database(name, inner)
def connect_database(self, name: str) -> DBConnection:
"""Connect to an existing database by its logical name."""
return self._wrap_database(name, LOOP.run(self._inner.connect_database(name)))
def _wrap_database(self, name: str, inner: AsyncConnection) -> DBConnection:
return RemoteDBConnection._from_catalog(
inner,
name,
self.uri,
self._api_key,
self._client_config,
self._oauth_config,
self._sql_host_override,
)
def list_databases(
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
) -> ListDatabasesResponse:
"""List a page of databases. Pass the returned token for the next page."""
return LOOP.run(self._inner.list_databases(limit=limit, page_token=page_token))
def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
"""Drop an empty database. A nonempty database is an error."""
LOOP.run(self._inner.drop_database(name, ignore_missing=ignore_missing))
async def connect_catalog_async(
endpoint: str,
*,
api_key: Optional[str] = None,
client_config: Optional[Union[ClientConfig, dict[str, Any]]] = None,
sql_host_override: Optional[str] = None,
read_consistency_interval: Optional[timedelta] = None,
oauth_config: Optional[OAuthConfig] = None,
) -> AsyncCatalog:
"""Connect to an HTTP(S) server's root catalog.
Root requests omit database-selection headers. API key, client configuration,
OAuth, and table read consistency settings are inherited by opened databases.
Database names containing slashes remain single logical names.
Set ``sql_host_override`` to the SQL service endpoint to execute SQL through
returned connections when the catalog endpoint uses HTTPS.
"""
if isinstance(client_config, dict):
client_config = ClientConfig(**client_config)
if client_config is None:
client_config = ClientConfig()
inner = await _lancedb.connect_catalog(
endpoint,
api_key=api_key,
client_config=client_config,
sql_host_override=sql_host_override,
read_consistency_interval=(
read_consistency_interval.total_seconds()
if read_consistency_interval is not None
else None
),
oauth_config=oauth_config,
)
return AsyncCatalog(inner)
def connect_catalog(
endpoint: str,
*,
api_key: Optional[str] = None,
client_config: Optional[Union[ClientConfig, dict[str, Any]]] = None,
sql_host_override: Optional[str] = None,
read_consistency_interval: Optional[timedelta] = None,
oauth_config: Optional[OAuthConfig] = None,
) -> Catalog:
"""Connect synchronously to an HTTP(S) server's root catalog.
See [connect_catalog_async][lancedb.connect_catalog_async] for options.
Local filesystem and object-store catalogs are not supported.
"""
return Catalog(
LOOP.run(
connect_catalog_async(
endpoint,
api_key=api_key,
client_config=client_config,
sql_host_override=sql_host_override,
read_consistency_interval=read_consistency_interval,
oauth_config=oauth_config,
)
),
api_key=api_key,
client_config=client_config,
sql_host_override=sql_host_override,
oauth_config=oauth_config,
)
python/python/lancedb/db.py
+531 -90
View File
@@ -17,8 +17,10 @@ from typing import (
List,
Literal,
Optional,
Sequence,
Union,
)
from uuid import UUID
if sys.version_info >= (3, 12):
from typing import override
@@ -47,12 +49,21 @@ from . import __version__
from ._lancedb import connect as lancedb_connect # type: ignore
from .functions import FunctionVersion, UdfDefinition
from .job import AsyncJob, Job, _typed_job
from .sql import AsyncQuery as AsyncSqlQuery
from .sql import Query as SqlQuery
from .sql import QueryDescription
from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
SelectArg,
normalize_select,
)
from .secrets import (
EnvVarSecret,
SecretInfo,
validate_namespace_path,
validate_secret_name,
)
from .table import (
AsyncTable,
LanceTable,
@@ -68,10 +79,11 @@ import deprecation
if TYPE_CHECKING:
import pyarrow as pa
from .arrow import AsyncRecordBatchReader
from .pydantic import LanceModel
from ._lancedb import Connection as LanceDbConnection
from ._lancedb import JobDescription, JobInfo
from ._lancedb import JobInfo
from .common import DATA, URI
from .embeddings import EmbeddingFunctionConfig
from ._lancedb import Session
@@ -524,13 +536,14 @@ class DBConnection(EnforceOverrides):
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> 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 view is populated before creation returns. Pass
``with_no_data=True`` to create only its definition and empty backing
table. The view is a normal table: it can be queried, indexed and
searched, and it appears in ``table_names``.
The source table must have stable row ids (create it with the
``new_table_enable_stable_row_ids`` storage option): they keep the
@@ -551,6 +564,8 @@ class DBConnection(EnforceOverrides):
SQL predicate; only matching source rows appear in the view.
limit: int, optional
Cap the view at this many rows, in materialization order.
with_no_data: bool, default False
Skip the initial refresh and leave the backing table empty.
Returns
-------
@@ -560,6 +575,27 @@ class DBConnection(EnforceOverrides):
"materialized views are not supported on this connection type"
)
def create_materialized_view_async(
self,
name: str,
source: str,
*,
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> Job[None]:
"""Submit materialized-view creation and return its job.
The job may already be complete for a local database. On LanceDB
Cloud and Enterprise, its ``id`` is the server job identifier from
the ``202 Accepted`` create response. Wait for the job before opening
or querying the view.
"""
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``.
@@ -580,6 +616,32 @@ class DBConnection(EnforceOverrides):
"materialized views are not supported on this connection type"
)
def drop_materialized_view(
self, name: str, namespace_path: Optional[List[str]] = None
) -> None:
"""Drop a materialized view.
The view may become unavailable before physical cleanup finishes. Use
:meth:`drop_materialized_view_async` to retain and wait for the cleanup
job.
"""
raise NotImplementedError(
"materialized views are not supported on this connection type"
)
def drop_materialized_view_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Job[None]:
"""Start dropping a materialized view and return its cleanup job.
The job may already be complete for a local database. On LanceDB Cloud
and Enterprise, its ``id`` is the server job identifier from the
``202 Accepted`` drop response.
"""
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.
@@ -687,20 +749,55 @@ 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.
def create_function(
self,
definition: UdfDefinition,
*,
secrets: Optional[Sequence[EnvVarSecret]] = None,
) -> FunctionVersion:
"""Build and register a scalar Python UDF, then return its version.
The server builds the OCI image and registers the completed artifact.
This is the blocking counterpart of :meth:`create_function_async`.
Local connections raise ``NotImplementedError``.
Parameters
----------
definition : UdfDefinition
A callable decorated with [udf][lancedb.udf].
secrets : sequence of EnvVarSecret, optional
One [EnvVarSecret][lancedb.secrets.EnvVarSecret] per credential the
Function needs, each naming a Secret and the environment variable
its value arrives in. The Function's source is unchanged by this;
it reads the variable the way it already did.
Examples
--------
```python
db.create_secret("openai-prod", os.environ["OPENAI_API_KEY"])
db.create_function(
analyze_caption,
secrets=[
EnvVarSecret(
secret_name="openai-prod", env_variable="OPENAI_API_KEY"
)
],
)
```
"""
return self.create_function_async(definition).wait()
return self.create_function_async(definition, secrets=secrets).wait()
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
"""Register a scalar Python UDF through the remote Function catalog.
def create_function_async(
self,
definition: UdfDefinition,
*,
secrets: Optional[Sequence[EnvVarSecret]] = None,
) -> Job[FunctionVersion]:
"""Submit a scalar Python UDF for building and registration.
Submission returns a typed job. The immutable Function version becomes
available only when :meth:`Job.wait` succeeds. Local connections raise
``NotImplementedError``.
The server-side job builds the OCI image, then registers the completed
artifact. Waiting on the job returns the immutable Function version.
Local connections raise ``NotImplementedError``.
"""
raise NotImplementedError(
"Function catalog operations are not supported for this connection type"
@@ -712,26 +809,117 @@ class DBConnection(EnforceOverrides):
"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.
def list_functions(self) -> List[FunctionVersion]:
"""List every published immutable Function version.
The handle is constructed without a server round trip; an unknown id
surfaces when the handle is used. Dropping the handle has no effect
on the job itself.
Results are ordered by Function name then version. Local connections
raise ``NotImplementedError``.
Examples
--------
List the identities available to use in Function-backed columns:
```python
[(function.name, function.version) for function in db.list_functions()]
```
"""
raise NotImplementedError("job is not supported for this connection type")
raise NotImplementedError(
"Function catalog operations are not supported for this connection type"
)
def drop_function(self, name: str, *, version: str) -> bool:
"""Remove the current Function name binding from the remote catalog.
The requested version must exist in the currently named object.
Object history and existing computed-column references are retained.
Returns True when the name was removed and False when it was absent.
Local connections raise NotImplementedError.
"""
raise NotImplementedError(
"Function catalog operations are not supported for this connection type"
)
def create_secret(
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
) -> None:
"""Create a named Secret in this database.
Fails if the name is taken, so a create never silently becomes a
rotation. Nothing reads the value back: it is bound to a Function by
name and resolved by the service when that Function runs. Local
connections raise ``NotImplementedError``.
"""
raise NotImplementedError(
"Secret operations are not supported for this connection type"
)
def alter_secret(
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
) -> None:
"""Replace the credential behind an existing Secret.
Fails if it does not exist. Every Function bound to the Secret uses the
new value from its next job, and no new Function version is created --
which is how a rotation reaches columns pinned to a version registered
before it. Local connections raise ``NotImplementedError``.
"""
raise NotImplementedError(
"Secret operations are not supported for this connection type"
)
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
"""The names of every Secret in this database.
Names only. No method returns a stored credential, by construction
rather than by policy. Local connections raise ``NotImplementedError``.
"""
raise NotImplementedError(
"Secret operations are not supported for this connection type"
)
def drop_secret(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> None:
"""Drop a Secret.
Functions bound to it fail at their next job, naming the Secret; that
is the revocation path. The name becomes free to reuse, and a new
Secret under it is picked up by everything still bound to that name.
Local connections raise ``NotImplementedError``.
"""
raise NotImplementedError(
"Secret operations are not supported for this connection type"
)
def describe_secret(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> SecretInfo:
"""What this database records about a Secret: name and timestamps.
Never the value -- there is no code path that could return one. Local
connections raise ``NotImplementedError``.
"""
raise NotImplementedError(
"Secret operations are not supported for this connection type"
)
def open_job(self, job_id: str) -> Job:
"""Open a server-side job by id, returning a handle with its record
already populated.
The returned [Job][lancedb.job.Job] answers for its own state,
specification, result, failure and event history, so there is no
separate connection-level call for any of them.
Raises `JobNotFoundError` when the server has no such job, the way
`open_table` does for a missing table.
"""
raise NotImplementedError("open_job is not supported for this connection type")
def list_jobs(self) -> List[JobInfo]:
"""List server-side jobs across the database's tables."""
raise NotImplementedError("list_jobs is not supported for this connection type")
def get_job(self, job_id: str) -> Optional[JobDescription]:
"""Describe a single server-side job by id.
Returns None when the server has no such job.
"""
raise NotImplementedError("get_job is not supported for this connection type")
def cancel_job(self, job_id: str) -> bool:
"""Request cancellation of a server-side job by id.
@@ -743,14 +931,38 @@ class DBConnection(EnforceOverrides):
"cancel_job is not supported for this connection type"
)
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
"""The lifecycle event history of a server-side job, as Arrow batches.
def execute_query(
self,
query: str,
*,
default_namespace_path: Optional[List[str]] = None,
) -> pa.RecordBatchReader:
"""Execute SQL and return a blocking Arrow reader.
Lists history across all jobs when `job_id` is None.
This submits through :meth:`execute_query_async` and waits until the
initial result stream is readable. It does not wait for the full query
to finish.
"""
raise NotImplementedError(
"job_history is not supported for this connection type"
)
return self.execute_query_async(
query,
default_namespace_path=default_namespace_path,
).reader()
def execute_query_async(
self,
query: str,
*,
default_namespace_path: Optional[List[str]] = None,
) -> SqlQuery:
"""Start executing SQL and return its query handle.
Local connections do not support SQL.
"""
raise NotImplementedError("SQL is not supported for this connection type")
def describe_query(self, query_id: UUID) -> QueryDescription:
"""Describe a submitted SQL query by its connection-scoped id."""
raise NotImplementedError("SQL is not supported for this connection type")
class LanceDBConnection(DBConnection):
@@ -847,6 +1059,7 @@ class LanceDBConnection(DBConnection):
None,
None,
None,
None,
read_consistency_interval_secs,
None,
storage_options,
@@ -1215,6 +1428,7 @@ class LanceDBConnection(DBConnection):
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> MaterializedView:
"""Define a materialized view named ``name`` over the table ``source``.
See
@@ -1235,17 +1449,44 @@ class LanceDBConnection(DBConnection):
... select=["name", ("shout", "upper(name)")],
... where="age >= 18",
... )
>>> result = view.refresh()
>>> result.rows_written
>>> view.table.count_rows()
1
"""
LOOP.run(
self._conn.create_materialized_view(
name, source, select=select, where=where, limit=limit
name,
source,
select=select,
where=where,
limit=limit,
with_no_data=with_no_data,
)
)
return MaterializedView(self.open_table(name))
@override
def create_materialized_view_async(
self,
name: str,
source: str,
*,
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> Job[None]:
job = LOOP.run(
self._conn.create_materialized_view_async(
name,
source,
select=select,
where=where,
limit=limit,
with_no_data=with_no_data,
)
)
return Job(job)
@override
def open_materialized_view(self, name: str) -> MaterializedView:
"""Open the materialized view named ``name``."""
@@ -1258,6 +1499,25 @@ class LanceDBConnection(DBConnection):
"""The names of the materialized views in this database."""
return LOOP.run(self._conn.list_materialized_views())
@override
def drop_materialized_view(
self, name: str, namespace_path: Optional[List[str]] = None
) -> None:
if namespace_path is None:
namespace_path = []
LOOP.run(self._conn.drop_materialized_view(name, namespace_path=namespace_path))
@override
def drop_materialized_view_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Job[None]:
if namespace_path is None:
namespace_path = []
job = LOOP.run(
self._conn.drop_materialized_view_async(name, namespace_path=namespace_path)
)
return Job(job)
def clone_table(
self,
target_table_name: str,
@@ -1395,37 +1655,67 @@ class LanceDBConnection(DBConnection):
)
@override
def job(self, job_id: str) -> Job:
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
The handle is constructed without a server round trip; an unknown id
surfaces when the handle is used. Dropping the handle has no effect
on the job itself.
def open_job(self, job_id: str) -> Job:
"""Open a server-side job by id. See
[DBConnection.open_job][lancedb.db.DBConnection.open_job].
"""
return Job(self._conn.job(job_id))
return Job(LOOP.run(self._conn.open_job(job_id)))
@override
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
job = LOOP.run(self._conn.create_function_async(definition))
def create_function_async(
self,
definition: UdfDefinition,
*,
secrets: Optional[Sequence[EnvVarSecret]] = None,
) -> Job[FunctionVersion]:
job = LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
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_functions(self) -> List[FunctionVersion]:
return LOOP.run(self._conn.list_functions())
@override
def drop_function(self, name: str, *, version: str) -> bool:
return LOOP.run(self._conn.drop_function(name, version=version))
@override
def create_secret(
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
) -> None:
LOOP.run(self._conn.create_secret(name, value, namespace_path=namespace_path))
@override
def alter_secret(
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
) -> None:
LOOP.run(self._conn.alter_secret(name, value, namespace_path=namespace_path))
@override
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
return LOOP.run(self._conn.list_secrets(namespace_path=namespace_path))
@override
def drop_secret(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> None:
LOOP.run(self._conn.drop_secret(name, namespace_path=namespace_path))
@override
def describe_secret(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> SecretInfo:
return LOOP.run(self._conn.describe_secret(name, namespace_path=namespace_path))
@override
def list_jobs(self) -> List[JobInfo]:
"""List server-side jobs across the database's tables."""
return LOOP.run(self._conn.list_jobs())
@override
def get_job(self, job_id: str) -> Optional[JobDescription]:
"""Describe a single server-side job by id.
Returns None when the server has no such job.
"""
return LOOP.run(self._conn.get_job(job_id))
@override
def cancel_job(self, job_id: str) -> bool:
"""Request cancellation of a server-side job by id.
@@ -1436,14 +1726,6 @@ class LanceDBConnection(DBConnection):
"""
return LOOP.run(self._conn.cancel_job(job_id))
@override
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
"""The lifecycle event history of a server-side job, as Arrow batches.
Lists history across all jobs when `job_id` is None.
"""
return LOOP.run(self._conn.job_history(job_id))
@override
def namespace_client(self) -> LanceNamespace:
"""Get the equivalent namespace client for this connection.
@@ -2036,6 +2318,7 @@ class AsyncConnection(object):
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> AsyncMaterializedView:
"""Define a materialized view named ``name`` over the table ``source``.
See
@@ -2047,19 +2330,40 @@ class AsyncConnection(object):
projections=normalize_select(select),
filter=where,
limit=limit,
with_no_data=with_no_data,
)
return AsyncMaterializedView(AsyncTable(inner))
async def create_materialized_view_async(
self,
name: str,
source: str,
*,
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> AsyncJob[None]:
"""Submit materialized-view creation and return its job.
Wait for the returned job before opening or querying the view.
"""
inner = await self._inner.create_materialized_view_async(
name,
source,
projections=normalize_select(select),
filter=where,
limit=limit,
with_no_data=with_no_data,
)
return AsyncJob(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
@@ -2072,6 +2376,41 @@ class AsyncConnection(object):
"""
return await self._inner.list_materialized_views()
async def drop_materialized_view(
self,
name: str,
*,
namespace_path: Optional[List[str]] = None,
) -> None:
"""Drop a materialized view.
The view may become unavailable before physical cleanup finishes. Use
:meth:`drop_materialized_view_async` to retain and wait for the cleanup
job.
"""
if namespace_path is None:
namespace_path = []
await self._inner.drop_materialized_view(name, namespace_path=namespace_path)
async def drop_materialized_view_async(
self,
name: str,
*,
namespace_path: Optional[List[str]] = None,
) -> AsyncJob[None]:
"""Start dropping a materialized view and return its cleanup job.
Await :meth:`AsyncJob.wait` before assuming physical cleanup has
finished.
"""
if namespace_path is None:
namespace_path = []
return AsyncJob(
await self._inner.drop_materialized_view_async(
name, namespace_path=namespace_path
)
)
async def clone_table(
self,
target_table_name: str,
@@ -2214,46 +2553,114 @@ class AsyncConnection(object):
namespace_path = []
await self._inner.drop_all_tables(namespace_path=namespace_path)
def job(self, job_id: str) -> AsyncJob:
"""An [AsyncJob][lancedb.job.AsyncJob] handle for a server-side job
by id.
The handle is constructed without a server round trip; an unknown id
surfaces when the handle is used. Dropping the handle has no effect
on the job itself.
async def open_job(self, job_id: str) -> AsyncJob:
"""Open a server-side job by id. See
[DBConnection.open_job][lancedb.db.DBConnection.open_job].
"""
return AsyncJob(self._inner.job(job_id))
return AsyncJob(await self._inner.open_job(job_id))
async def create_function_async(
self, definition: UdfDefinition
self,
definition: UdfDefinition,
*,
secrets: Optional[Sequence[EnvVarSecret]] = None,
) -> AsyncJob[FunctionVersion]:
"""Register a scalar Python UDF through the remote Function catalog.
"""Submit a scalar Python UDF for building and registration.
The returned typed job resolves to the immutable Function version.
Local connections raise ``NotImplementedError``.
The server-side job builds the OCI image, then registers the completed
artifact. Waiting on the job returns the immutable Function version.
``secrets`` is a sequence of
[EnvVarSecret][lancedb.secrets.EnvVarSecret], each naming a Secret and
the environment variable its value arrives in. 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()
)
request = definition.bind_secrets(secrets)
inner = await self._inner.create_function_async(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_functions(self) -> List[FunctionVersion]:
"""List every published immutable Function version.
Results are ordered by Function name then version. Local connections
raise ``NotImplementedError``.
"""
return [
FunctionVersion.from_json(value)
for value in await self._inner.list_functions()
]
async def drop_function(self, name: str, *, version: str) -> bool:
"""Remove the current name binding, retaining the object and its history."""
return await self._inner.drop_function(name, version)
async def create_secret(
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
) -> None:
"""Create a named Secret in this database.
Fails if the name is taken, so a create never silently becomes a
rotation. Nothing reads the value back.
"""
await self._inner.create_secret(
validate_secret_name(name),
value,
list(validate_namespace_path(namespace_path)),
)
async def alter_secret(
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
) -> None:
"""Replace the credential behind an existing Secret.
Fails if it does not exist. Bound Functions use the new value from
their next job, with no new Function version.
"""
await self._inner.alter_secret(
validate_secret_name(name),
value,
list(validate_namespace_path(namespace_path)),
)
async def list_secrets(
self, *, namespace_path: Optional[List[str]] = None
) -> List[str]:
"""The names of every Secret in this database. Names only."""
return await self._inner.list_secrets(
list(validate_namespace_path(namespace_path))
)
async def drop_secret(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> None:
"""Drop a Secret. Bound Functions fail at their next job."""
await self._inner.drop_secret(
validate_secret_name(name), list(validate_namespace_path(namespace_path))
)
async def describe_secret(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> SecretInfo:
"""What this database records about a Secret. Never the value."""
name, created_at_millis, updated_at_millis = await self._inner.describe_secret(
validate_secret_name(name),
list(validate_namespace_path(namespace_path)),
)
return SecretInfo(
name=name,
created_at_millis=created_at_millis,
updated_at_millis=updated_at_millis,
)
async def list_jobs(self) -> List[JobInfo]:
"""List server-side jobs across the database's tables."""
return await self._inner.list_jobs()
async def get_job(self, job_id: str) -> Optional[JobDescription]:
"""Describe a single server-side job by id.
Returns None when the server has no such job.
"""
return await self._inner.get_job(job_id)
async def cancel_job(self, job_id: str) -> bool:
"""Request cancellation of a server-side job by id.
@@ -2263,12 +2670,46 @@ class AsyncConnection(object):
"""
return await self._inner.cancel_job(job_id)
async def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
"""The lifecycle event history of a server-side job, as Arrow batches.
async def execute_query(
self,
query: str,
*,
default_namespace_path: Optional[List[str]] = None,
) -> AsyncRecordBatchReader:
"""Execute SQL and return an asynchronous Arrow reader.
Lists history across all jobs when `job_id` is None.
This submits through :meth:`execute_query_async` and waits until the
initial result stream is readable. It does not wait for the full query
to finish.
"""
return await self._inner.job_history(job_id)
submitted = await self.execute_query_async(
query,
default_namespace_path=default_namespace_path,
)
return await submitted.reader()
async def execute_query_async(
self,
query: str,
*,
default_namespace_path: Optional[List[str]] = None,
) -> AsyncSqlQuery:
"""Start executing SQL and return its query handle.
The database from ``connect_async`` is used for unqualified database
references. The namespace defaults to ``["public"]``. Local
connections raise ``NotImplementedError``.
"""
return AsyncSqlQuery(
await self._inner.execute_query_async(
query,
default_namespace_path=default_namespace_path,
)
)
async def describe_query(self, query_id: UUID) -> QueryDescription:
"""Describe a submitted SQL query by its connection-scoped id."""
return await self._inner.describe_query(query_id)
async def namespace_client(self) -> LanceNamespace:
"""Get the equivalent namespace client for this connection.
python/python/lancedb/embeddings/gte.py
+1 -1
View File
@@ -21,7 +21,7 @@ class GteEmbeddings(TextEmbeddingFunction):
An embedding function that uses GTE-LARGE MLX format(for Apple silicon devices only)
as well as the standard cpu/gpu version from: https://huggingface.co/thenlper/gte-large.
For Apple users, you will need the mlx package insalled, which can be done with:
For Apple users, you will need the mlx package installed, which can be done with:
pip install mlx
Parameters
python/python/lancedb/embeddings/instructor.py
+1 -1
View File
@@ -60,7 +60,7 @@ class InstructorEmbeddingFunction(TextEmbeddingFunction):
import lancedb
from lancedb.pydantic import LanceModel, Vector
from lancedb.embeddings import get_registry, InstuctorEmbeddingFunction
from lancedb.embeddings import get_registry, InstructorEmbeddingFunction
instructor = get_registry().get("instructor").create(
source_instruction="represent the document for retrieval",
python/python/lancedb/embeddings/utils.py
+1 -1
View File
@@ -249,7 +249,7 @@ def retry_with_exponential_backoff(
initial_delay (float): Initial delay in seconds (default is 1).
exponential_base (float): The base for exponential backoff (default is 2).
jitter (bool): Whether to add jitter to the delay (default is True).
max_retries (int): Maximum number of retries (default is 10).
max_retries (int): Maximum number of retries (default is 7).
Returns:
function: The decorated function.
python/python/lancedb/exceptions.py
+6
View File
@@ -35,3 +35,9 @@ class JobCancelledError(RuntimeError):
"""Exception raised when an asynchronous job was cancelled."""
pass
class JobNotFoundError(ValueError):
"""Exception raised when opening a job the server does not have."""
pass
python/python/lancedb/functions.py
+489 -47
View File
@@ -4,7 +4,7 @@
"""Canonical Function values exchanged with LanceDB Enterprise services.
These immutable models contain client/wire state only. Catalog persistence,
environment bake, and execution are owned by Sophon.
environment bake, secret resolution, and execution are owned by Sophon.
``RefreshColumnResult`` is also the backend-neutral result of a local
expression-backed refresh job.
"""
@@ -25,7 +25,7 @@ import re
import sys
import textwrap
import types
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import date, datetime
from typing import (
Annotated,
@@ -41,6 +41,7 @@ from typing import (
import pyarrow as pa
from pydantic import (
AfterValidator,
BaseModel,
ConfigDict,
Field,
@@ -49,11 +50,26 @@ from pydantic import (
model_validator,
)
from .schema import is_blob_v2_field as _is_blob_v2_field
from .secrets import EnvVarSecret
_Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1)
_UInt32 = conint(strict=True, ge=0, le=2**32 - 1)
_UInt64 = conint(strict=True, ge=0, le=2**64 - 1)
def _validate_gpu_wire_marker(value: Any) -> bool:
if value is not True:
raise ValueError("runtime.gpu must be true")
return True
def _normalize_gpu_marker(value: bool) -> Optional[bool]:
if not isinstance(value, bool):
raise ValueError("gpu must be a boolean")
return True if value else None
class _FrozenDict(dict):
def _immutable(self, *args, **kwargs):
raise TypeError("remote canonical values are immutable")
@@ -212,6 +228,33 @@ class FunctionOutput(_OpenRemoteValue):
fields: tuple[FunctionResultField, ...] = ()
class SecretReference(_RemoteValue):
"""Where a Secret lives, carried as its parts rather than as one string.
A joined id would need a delimiter, and a delimiter has to be excluded from
every name and segment forever, agreed on by both sides, and re-agreed each
time either grows a new way to be configured. Naming the parts settles all
of that: nothing here is parsed, so nothing can parse two ways.
"""
name: str
namespace_path: tuple[str, ...] = ()
class SecretBinding(_RemoteValue):
"""How a Secret reaches the Function that binds it.
One list rather than a field per delivery mode: a binding is the concept,
and how it arrives is a property of one. ``kind`` is open, so a binding a
newer service introduces decodes here instead of failing the whole
FunctionVersion.
"""
kind: str
variable: Optional[str] = None
secret_ref: Optional[SecretReference] = None
class FunctionSignature(_RemoteValue):
inputs: tuple[FunctionParameter, ...]
output: FunctionOutput
@@ -239,6 +282,23 @@ class PythonRuntimeSpec(_RemoteValue):
python_version: Optional[str] = None
environment: Optional[PythonEnvironmentSpec] = None
env: Optional[Mapping[str, str]] = None
gpu: Optional[bool] = None
@model_validator(mode="before")
@classmethod
def _discard_unknown_runtime_payload(cls, value):
if isinstance(value, Mapping):
kind = value.get("kind")
if isinstance(kind, str) and kind not in {"python", "python_v2"}:
return {"kind": kind}
return value
@field_validator("gpu", mode="before")
@classmethod
def _validate_gpu_marker(cls, value):
if value is None:
return None
return _validate_gpu_wire_marker(value)
@model_validator(mode="after")
def _validate_runtime_kind(self):
@@ -247,28 +307,57 @@ class PythonRuntimeSpec(_RemoteValue):
raise ValueError("python runtime requires python_version")
if self.environment is None:
raise ValueError("python runtime requires environment")
if self.gpu is not None:
raise ValueError("python runtime with gpu requires kind='python_v2'")
elif self.kind == "python_v2":
if self.python_version is None:
raise ValueError("python_v2 runtime requires python_version")
if self.environment is None:
raise ValueError("python_v2 runtime requires environment")
if self.gpu is None:
raise ValueError("python_v2 runtime requires gpu")
else:
object.__setattr__(self, "python_version", None)
object.__setattr__(self, "environment", None)
object.__setattr__(self, "env", None)
object.__setattr__(self, "gpu", None)
return self
class FunctionVersion(_RemoteValue):
"""An exact immutable Function version returned by Enterprise.
class FunctionImage(_RemoteValue):
"""A complete OCI Function image identified by its exact manifest digest."""
Scheduling resources, priority, concurrency, and retry policy belong to
the submitting Job and are not part of this identity.
"""
manifest_digest: str
descriptor: Mapping[str, Any]
source: bool
def _validate_object_version(value: str) -> str:
if int(value) > 2**64 - 1:
raise ValueError("Function version exceeds uint64")
return value
_ObjectVersion = Annotated[
str,
Field(strict=True, pattern=r"^[1-9][0-9]*$"),
AfterValidator(_validate_object_version),
]
class FunctionVersion(_RemoteValue):
"""A pinned object revision, independent of its executable image digest."""
name: str
version: str
artifact: FunctionArtifact
object_id: str
location: str
version: _ObjectVersion
image: FunctionImage
signature: FunctionSignature
runtime: PythonRuntimeSpec
runtime_digest: str
environment_digest: str
secret_bindings: tuple[SecretBinding, ...] = ()
created_at: str
metadata: Mapping[str, str]
disabled: bool
def __call__(self, **inputs: Any) -> FunctionApplication:
"""Bind this exact version to named table columns.
@@ -322,24 +411,39 @@ class FunctionVersion(_RemoteValue):
)
)
return FunctionApplication(
function=FunctionVersionRef(name=self.name, version=self.version),
function=FunctionVersionRef(
name=self.name,
object_id=self.object_id,
location=self.location,
version=self.version,
manifest_digest=self.image.manifest_digest,
),
inputs=tuple(bindings),
output=self.signature.output,
)
class FunctionRegistrationRequest(_RemoteValue):
"""Stable remote registration envelope produced by :func:`udf`."""
"""Stable remote registration envelope produced by :func:`udf`.
Credential values deliberately have no field here. The only secret-shaped
thing a client sends is ``secret_bindings``: the name of a Secret the
database already holds, which the remote service resolves at execution.
"""
name: str
artifact: FunctionArtifactRequest
signature: FunctionSignature
runtime: PythonRuntimeSpec
secret_bindings: tuple[SecretBinding, ...] = ()
class FunctionVersionRef(_OpenRemoteValue):
name: str
version: str
object_id: str
location: str
version: _ObjectVersion
manifest_digest: str
class ApplicationInput(_OpenRemoteValue):
@@ -429,11 +533,7 @@ class InputBinding(_RemoteValue):
class OutputMapping(_RemoteValue):
"""One stable result-field mapping.
Assignment state is outside the Slice 1 client contract. During the NULL
transition Lance exposes no public cell-flag identifier to persist here.
"""
"""One stable result-field mapping."""
result_field: str
output_name: str
@@ -443,6 +543,13 @@ class OutputMapping(_RemoteValue):
nullable: bool
class AssignmentMapping(_RemoteValue):
"""Internal physical column preserving flattened struct validity."""
output_name: str
output_field_id: _Int32
class FunctionBinding(_RemoteValue):
"""Immutable Function binding persisted by the Enterprise table service."""
@@ -450,6 +557,7 @@ class FunctionBinding(_RemoteValue):
function: FunctionVersionRef
inputs: tuple[InputBinding, ...]
outputs: tuple[OutputMapping, ...]
assignment: Optional[AssignmentMapping] = None
input_schema: Optional[Mapping[str, Any]] = None
output_schema: Optional[Mapping[str, Any]] = None
@@ -480,6 +588,14 @@ class RefreshColumnResult(_RemoteValue):
_FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")
_FUNCTION_BLOB_V2_TYPE = "blob_v2"
_ARROW_EXTENSION_NAME_KEY = "ARROW:extension:name"
_BLOB_V2_EXTENSION_NAME = "lance.blob.v2"
_NESTED_BLOB_COLLECTION_ERROR = (
"unsupported Arrow type for Function signature: Blob v2 fields nested under "
"collection types are not supported"
)
_GRAMMAR_PRIMITIVES = (
(pa.bool_(), "bool"),
@@ -495,6 +611,7 @@ _GRAMMAR_PRIMITIVES = (
(pa.float32(), "float32"),
(pa.float64(), "float64"),
(pa.string(), "utf8"),
(pa.large_string(), "large_utf8"),
(pa.binary(), "binary"),
(pa.date32(), "date32"),
(pa.date64(), "date64"),
@@ -502,31 +619,258 @@ _GRAMMAR_PRIMITIVES = (
def _canonical_arrow_type(data_type: pa.DataType) -> str:
"""The server's V1 Function type grammar. Anything outside it is rejected
here rather than at registration."""
"""The compact Function grammar, or canonical exact JSON for nested types."""
grammar = _grammar_arrow_type(data_type)
if grammar is not None:
return grammar
exact = _exact_arrow_type(data_type)
return json.dumps(exact, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
def _grammar_arrow_type(data_type: pa.DataType) -> Optional[str]:
for candidate, name in _GRAMMAR_PRIMITIVES:
if data_type == candidate:
return name
if pa.types.is_list(data_type) or pa.types.is_large_list(data_type):
item = _grammar_list_item(data_type)
if item is None:
return None
prefix = "list" if pa.types.is_list(data_type) else "large_list"
return f"{prefix}<{_canonical_list_item(data_type)}>"
return f"{prefix}<{item}>"
if pa.types.is_fixed_size_list(data_type) and data_type.list_size > 0:
return (
f"fixed_size_list<{_canonical_list_item(data_type)}, {data_type.list_size}>"
)
raise TypeError(f"unsupported Arrow type for Function signature: {data_type}")
item = _grammar_list_item(data_type)
if item is not None:
return f"fixed_size_list<{item}, {data_type.list_size}>"
return None
def _canonical_list_item(data_type: pa.DataType) -> str:
def _grammar_list_item(data_type: pa.DataType) -> Optional[str]:
"""The grammar names only the item type; it always means a non-nullable
child called `item`, so any other child metadata cannot be represented."""
child called `item`, so other child properties require exact JSON."""
child = data_type.value_field
if child.name != "item" or child.nullable or child.metadata:
return None
return _grammar_arrow_type(child.type)
def _validate_exact_arrow_field(field: pa.Field) -> None:
if not field.name:
raise TypeError(
"unsupported Arrow type for Function signature: list items must be a "
f"non-nullable field named 'item', got {child}"
"unsupported Arrow type for Function signature: field names "
"must not be empty"
)
return _canonical_arrow_type(child.type)
if _is_blob_v2_field(field):
if not _has_supported_blob_v2_layout(field):
raise TypeError(
"unsupported Arrow type for Function signature: lance.blob.v2 "
f"requires a supported Blob storage layout, got {field}"
)
metadata = {
(key.decode() if isinstance(key, bytes) else key): (
value.decode() if isinstance(value, bytes) else value
)
for key, value in (field.metadata or {}).items()
}
if metadata and metadata != {
_ARROW_EXTENSION_NAME_KEY: _BLOB_V2_EXTENSION_NAME
}:
raise TypeError(
"unsupported Arrow type for Function signature: lance.blob.v2 "
"field metadata must contain only its canonical extension marker"
)
elif field.metadata:
raise TypeError(
"unsupported Arrow type for Function signature: field metadata "
f"is not supported, got {field}"
)
def _has_supported_blob_v2_layout(field: pa.Field) -> bool:
data_type = field.type
if isinstance(data_type, pa.ExtensionType):
data_type = data_type.storage_type
if not pa.types.is_struct(data_type):
return False
fields = tuple(data_type)
def matches(spec, compare_nullable) -> bool:
return len(fields) == len(spec) and all(
actual.name == name
and actual.type == expected_type
and (not check_nullable or actual.nullable == nullable)
for actual, (name, expected_type, nullable), check_nullable in zip(
fields, spec, compare_nullable
)
)
logical_minimal = (
("data", pa.large_binary(), True),
("uri", pa.utf8(), True),
)
logical_full = logical_minimal + (
("position", pa.uint64(), True),
("size", pa.uint64(), True),
)
prepared = (
("kind", pa.uint8(), True),
("data", pa.large_binary(), True),
("uri", pa.utf8(), True),
("blob_id", pa.uint32(), True),
("blob_size", pa.uint64(), True),
("position", pa.uint64(), True),
)
descriptor = (
("kind", pa.uint8(), False),
("position", pa.uint64(), False),
("size", pa.uint64(), False),
("blob_id", pa.uint32(), False),
("blob_uri", pa.utf8(), False),
)
return (
matches(logical_minimal, (True, True))
or matches(logical_full, (True, True, False, False))
or matches(prepared, (True,) * len(prepared))
or matches(descriptor, (False,) * len(descriptor))
)
def _canonical_arrow_field(field: pa.Field) -> str:
_validate_exact_arrow_field(field)
if _is_blob_v2_field(field):
return _FUNCTION_BLOB_V2_TYPE
return _canonical_arrow_type(field.type)
def _blob_storage_type(field: pa.Field) -> pa.DataType:
data_type = field.type
if isinstance(data_type, pa.ExtensionType):
return data_type.storage_type
return data_type
def _exact_blob_storage_type(field: pa.Field) -> dict[str, Any]:
storage = _blob_storage_type(field)
if not pa.types.is_struct(storage):
raise TypeError(
"unsupported Arrow type for Function signature: lance.blob.v2 "
"requires struct storage"
)
return {
"type": "struct",
"fields": [
{
"name": child.name,
"nullable": child.nullable,
"type": (
{"type": "large_binary"}
if pa.types.is_large_binary(child.type)
else _exact_arrow_type(child.type)
),
}
for child in storage
],
}
def _data_type_has_blob_v2(data_type: pa.DataType) -> bool:
if pa.types.is_struct(data_type):
return any(
_is_blob_v2_field(field) or _data_type_has_blob_v2(field.type)
for field in data_type
)
if (
pa.types.is_list(data_type)
or pa.types.is_large_list(data_type)
or pa.types.is_fixed_size_list(data_type)
):
field = data_type.value_field
return _is_blob_v2_field(field) or _data_type_has_blob_v2(field.type)
if pa.types.is_map(data_type):
return any(
_is_blob_v2_field(field) or _data_type_has_blob_v2(field.type)
for field in (data_type.key_field, data_type.item_field)
)
return False
def _exact_arrow_field(
field: pa.Field, *, inside_collection: bool = False
) -> dict[str, Any]:
_validate_exact_arrow_field(field)
if _is_blob_v2_field(field):
if inside_collection:
raise TypeError(_NESTED_BLOB_COLLECTION_ERROR)
return {
"name": field.name,
"nullable": field.nullable,
"type": _exact_blob_storage_type(field),
"metadata": {
_ARROW_EXTENSION_NAME_KEY: _BLOB_V2_EXTENSION_NAME,
},
}
value = {
"name": field.name,
"nullable": field.nullable,
"type": _exact_arrow_type(field.type, inside_collection=inside_collection),
}
return value
def _exact_arrow_type(
data_type: pa.DataType, *, inside_collection: bool = False
) -> dict[str, Any]:
for candidate, name in _GRAMMAR_PRIMITIVES:
if data_type == candidate:
return {"type": name}
if pa.types.is_struct(data_type):
fields = list(data_type)
names = [field.name for field in fields]
if not fields or len(set(names)) != len(names):
raise TypeError(
"unsupported Arrow type for Function signature: structs must have "
"non-empty, uniquely named fields"
)
return {
"type": "struct",
"fields": [
_exact_arrow_field(field, inside_collection=inside_collection)
for field in fields
],
}
if (
pa.types.is_list(data_type)
or pa.types.is_large_list(data_type)
or pa.types.is_fixed_size_list(data_type)
):
if pa.types.is_fixed_size_list(data_type):
if data_type.value_field.name != "item":
raise TypeError(
"unsupported Arrow type for Function signature: fixed-size list "
"items must be named 'item'"
)
if data_type.list_size <= 0:
raise TypeError(
f"unsupported Arrow type for Function signature: {data_type}"
)
value: dict[str, Any] = {
"type": (
"list"
if pa.types.is_list(data_type)
else "large_list"
if pa.types.is_large_list(data_type)
else "fixed_size_list"
),
"fields": [
_exact_arrow_field(data_type.value_field, inside_collection=True)
],
}
if pa.types.is_fixed_size_list(data_type):
value["length"] = data_type.list_size
return value
if pa.types.is_map(data_type) and _data_type_has_blob_v2(data_type):
raise TypeError(_NESTED_BLOB_COLLECTION_ERROR)
raise TypeError(f"unsupported Arrow type for Function signature: {data_type}")
def _list_of(item: pa.DataType) -> pa.DataType:
@@ -600,8 +944,15 @@ def _callable_parameters(function: Callable[..., Any]) -> tuple[inspect.Paramete
def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutput:
if isinstance(output, pa.Schema):
if output.metadata:
raise TypeError("Function output schema metadata is not supported")
fields = tuple(output)
elif isinstance(output, pa.Field) and pa.types.is_struct(output.type):
elif (
isinstance(output, pa.Field)
and not _is_blob_v2_field(output)
and pa.types.is_struct(output.type)
):
_validate_exact_arrow_field(output)
if output.nullable:
raise ValueError("Function output must be non-nullable")
fields = tuple(output.type)
@@ -617,18 +968,19 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp
raise TypeError(
"output_schema must be a PyArrow DataType, Field, or Schema"
)
_validate_exact_arrow_field(field)
if field.nullable:
raise ValueError("Function output must be non-nullable")
return FunctionOutput(
kind="scalar",
arrow_type=_canonical_arrow_type(field.type),
arrow_type=_canonical_arrow_field(field),
nullable=False,
)
if not fields:
raise ValueError("named-struct Function output must contain at least one field")
if any(field.nullable for field in fields):
raise ValueError("Function output fields must be non-nullable")
for field in fields:
_validate_exact_arrow_field(field)
names = [field.name for field in fields]
if len(set(names)) != len(names):
raise ValueError("Function output field names must be unique")
@@ -637,8 +989,8 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp
fields=tuple(
FunctionResultField(
name=field.name,
arrow_type=_canonical_arrow_type(field.type),
nullable=False,
arrow_type=_canonical_arrow_field(field),
nullable=field.nullable,
)
for field in fields
),
@@ -657,6 +1009,10 @@ def _infer_signature(
if input_schema is not None:
if not isinstance(input_schema, pa.Schema):
raise TypeError("input_schema must be a PyArrow Schema")
if input_schema.metadata:
raise TypeError("Function input schema metadata is not supported")
for field in input_schema:
_validate_exact_arrow_field(field)
expected = tuple(parameter.name for parameter in parameters)
actual = tuple(input_schema.names)
if actual != expected:
@@ -667,7 +1023,7 @@ def _infer_signature(
inputs = tuple(
FunctionParameter(
name=field.name,
arrow_type=_canonical_arrow_type(field.type),
arrow_type=_canonical_arrow_field(field),
nullable=field.nullable,
)
for field in input_schema
@@ -690,7 +1046,9 @@ def _infer_signature(
inputs.append(
FunctionParameter(
name=parameter.name,
arrow_type=_canonical_arrow_type(data_type),
arrow_type=_canonical_arrow_field(
pa.field(parameter.name, data_type, nullable=nullable)
),
nullable=nullable,
)
)
@@ -910,6 +1268,7 @@ class UdfDefinition:
pip: tuple[str, ...],
env: Mapping[str, str],
python_version: Optional[str],
gpu: bool = False,
conda: tuple[str, ...] = (),
conda_channels: tuple[str, ...] = (),
):
@@ -938,12 +1297,14 @@ class UdfDefinition:
signature = _infer_signature(function, input_schema, output_schema)
source = _package_source(function)
digest = f"sha256:{hashlib.sha256(source).hexdigest()}"
gpu_marker = _normalize_gpu_marker(gpu)
runtime = PythonRuntimeSpec(
kind="python",
kind="python_v2" if gpu_marker is not None else "python",
python_version=python_version
or f"{sys.version_info.major}.{sys.version_info.minor}",
environment=environment_spec,
env=environment,
gpu=gpu_marker,
)
self._function = function
self._request = FunctionRegistrationRequest(
@@ -968,9 +1329,73 @@ class UdfDefinition:
@property
def registration_request(self) -> FunctionRegistrationRequest:
"""The immutable request sent by ``create_function_async``."""
"""The immutable request sent by ``create_function_async``.
Carries no secret bindings. Binding is a registration-time decision,
so a Function bound to Secrets is registered through :meth:`bind_secrets`,
which is what ``create_function`` calls.
"""
return self._request
def bind_secrets(
self, secrets: Optional[Sequence[EnvVarSecret]]
) -> FunctionRegistrationRequest:
"""The registration request for this definition bound to ``secrets``.
Binding does not change the Function's source: each
[EnvVarSecret][lancedb.secrets.EnvVarSecret] names a Secret and the
environment variable its value should arrive in, and the Function reads
that variable the way it already did. Whether the named Secrets exist is
the server's answer, not this one.
"""
bindings = () if secrets is None else tuple(secrets)
wrong_type = [
binding for binding in bindings if not isinstance(binding, EnvVarSecret)
]
if wrong_type:
kinds = sorted({type(binding).__name__ for binding in wrong_type})
raise TypeError(
f"Function secrets must be EnvVarSecret values, not {kinds!r}; a "
"credential value is never sent to this API"
)
variables = [binding.env_variable for binding in bindings]
duplicates = sorted({name for name in variables if variables.count(name) > 1})
if duplicates:
raise ValueError(
"a Function binds each environment variable once; duplicated: "
f"{duplicates!r}"
)
# `env` is ordinary configuration carried in the definition, so a name in
# both would have a value visible in the Function's record and a value
# that is not. Refuse rather than pick.
environment = self._request.runtime.env or {}
overlap = sorted(set(environment) & set(variables))
if overlap:
raise ValueError(
f"Function env and secret bindings must be disjoint: {overlap!r}"
)
if not bindings:
return self._request
# Sorted, because the list is carried in the FunctionVersion hash and a
# caller's argument order is not part of what a Function is.
resolved = tuple(
sorted(
(
SecretBinding(
kind="env",
variable=binding.env_variable,
secret_ref=SecretReference(
name=binding.secret_name,
namespace_path=tuple(binding.secret_namespace_path),
),
)
for binding in bindings
),
key=lambda binding: (binding.kind, binding.variable or ""),
)
)
return self._request._copy(update={"secret_bindings": resolved})
def __call__(self, *args, **kwargs):
return self._function(*args, **kwargs)
@@ -989,6 +1414,7 @@ def udf(
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
python_version: Optional[str] = None,
gpu: bool = False,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
) -> Callable[[Callable[..., Any]], UdfDefinition]: ...
@@ -1003,6 +1429,7 @@ def udf(
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
python_version: Optional[str] = None,
gpu: bool = False,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
):
@@ -1010,8 +1437,9 @@ def udf(
Input and output signatures are inferred from supported annotations. For
Arrow types annotations cannot express precisely, pass ``input_schema``
and ``output_schema`` together. Nullable outputs are rejected because V1
uses physical NULL to represent unassigned computed-column rows.
and ``output_schema`` together. Scalar outputs must be non-nullable. Every
named-struct field may be nullable; Enterprise preserves the struct's
validity when the result is expanded into sibling columns.
Parameters
----------
@@ -1023,8 +1451,8 @@ def udf(
Explicit input fields in the exact order of the callable parameters.
Must be provided together with ``output_schema``.
output_schema : pyarrow.DataType, pyarrow.Field, or pyarrow.Schema, optional
Explicit scalar or named-struct output. Must be non-nullable and be
provided together with ``input_schema``.
Explicit scalar or named-struct output. Scalar outputs must be
non-nullable. Must be provided together with ``input_schema``.
pip : sequence of str, optional
Pip requirements for the remote environment.
conda : sequence of str, optional
@@ -1032,9 +1460,15 @@ def udf(
conda_channels : sequence of str, optional
Conda channels in priority order; requires ``conda``.
env : mapping of str to str, optional
Environment variables included in the Function definition.
Environment variables included in the Function definition. Not for
credentials -- these are ordinary configuration, stored with the
Function and visible wherever it is.
python_version : str, optional
Remote Python major/minor version. Defaults to the client version.
gpu : bool, default False
Whether every remote execution requires a GPU. The execution platform
selects one compatible GPU for each worker. The requirement is part of
the immutable Function version.
The packaged artifact is a snapshot: the function source plus exactly
the module-level names it references (modules as imports, importable
@@ -1059,6 +1493,11 @@ def udf(
... return value * 2
>>> score(1.5)
3.0
>>> @udf(pip=["cupy-cuda12x"], gpu=True)
... def gpu_score(value: int) -> int:
... return value * 2
>>> gpu_score.registration_request.runtime.gpu
True
"""
def decorate(target: Callable[..., Any]) -> UdfDefinition:
@@ -1070,6 +1509,7 @@ def udf(
pip=tuple(pip),
env={} if env is None else env,
python_version=python_version,
gpu=gpu,
conda=tuple(conda),
conda_channels=tuple(conda_channels),
)
@@ -1080,6 +1520,7 @@ def udf(
__all__ = [
"AssignmentMapping",
"ApplicationInput",
"FunctionApplication",
"FunctionArtifact",
@@ -1091,6 +1532,7 @@ __all__ = [
"FunctionRegistrationRequest",
"FunctionResultField",
"FunctionSignature",
"FunctionImage",
"FunctionVersion",
"FunctionVersionRef",
"InputBinding",
python/python/lancedb/index.py
+1 -1
View File
@@ -751,7 +751,7 @@ class IvfPq:
This value controls how much the vector is compressed during the
quantization step. The more sub vectors there are the less the vector is
compressed. The default is the dimension of the vector divided by 16. If
the dimension is not evenly divisible by 16 we use the dimension divded by
the dimension is not evenly divisible by 16 we use the dimension divided by
8.
The above two cases are highly preferred. Having 8 or 16 values per
python/python/lancedb/job.py
+224
View File
@@ -4,15 +4,27 @@
"""Handles to operations a server may run asynchronously."""
import asyncio
import json
from datetime import timedelta
from typing import Any, Callable, Generic, Optional, TypeVar, cast
import pyarrow as pa
from lancedb.background_loop import LOOP
from . import _lancedb
from ._lancedb import JobDescription, JobFailureInfo, JobInfo
T = TypeVar("T")
__all__ = [
"AsyncJob",
"Job",
"JobDescription",
"JobFailureInfo",
"JobInfo",
]
class AsyncJob(Generic[T]):
"""A handle to an operation that may still be running.
@@ -78,6 +90,149 @@ class AsyncJob(Generic[T]):
return
await self._inner.cancel()
async def refresh(self) -> None:
"""Ask the backend for this job's current state, and for a server-side
job its full record, then cache it for the properties below.
The properties are all `None` until this runs, because submitting an
operation returns only a job id. `status` fetches the whole record too;
`wait` records only the terminal state it establishes.
"""
if self._inner is None:
return
await self._inner.refresh()
@property
def state(self) -> Optional[str]:
"""The last observed lifecycle state, without contacting the backend.
`None` until the handle has talked to it. See :meth:`AsyncJob.refresh`.
"""
if self._inner is None:
return "finished"
return self._inner._state
@property
def job_type(self) -> Optional[str]:
"""The job's type, as the server names it.
`None` for an in-process job, which has no server-side record.
"""
return self._field("job_type")
@property
def creation_ms(self) -> Optional[int]:
"""When the job was created, in milliseconds since the epoch."""
return self._field("creation_ms")
@property
def spec(self) -> Optional[Any]:
"""The job-type-specific specification it was submitted with."""
return self._field("spec")
@property
def result(self) -> Optional[Any]:
"""The job-type-specific terminal result, as reported data rather than
the typed model :meth:`AsyncJob.wait` returns.
`None` until the job succeeds, so a job that never terminates reports
its progress through :meth:`AsyncJob.events` instead.
"""
return self._field("result")
@property
def failure(self) -> Optional[JobFailureInfo]:
"""Why the job failed, when it failed and the server reports a reason."""
return self._field("failure")
@property
def _spec_json(self) -> Optional[str]:
return self._field("_spec_json")
@property
def _result_json(self) -> Optional[str]:
return self._field("_result_json")
def _field(self, name: str) -> Optional[Any]:
description = self._inner._description if self._inner is not None else None
return getattr(description, name) if description is not None else None
async def events(
self,
*,
limit: Optional[int] = None,
filter: Optional[str] = None,
) -> "pa.Table":
"""This job's recorded lifecycle events.
Where the properties above report a terminal result only once the job
reaches one, events are written as the job runs and outlive the workers
that produced them. A distributed job records a `claim`/`claim_complete`
pair per unit of work, each carrying `rows_processed`, so a job that
never finishes still accounts for what it did.
Parameters
----------
limit: int, optional
Maximum event rows to return. The server caps results at 1000 by
default and 10,000 at most, and truncates without saying so, so
pass this for a job that emits an event per fragment.
filter: str, optional
SQL-like expression over the `state`, `updated_by`, `emitted_from`,
`emitted_by`, and `claim_entity` columns, such as
``state = 'claim_complete'``.
"""
if self._inner is None:
raise NotImplementedError(
"job event history is only available for server-side jobs"
)
return await self._inner.events(limit=limit, filter=filter)
def __repr__(self) -> str:
return _job_repr("AsyncJob", self)
_REPR_INDENT = " " * 4
def _repr_payload(value: Any) -> str:
"""Render a job payload as indented JSON, aligned under its field."""
try:
rendered = json.dumps(value, indent=4)
except TypeError:
return repr(value)
return rendered.replace("\n", "\n" + _REPR_INDENT)
def _job_repr(kind: str, job: Any) -> str:
"""Render every field the handle currently knows, omitting the rest.
One field per line, with the JSON payloads indented, because a refresh
job's spec and result are the point of printing it.
"""
state = job.state
if state is None:
# Nothing has been fetched yet, so there is nothing to lay out.
known = f"id={job.id!r}, " if job.id is not None else ""
return f"{kind}({known}not refreshed)"
fields = []
if job.id is not None:
fields.append(f"id={job.id!r}")
fields.append(f"state={state!r}")
for name in ("job_type", "creation_ms"):
value = getattr(job, name)
if value is not None:
fields.append(f"{name}={value!r}")
for name in ("spec", "result"):
value = getattr(job, name)
if value is not None:
fields.append(f"{name}={_repr_payload(value)}")
if job.failure is not None:
fields.append(f"failure={job.failure!r}")
body = "".join(f"\n{_REPR_INDENT}{field}," for field in fields)
return f"{kind}({body}\n)"
class Job(Generic[T]):
"""Synchronous counterpart of `AsyncJob` with the same result type."""
@@ -122,6 +277,75 @@ class Job(Generic[T]):
return
LOOP.run(self._inner.cancel())
def refresh(self) -> None:
"""Ask the backend for this job's current state and record.
See :meth:`AsyncJob.refresh`.
"""
if self._inner is None:
return
LOOP.run(self._inner.refresh())
@property
def state(self) -> Optional[str]:
"""The last observed lifecycle state. See :attr:`AsyncJob.state`."""
return self._inner.state if self._inner is not None else "finished"
@property
def job_type(self) -> Optional[str]:
"""The job's type. See :attr:`AsyncJob.job_type`."""
return self._field("job_type")
@property
def creation_ms(self) -> Optional[int]:
"""When the job was created. See :attr:`AsyncJob.creation_ms`."""
return self._field("creation_ms")
@property
def spec(self) -> Optional[Any]:
"""The job's specification. See :attr:`AsyncJob.spec`."""
return self._field("spec")
@property
def result(self) -> Optional[Any]:
"""The job's terminal result. See :attr:`AsyncJob.result`."""
return self._field("result")
@property
def failure(self) -> Optional[JobFailureInfo]:
"""Why the job failed. See :attr:`AsyncJob.failure`."""
return self._field("failure")
@property
def _spec_json(self) -> Optional[str]:
return self._field("_spec_json")
@property
def _result_json(self) -> Optional[str]:
return self._field("_result_json")
def _field(self, name: str) -> Optional[Any]:
return getattr(self._inner, name) if self._inner is not None else None
def events(
self,
*,
limit: Optional[int] = None,
filter: Optional[str] = None,
) -> "pa.Table":
"""This job's recorded lifecycle events.
See :meth:`AsyncJob.events`.
"""
if self._inner is None:
raise NotImplementedError(
"job event history is only available for server-side jobs"
)
return LOOP.run(self._inner.events(limit=limit, filter=filter))
def __repr__(self) -> str:
return _job_repr("Job", self)
def _typed_job(
inner: "_lancedb.Job", result_decoder: Callable[[str], T]
python/python/lancedb/materialized_view.py
+57 -5
View File
@@ -11,6 +11,7 @@ from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
from .background_loop import LOOP
from .job import AsyncJob, Job, _typed_job
if TYPE_CHECKING:
import pyarrow as pa
@@ -42,6 +43,8 @@ class MaterializedViewDefinition:
"""Cap on the number of rows the view holds."""
inputs: List[str] = field(default_factory=list)
"""Source columns the projections and filter read."""
source_namespace: List[str] = field(default_factory=list)
"""Namespace holding the source table; empty is the root namespace."""
def _definition_from_schema(
@@ -53,7 +56,8 @@ def _definition_from_schema(
raise ValueError(f"Table '{name}' is not a materialized view")
value = json.loads(raw)
kind = value.get("kind")
if kind != "select":
# "namespaced_select" keeps older readers from resolving the source at root.
if kind not in ("select", "namespaced_select"):
raise NotImplementedError(
f"materialized view '{name}' is defined by '{kind}', which this "
"version of lancedb cannot refresh"
@@ -66,6 +70,21 @@ def _definition_from_schema(
filter=value.get("filter"),
limit=value.get("limit"),
inputs=value.get("inputs", []),
source_namespace=value.get("source_namespace", []),
)
def _definition_from_json(raw: str) -> MaterializedViewDefinition:
value = json.loads(raw)
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", []),
source_namespace=value.get("source_namespace", []),
)
@@ -122,8 +141,9 @@ class AsyncMaterializedView:
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)
"""The query that defines the view."""
raw = await self._table._inner.materialized_view_definition()
return _definition_from_json(raw)
async def refresh(
self, *, full: bool = False, source_version: Optional[int] = None
@@ -144,6 +164,24 @@ class AsyncMaterializedView:
full=full, source_version=source_version
)
async def refresh_async(
self, *, full: bool = False, source_version: Optional[int] = None
) -> "AsyncJob[RefreshMaterializedViewResult]":
"""Submit a refresh and return its job without waiting.
The job may already be complete for a local view. On LanceDB Cloud
and Enterprise, its ``id`` is the server job identifier returned by
the refresh endpoint.
"""
from ._lancedb import RefreshMaterializedViewResult
return _typed_job(
await self._table._inner.refresh_materialized_view_async(
full=full, source_version=source_version
),
RefreshMaterializedViewResult.from_json,
)
class MaterializedView:
"""Synchronous variant of
@@ -167,8 +205,8 @@ class MaterializedView:
@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)
"""The query that defines the view."""
return LOOP.run(self._async.definition())
def refresh(
self, *, full: bool = False, source_version: Optional[int] = None
@@ -176,3 +214,17 @@ class MaterializedView:
"""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))
def refresh_async(
self, *, full: bool = False, source_version: Optional[int] = None
) -> "Job[RefreshMaterializedViewResult]":
"""Submit a refresh and return its job without waiting.
See
[AsyncMaterializedView.refresh_async][lancedb.materialized_view.AsyncMaterializedView.refresh_async].
"""
return Job(
LOOP.run(
self._async.refresh_async(full=full, source_version=source_version)
)
)
python/python/lancedb/namespace.py
+134 -2
View File
@@ -12,6 +12,7 @@ from __future__ import annotations
import sys
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
from uuid import UUID
if sys.version_info >= (3, 12):
from typing import override
@@ -48,8 +49,11 @@ from lancedb._lancedb import (
connect_namespace_client as _connect_namespace_client,
)
from lancedb.background_loop import LOOP
from lancedb.arrow import AsyncRecordBatchReader
from lancedb.db import AsyncConnection, DBConnection
from lancedb.job import AsyncJob, Job
from lancedb.sql import AsyncQuery as AsyncSqlQuery
from lancedb.sql import QueryDescription
from lance_namespace import (
LanceNamespace,
connect as namespace_connect,
@@ -633,6 +637,7 @@ class LanceNamespaceDBConnection(DBConnection):
select: "SelectArg" = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> "MaterializedView":
"""Define a materialized view over a table in the root namespace.
See
@@ -642,12 +647,40 @@ class LanceNamespaceDBConnection(DBConnection):
self.open_table(
LOOP.run(
self._inner.create_materialized_view(
name, source, select=select, where=where, limit=limit
name,
source,
select=select,
where=where,
limit=limit,
with_no_data=with_no_data,
)
).name
)
)
@override
def create_materialized_view_async(
self,
name: str,
source: str,
*,
select: "SelectArg" = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> Job[None]:
job = LOOP.run(
self._inner.create_materialized_view_async(
name,
source,
select=select,
where=where,
limit=limit,
with_no_data=with_no_data,
)
)
return Job(job)
@override
def open_materialized_view(self, name: str) -> "MaterializedView":
"""Open the materialized view named ``name``."""
@@ -660,6 +693,30 @@ class LanceNamespaceDBConnection(DBConnection):
"""The names of the materialized views in the root namespace."""
return LOOP.run(self._inner.list_materialized_views())
@override
def drop_materialized_view(
self, name: str, namespace_path: Optional[List[str]] = None
) -> None:
if namespace_path is None:
namespace_path = []
LOOP.run(
self._inner.drop_materialized_view(name, namespace_path=namespace_path)
)
@override
def drop_materialized_view_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Job[None]:
if namespace_path is None:
namespace_path = []
return Job(
LOOP.run(
self._inner.drop_materialized_view_async(
name, namespace_path=namespace_path
)
)
)
@override
def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
if namespace_path is None:
@@ -1190,15 +1247,41 @@ class AsyncLanceNamespaceDBConnection:
select: "SelectArg" = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> "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
name,
source,
select=select,
where=where,
limit=limit,
with_no_data=with_no_data,
)
# 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 create_materialized_view_async(
self,
name: str,
source: str,
*,
select: "SelectArg" = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> AsyncJob[None]:
"""Submit materialized-view creation and return its job."""
return await self._inner.create_materialized_view_async(
name,
source,
select=select,
where=where,
limit=limit,
with_no_data=with_no_data,
)
async def open_materialized_view(self, name: str) -> "AsyncMaterializedView":
"""Open the materialized view named ``name``."""
view = AsyncMaterializedView(await self.open_table(name))
@@ -1209,6 +1292,24 @@ class AsyncLanceNamespaceDBConnection:
"""The names of the materialized views in the root namespace."""
return await self._inner.list_materialized_views()
async def drop_materialized_view(
self, name: str, namespace_path: Optional[List[str]] = None
) -> None:
"""Drop a materialized view from the namespace."""
if namespace_path is None:
namespace_path = []
await self._inner.drop_materialized_view(name, namespace_path=namespace_path)
async def drop_materialized_view_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> AsyncJob[None]:
"""Start dropping a materialized view and return its cleanup job."""
if namespace_path is None:
namespace_path = []
return await self._inner.drop_materialized_view_async(
name, namespace_path=namespace_path
)
async def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
"""Drop a table from the namespace."""
if namespace_path is None:
@@ -1447,6 +1548,37 @@ class AsyncLanceNamespaceDBConnection:
namespace_path=namespace_path, page_token=page_token, limit=limit
)
async def execute_query(
self,
query: str,
*,
default_namespace_path: Optional[List[str]] = None,
) -> AsyncRecordBatchReader:
"""Execute SQL when supported by the underlying connection."""
return await self._inner.execute_query(
query,
default_namespace_path=default_namespace_path,
)
async def execute_query_async(
self,
query: str,
*,
default_namespace_path: Optional[List[str]] = None,
) -> AsyncSqlQuery:
"""Start executing SQL when supported by the underlying connection.
Namespace-backed local connections do not support SQL.
"""
return await self._inner.execute_query_async(
query,
default_namespace_path=default_namespace_path,
)
async def describe_query(self, query_id: UUID) -> QueryDescription:
"""Describe a submitted SQL query when supported."""
return await self._inner.describe_query(query_id)
async def namespace_client(self) -> LanceNamespace:
"""Get the namespace client for this connection.
python/python/lancedb/query.py
+64 -17
View File
@@ -78,6 +78,10 @@ if TYPE_CHECKING:
T = TypeVar("T", bound="LanceModel")
AnalyzePlanDistributedMetrics = Literal["aggregate", "per_worker", "full"]
# Number of rows a hybrid query returns when no limit was set on it. This
# mirrors the default the Rust query builder applies to its sub-queries.
DEFAULT_HYBRID_LIMIT = 10
@runtime_checkable
class _LanceScanner(Protocol):
@@ -109,6 +113,7 @@ def _query_is_plain_scan(query: Query) -> bool:
return (
query.vector is None
and query.full_text_query is None
and query.take_offsets is None
and not query.postfilter
and not query.order_by
)
@@ -804,6 +809,10 @@ class Query(pydantic.BaseModel):
# offset to start fetching results from
offset: Optional[int] = None
# Dataset offsets whose duplicate occurrences must be restored after lookup.
# This is populated when a take query is converted to this serializable form.
take_offsets: Optional[List[int]] = None
# if true, will only search the indexed data
fast_search: Optional[bool] = None
@@ -825,6 +834,7 @@ class Query(pydantic.BaseModel):
query = cls()
query.limit = req.limit
query.offset = req.offset
query.take_offsets = req.take_offsets
query.filter = req.filter
query.full_text_query = req.full_text_search
query.columns = req.select
@@ -853,7 +863,7 @@ class Query(pydantic.BaseModel):
return query
# This tells pydantic to allow custom types (needed for the `vector` query since
# pa.Array wouln't be allowed otherwise)
# pa.Array wouldn't be allowed otherwise)
model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
@@ -3887,14 +3897,54 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
return self
def _create_child_queries(
self,
) -> Tuple["AsyncFTSQuery", "AsyncVectorQuery", int, int]:
"""Build the sub-queries that make up this hybrid query.
Execution, `explain_plan` and `analyze_plan` all go through here so that
the plans that are reported are the plans that actually run.
Returns the two sub-queries along with the effective limit and offset of
the hybrid query itself.
"""
fts_query = AsyncFTSQuery(self._inner.to_fts_query(), self._table)
vec_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table)
fts_req = fts_query._inner.to_query_request()
vec_req = vec_query._inner.to_query_request()
# Only one of the two sub-queries carries the limit when it was never
# set explicitly: nearest_to()/nearest_to_text() build the sibling query
# from scratch, and that is where the default gets filled in. Which one
# that is depends on the order the hybrid query was built in, so look at
# both rather than at a single side.
limit = fts_req.limit if fts_req.limit is not None else vec_req.limit
if limit is None:
limit = DEFAULT_HYBRID_LIMIT
offset = fts_req.offset or vec_req.offset or 0
fts_query.with_row_id()
vec_query.with_row_id()
# offset() pushes the offset down into both sub-queries, which would make
# each of them skip its own first `offset` rows. The window has to be
# taken out of the combined, reranked results instead, so fetch the
# skipped prefix here too and slice it off afterwards.
fts_query.limit(limit + offset)
vec_query.limit(limit + offset)
fts_query.offset(0)
vec_query.offset(0)
return fts_query, vec_query, limit, offset
async def to_batches(
self,
*,
max_batch_length: Optional[int] = None,
timeout: Optional[timedelta] = None,
) -> AsyncRecordBatchReader:
fts_query = AsyncFTSQuery(self._inner.to_fts_query(), self._table)
vec_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table)
fts_query, vec_query, limit, offset = self._create_child_queries()
req = fts_query._inner.to_query_request()
blob_auto_row_id = False
@@ -3914,9 +3964,6 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
self._blob_auto_row_id = blob_auto_row_id
self._blob_paths = blob_paths
fts_query.with_row_id()
vec_query.with_row_id()
fts_results, vector_results = await asyncio.gather(
fts_query.to_arrow(timeout=timeout),
vec_query.to_arrow(timeout=timeout),
@@ -3928,8 +3975,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
norm=self._norm,
fts_query=fts_query.get_query(),
reranker=self._reranker,
limit=self._inner.get_limit(),
limit=limit,
with_row_ids=True,
offset=offset,
)
if (
not self._user_requested_row_id()
@@ -3958,14 +4006,14 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
... print(plan)
>>> asyncio.run(doctest_example()) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
RRFReranker(K=60)
ProjectionExec: expr=[vector@0 as vector, text@3 as text, _distance@2 as _distance]
ProjectionExec: expr=[vector@0 as vector, text@3 as text, _distance@2 as _distance, _rowid@1 as _rowid]
LanceRead: uri=..., projection=[text], source=stream(_rowid)
GlobalLimitExec: skip=0, fetch=10
FilterExec: _distance@2 IS NOT NULL
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
KNNVectorDistance: metric=l2
LanceRead: uri=..., projection=[vector], ...
ProjectionExec: expr=[vector@2 as vector, text@3 as text, _score@1 as _score]
ProjectionExec: expr=[vector@2 as vector, text@3 as text, _score@1 as _score, _rowid@0 as _rowid]
LanceRead: uri=..., projection=[vector, text], source=stream(_rowid)
GlobalLimitExec: skip=0, fetch=10
MatchQuery: column=text, query=[hello]
@@ -3980,8 +4028,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
plan : str
""" # noqa: E501
vector_plan = await self._inner.to_vector_query().explain_plan(verbose)
fts_plan = await self._inner.to_fts_query().explain_plan(verbose)
fts_query, vec_query, _, _ = self._create_child_queries()
vector_plan = await vec_query.explain_plan(verbose)
fts_plan = await fts_query.explain_plan(verbose)
# Indent sub-plans under the reranker
indented_vector = "\n".join(" " + line for line in vector_plan.splitlines())
indented_fts = "\n".join(" " + line for line in fts_plan.splitlines())
@@ -4008,14 +4057,12 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
-------
plan : str
"""
fts_query, vec_query, _, _ = self._create_child_queries()
results = ["Vector Search Query:"]
results.append(
await self._inner.to_vector_query().analyze_plan(distributed_metrics)
)
results.append(await vec_query.analyze_plan(distributed_metrics))
results.append("FTS Search Query:")
results.append(
await self._inner.to_fts_query().analyze_plan(distributed_metrics)
)
results.append(await fts_query.analyze_plan(distributed_metrics))
return "\n".join(results)
python/python/lancedb/remote/__init__.py
+14 -2
View File
@@ -9,7 +9,13 @@ from typing import List, Optional
from lancedb import __version__
from .header import HeaderProvider
from .oauth import OAuthConfig, OAuthFlowType
from .oauth import (
ClientAuthMethod,
OAuthConfig,
OAuthFlowType,
OAuthSession,
TokenCacheOptions,
)
# The API reference renders this module with a single mkdocstrings directive,
# which only picks up names listed here. New public names must be added to this
@@ -22,6 +28,9 @@ __all__ = [
"HeaderProvider",
"OAuthConfig",
"OAuthFlowType",
"ClientAuthMethod",
"OAuthSession",
"TokenCacheOptions",
]
@@ -164,7 +173,10 @@ class ClientConfig:
extra_headers: Optional[dict]
Additional headers to include in requests.
id_delimiter: Optional[str]
The delimiter to use when constructing object identifiers.
The delimiter joining a namespace path and a name into one object
identifier. ``"$"`` is the only supported value, and leaving this
unset is how to get it; anything else is rejected when the connection
is created.
tls_config: Optional[TlsConfig]
TLS/mTLS configuration for secure connections.
header_provider: Optional[HeaderProvider]
python/python/lancedb/remote/db.py
+200 -30
View File
@@ -2,13 +2,24 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from dataclasses import replace
from datetime import timedelta
import json
import logging
from concurrent.futures import ThreadPoolExecutor
import sys
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
from typing import (
TYPE_CHECKING,
Any,
Dict,
Iterable,
List,
Optional,
Sequence,
Union,
)
from urllib.parse import urlparse
from uuid import UUID
import warnings
if sys.version_info >= (3, 12):
@@ -25,10 +36,13 @@ from ..common import DATA
from ..db import DBConnection, LOOP
from ..functions import FunctionVersion, UdfDefinition
from ..job import AsyncJob, Job
from ..sql import Query as SqlQuery
from ..sql import QueryDescription
from ..materialized_view import MaterializedView, SelectArg
from ..secrets import EnvVarSecret, SecretInfo
if TYPE_CHECKING:
from .._lancedb import JobDescription, JobInfo
from .._lancedb import JobInfo
from ..embeddings import EmbeddingFunctionConfig
from lance_namespace import (
LanceNamespace,
@@ -116,6 +130,7 @@ class RemoteDBConnection(DBConnection):
read_timeout: Optional[float] = None,
storage_options: Optional[Dict[str, str]] = None,
read_consistency_interval: Optional[timedelta] = None,
sql_host_override: Optional[str] = None,
):
"""Connect to a remote LanceDB database."""
if isinstance(client_config, dict):
@@ -161,6 +176,7 @@ class RemoteDBConnection(DBConnection):
self.api_key = api_key
self.region = region
self.host_override = host_override
self.sql_host_override = sql_host_override
self.storage_options = storage_options
self.db_name = parsed.netloc
@@ -175,17 +191,58 @@ class RemoteDBConnection(DBConnection):
api_key=api_key,
region=region,
host_override=host_override,
sql_host_override=sql_host_override,
client_config=client_config,
storage_options=storage_options,
read_consistency_interval=read_consistency_interval,
)
)
@classmethod
def _from_catalog(
cls,
inner,
name,
endpoint,
api_key,
client_config,
oauth_config,
sql_host_override,
):
config = (
ClientConfig(**client_config)
if isinstance(client_config, dict)
else (client_config or ClientConfig())
)
headers = {
key: value
for key, value in (config.extra_headers or {}).items()
if key.lower() not in ("x-lancedb-database", "x-lancedb-database-prefix")
}
headers["x-lancedb-database"] = name
result = cls.__new__(cls)
result.db_url = inner.uri
result.db_name = name
result.api_key = api_key or ""
result.region = "us-east-1"
result.host_override = endpoint
result.sql_host_override = sql_host_override
result.storage_options = None
result.client_config = replace(config, extra_headers=headers)
result._catalog_oauth = oauth_config is not None
result._conn = inner
return result
def __repr__(self) -> str:
return f"RemoteConnect(name={self.db_name})"
@override
def serialize(self) -> str:
if getattr(self, "_catalog_oauth", False):
raise ValueError(
"Cannot serialize a catalog connection using OAuth; "
"provide a worker-side connection factory"
)
return json.dumps(
{
"connection_type": "remote",
@@ -193,6 +250,7 @@ class RemoteDBConnection(DBConnection):
"api_key": self.api_key,
"region": self.region,
"host_override": self.host_override,
"sql_host_override": self.sql_host_override,
"client_config": _client_config_to_dict(self.client_config),
"storage_options": self.storage_options,
}
@@ -658,22 +716,80 @@ class RemoteDBConnection(DBConnection):
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> MaterializedView:
raise NotImplementedError(
"materialized views are supported only on local databases"
from .table import RemoteTable
view = LOOP.run(
self._conn.create_materialized_view(
name,
source,
select=select,
where=where,
limit=limit,
with_no_data=with_no_data,
)
)
return MaterializedView(
RemoteTable(
view.table,
self.db_name,
connection_state=self.serialize,
namespace_path=[],
)
)
@override
def create_materialized_view_async(
self,
name: str,
source: str,
*,
select: SelectArg = None,
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
) -> Job[None]:
job = LOOP.run(
self._conn.create_materialized_view_async(
name,
source,
select=select,
where=where,
limit=limit,
with_no_data=with_no_data,
)
)
return Job(job)
@override
def open_materialized_view(self, name: str) -> MaterializedView:
raise NotImplementedError(
"materialized views are supported only on local databases"
)
view = MaterializedView(self.open_table(name))
view.definition
return view
@override
def list_materialized_views(self) -> List[str]:
raise NotImplementedError(
"materialized views are supported only on local databases"
return LOOP.run(self._conn.list_materialized_views())
@override
def drop_materialized_view(
self, name: str, namespace_path: Optional[List[str]] = None
) -> None:
if namespace_path is None:
namespace_path = []
LOOP.run(self._conn.drop_materialized_view(name, namespace_path=namespace_path))
@override
def drop_materialized_view_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Job[None]:
if namespace_path is None:
namespace_path = []
job = LOOP.run(
self._conn.drop_materialized_view_async(name, namespace_path=namespace_path)
)
return Job(job)
@override
def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
@@ -732,36 +848,67 @@ class RemoteDBConnection(DBConnection):
)
@override
def job(self, job_id: str) -> Job:
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
The handle is constructed without a server round trip; an unknown id
surfaces when the handle is used. Dropping the handle has no effect
on the job itself.
def open_job(self, job_id: str) -> Job:
"""Open a server-side job by id. See
[DBConnection.open_job][lancedb.db.DBConnection.open_job].
"""
return Job(self._conn.job(job_id))
return Job(LOOP.run(self._conn.open_job(job_id)))
@override
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
return Job(LOOP.run(self._conn.create_function_async(definition)))
def create_function_async(
self,
definition: UdfDefinition,
*,
secrets: Optional[Sequence[EnvVarSecret]] = None,
) -> Job[FunctionVersion]:
job = LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
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_functions(self) -> List[FunctionVersion]:
return LOOP.run(self._conn.list_functions())
@override
def drop_function(self, name: str, *, version: str) -> bool:
return LOOP.run(self._conn.drop_function(name, version=version))
@override
def create_secret(
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
) -> None:
LOOP.run(self._conn.create_secret(name, value, namespace_path=namespace_path))
@override
def alter_secret(
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
) -> None:
LOOP.run(self._conn.alter_secret(name, value, namespace_path=namespace_path))
@override
def describe_secret(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> SecretInfo:
return LOOP.run(self._conn.describe_secret(name, namespace_path=namespace_path))
@override
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
return LOOP.run(self._conn.list_secrets(namespace_path=namespace_path))
@override
def drop_secret(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> None:
LOOP.run(self._conn.drop_secret(name, namespace_path=namespace_path))
@override
def list_jobs(self) -> List["JobInfo"]:
"""List server-side jobs across the database's tables."""
return LOOP.run(self._conn.list_jobs())
@override
def get_job(self, job_id: str) -> Optional["JobDescription"]:
"""Describe a single server-side job by id.
Returns None when the server has no such job.
"""
return LOOP.run(self._conn.get_job(job_id))
@override
def cancel_job(self, job_id: str) -> bool:
"""Request cancellation of a server-side job by id.
@@ -773,12 +920,35 @@ class RemoteDBConnection(DBConnection):
return LOOP.run(self._conn.cancel_job(job_id))
@override
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
"""The lifecycle event history of a server-side job, as Arrow batches.
def execute_query_async(
self,
query: str,
*,
default_namespace_path: Optional[List[str]] = None,
) -> SqlQuery:
"""Start executing SQL through this remote connection.
Lists history across all jobs when `job_id` is None.
Unqualified tables use this connection's database and the
``["public"]`` namespace by default. Fully qualified table names may
reference other databases available to the same deployment.
"""
return LOOP.run(self._conn.job_history(job_id))
return SqlQuery(
LOOP.run(
self._conn.execute_query_async(
query,
default_namespace_path=default_namespace_path,
)
)
)
@override
def describe_query(self, query_id: UUID) -> QueryDescription:
"""Describe a submitted SQL query by its connection-scoped id."""
return LOOP.run(
self._conn.describe_query(
query_id,
)
)
@override
def namespace_client(self) -> LanceNamespace:
python/python/lancedb/remote/header.py
+4 -1
View File
@@ -177,4 +177,7 @@ class OAuthProvider(HeaderProvider):
if not self._current_token:
raise RuntimeError("Failed to obtain OAuth token")
return {"Authorization": f"Bearer {self._current_token}"}
return {
"Authorization": f"Bearer {self._current_token}",
"x-lancedb-credential-type": "oidc",
}
python/python/lancedb/remote/oauth.py
+185
View File
@@ -12,10 +12,72 @@ class OAuthFlowType(str, Enum):
CLIENT_CREDENTIALS = "client_credentials"
"""Client Credentials grant (service-to-service / M2M)."""
AUTHORIZATION_CODE = "authorization_code"
"""Interactive Authorization Code grant, using PKCE by default."""
DEVICE_CODE = "device_code"
"""Device Authorization grant for CLI and headless environments."""
AZURE_MANAGED_IDENTITY = "azure_managed_identity"
"""Azure Managed Identity via IMDS."""
class ClientAuthMethod(str, Enum):
"""How the client authenticates to the OAuth token endpoint.
The method applies to every OAuth request that carries client
authentication: client-credentials, authorization-code exchange,
refresh-token, and device-authorization requests. The Azure managed
identity flow ignores this option.
"""
NONE = "none"
"""No client authentication, for public clients using PKCE or the device
flow. Cannot be combined with ``client_secret``."""
CLIENT_SECRET_BASIC = "client_secret_basic"
"""HTTP Basic authentication. This is the RFC 6749 recommended method and
the normal default for confidential clients, including default Okta
applications. Requires ``client_secret``."""
CLIENT_SECRET_POST = "client_secret_post"
"""Credentials in the request body, for providers configured to require
it. Requires ``client_secret``."""
@dataclass
class TokenCacheOptions:
"""Options for the persistent OAuth token cache.
The cache is opt-in: it is only used when set as ``token_cache`` on
:class:`OAuthConfig`. Only refresh tokens are persisted, in a private
directory with owner-only permissions, so short-lived processes can reuse
an authenticated session instead of re-prompting on every start.
Parameters
----------
cache_dir : Optional[str]
Directory that holds cached credentials. Defaults to
``$XDG_CACHE_HOME/lancedb/oauth``, ``$HOME/.cache/lancedb/oauth`` on
Unix, or ``%LOCALAPPDATA%\\lancedb\\oauth`` on Windows. The directory
is created with owner-only permissions (``0700``) when missing.
lock_timeout_secs : Optional[int]
How long to wait for the cross-process refresh lock before failing
(default: 30 seconds).
Examples
--------
>>> opts = TokenCacheOptions(cache_dir="/tmp/my-app/oauth-cache")
Multiple identities (issuer, client, scopes, flow, client
authentication) get separate cache entries. Within one identity the most
recent login wins.
"""
cache_dir: Optional[str] = None
lock_timeout_secs: Optional[int] = None
@dataclass
class OAuthConfig:
"""OAuth configuration for LanceDB authentication.
@@ -38,12 +100,37 @@ class OAuthConfig:
Authentication flow to use. Default: CLIENT_CREDENTIALS.
client_secret : Optional[str]
Client secret (required for CLIENT_CREDENTIALS).
client_auth_method : Optional[ClientAuthMethod]
How the client authenticates to the token endpoint (default: auto).
With a ``client_secret`` the default is
``ClientAuthMethod.CLIENT_SECRET_BASIC``, which matches the RFC 6749
recommendation and the default configuration of Okta confidential
applications; without a secret the client is public and no client
authentication is sent.
redirect_uri : Optional[str]
Loopback redirect URI for AUTHORIZATION_CODE. The default is
``http://127.0.0.1:{callback_port}/callback``.
callback_port : Optional[int]
Port for the AUTHORIZATION_CODE loopback callback server (default: 8400).
use_pkce : bool
Protect AUTHORIZATION_CODE with S256 PKCE (default: True).
managed_identity_client_id : Optional[str]
Client ID for user-assigned managed identity (AZURE_MANAGED_IDENTITY).
resource : Optional[str]
Resource indicator (RFC 8707), forwarded verbatim to authorization and
token endpoints, including refresh requests. Must be an absolute URI
without a fragment. Not supported for Azure managed identity.
audience : Optional[str]
Provider-specific audience, forwarded to authorization and token
endpoints, including refresh requests. Not supported for Azure managed identity.
refresh_buffer_secs : Optional[int]
Seconds before expiry to trigger proactive refresh (default: 300).
Keep this well below the token TTL; if it is greater than or equal to
the TTL, each request refreshes the token.
token_cache : Optional[TokenCacheOptions]
Opt in to the persistent token cache so short-lived processes reuse
one session. Only supported by AUTHORIZATION_CODE and DEVICE_CODE;
azure managed identity is rejected. Default: None (memory only).
Examples
--------
@@ -56,6 +143,12 @@ class OAuthConfig:
... scopes=["api://lancedb-api/.default"],
... )
Providers that require an explicit target can use ``resource`` and/or
``audience`` (these are forwarded unchanged):
>>> config.resource = "https://api.example.com"
>>> config.audience = "lancedb-api"
Azure Managed Identity:
>>> config = OAuthConfig(
@@ -64,6 +157,30 @@ class OAuthConfig:
... scopes=["api://lancedb-api/.default"],
... flow=OAuthFlowType.AZURE_MANAGED_IDENTITY,
... )
Authorization Code with PKCE:
The authorization URL is written to standard error before LanceDB tries to
open a browser, so it can be copied in headless environments.
>>> config = OAuthConfig(
... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0",
... client_id="app-id",
... scopes=["openid", "api://lancedb-api/access"],
... flow=OAuthFlowType.AUTHORIZATION_CODE,
... )
Device Authorization, with a persistent cache so later processes reuse
the session without a new device prompt. The verification URL and user
code are written to standard error before polling begins:
>>> config = OAuthConfig(
... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0",
... client_id="app-id",
... scopes=["openid", "offline_access", "api://lancedb-api/access"],
... flow=OAuthFlowType.DEVICE_CODE,
... token_cache=TokenCacheOptions(),
... )
"""
issuer_url: str
@@ -71,5 +188,73 @@ class OAuthConfig:
scopes: List[str]
flow: OAuthFlowType = OAuthFlowType.CLIENT_CREDENTIALS
client_secret: Optional[str] = field(default=None, repr=False)
client_auth_method: Optional[ClientAuthMethod] = None
redirect_uri: Optional[str] = None
callback_port: Optional[int] = None
use_pkce: bool = True
managed_identity_client_id: Optional[str] = None
refresh_buffer_secs: Optional[int] = None
token_cache: Optional[TokenCacheOptions] = None
resource: Optional[str] = None
audience: Optional[str] = None
class OAuthSession:
"""Explicit OAuth session lifecycle for the persistent token cache.
Built from the same :class:`OAuthConfig` used for
:func:`lancedb.connect_async` (including its ``token_cache`` options).
A connection created with the same configuration shares the cache, so
logging in here prepares tokens for later processes without any database
request.
``login`` always runs the configured interactive flow and replaces the
cached session (the most recent login wins). ``logout`` removes only the
local credential; it does not revoke anything with the provider and does
not sign out of a browser SSO session.
Examples
--------
>>> config = OAuthConfig(
... issuer_url="https://issuer.example.com",
... client_id="my-app",
... scopes=["openid", "offline_access"],
... flow=OAuthFlowType.DEVICE_CODE,
... token_cache=TokenCacheOptions(),
... )
>>> session = OAuthSession(config) # doctest: +SKIP
>>> status = await session.login() # doctest: +SKIP
>>> status.refreshable # doctest: +SKIP
True
"""
def __init__(self, config: OAuthConfig):
from lancedb._lancedb import OAuthSession as PyOAuthSession
self._inner: PyOAuthSession = PyOAuthSession(config)
async def login(self):
"""Eagerly run the configured flow and store the session.
Returns a :class:`lancedb._lancedb.SessionStatus` describing the
cached session. A successful login always replaces any prior cached
session for this identity; if the provider does not issue a refresh
token (for example without ``offline_access``), the previous record
is removed and ``refreshable`` is ``False``.
"""
return await self._inner.login()
async def status(self):
"""Report whether a cached session exists, with safe metadata.
Never contacts the identity provider and never exposes token values.
"""
return await self._inner.status()
async def logout(self):
"""Remove the matching local cached credential.
Returns a :class:`lancedb._lancedb.SessionLogout` whose ``removed``
flag reports whether a credential existed. Logout is idempotent.
"""
return await self._inner.logout()
python/python/lancedb/remote/table.py
+2 -1
View File
@@ -549,6 +549,7 @@ class RemoteTable(Table):
LOOP.run(
self._table.create_index(
column,
replace=replace,
config=config,
wait_timeout=wait_timeout,
name=name,
@@ -720,7 +721,7 @@ class RemoteTable(Table):
Parameters
----------
query: list/np.ndarray/str/PIL.Image.Image, default None
The targetted vector to search for.
The targeted vector to search for.
- *default None*.
Acceptable types are: list, np.ndarray, PIL.Image.Image
python/python/lancedb/rerankers/base.py
+1 -1
View File
@@ -175,7 +175,7 @@ class Reranker(ABC):
if the results haven't been executed yet or the results in arrow format.
query : str or None,
The input query. Some rerankers might not need the query to rerank.
In that case, it can be set to None explicitly. This is inteded to
In that case, it can be set to None explicitly. This is intended to
be handled by the reranker implementations.
deduplicate : bool, optional
Whether to deduplicate the results based on the `_rowid` column,
python/python/lancedb/secrets.py
+226
View File
@@ -0,0 +1,226 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Named Secrets, and the bindings that deliver them to Functions.
A Secret is a database-scoped named credential. Nothing in this module holds a
value: :class:`EnvVarSecret` names one and says which environment variable it
should arrive in, and the value is resolved by the remote service when a
Function bound to it runs. No API returns a stored credential, by construction
rather than by policy -- there is no code path that could.
"""
from __future__ import annotations
import re
# The same characters LanceDB already admits in a namespace or table name, and
# no positional rule on top of them: a segment may begin with `_`, `-` or `.`
# today, so anything narrower would put Secrets out of reach inside namespaces
# that already exist. Matches the service, which admits the same set.
_SECRET_NAME = re.compile(r"^[A-Za-z0-9_.-]{1,255}$")
_ENV_VARIABLE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
def validate_secret_name(name: str) -> str:
"""Check a Secret name locally and return it unchanged."""
if not isinstance(name, str):
raise TypeError(f"Secret name must be a string, not {type(name).__name__}")
if not _SECRET_NAME.fullmatch(name):
raise ValueError(f"invalid Secret name: {name!r}")
return name
def validate_namespace_path(namespace_path=None):
"""Check a namespace path locally and return it as a tuple.
``None`` and ``[]`` both mean the root namespace. Segments follow the same
rule as Secret names: a binding carries the path and the name as separate
fields, so neither is ever parsed out of the other.
"""
if namespace_path is None:
return ()
if isinstance(namespace_path, str):
raise TypeError(
"namespace_path must be a list of segments, not a string; "
f"did you mean [{namespace_path!r}]?"
)
segments = tuple(namespace_path)
for segment in segments:
if not isinstance(segment, str):
raise TypeError(
f"namespace path segment must be a string, not {type(segment).__name__}"
)
if not _SECRET_NAME.fullmatch(segment):
raise ValueError(f"invalid namespace path segment: {segment!r}")
return segments
def validate_env_variable(name: str) -> str:
"""Check an environment variable name locally and return it unchanged."""
if not isinstance(name, str):
raise TypeError(
f"environment variable name must be a string, not {type(name).__name__}"
)
if not _ENV_VARIABLE.fullmatch(name):
raise ValueError(f"invalid environment variable name: {name!r}")
return name
class EnvVarSecret:
"""A Secret bound to the environment variable a Function's library reads.
Pass these in the ``secrets`` sequence of
[DBConnection.create_function][lancedb.db.DBConnection.create_function]. The
Function's source is unchanged by binding: it reads ``OPENAI_API_KEY`` the
way it always did, and the binding is what puts a value there.
This is a local value. Constructing it contacts no server, so it always
succeeds and says nothing about whether the Secret exists; that is checked
at registration, where a mistyped Secret name surfaces as a clear "does not
exist" naming both the Secret and the variable bound to it. A mistyped
*variable* name cannot be caught anywhere -- nothing knows which variables a
Function reads -- so it surfaces on the first rows instead.
The type exists so a credential cannot be passed by accident. A bare string
in the same position is a plausible-looking mistake with the opposite
meaning, and it reads identically in a diff.
Parameters
----------
secret_name : str
The Secret's database-scoped name.
env_variable : str
The environment variable the Function reads it from.
secret_namespace_path : list of str, optional
The namespace the Secret is addressed within. ``None`` and ``[]`` both
mean the root namespace. Carried beside the name rather than joined
into it, so neither is ever parsed back out of the other.
Examples
--------
>>> from lancedb import EnvVarSecret
>>> binding = EnvVarSecret(
... secret_name="openai-prod", env_variable="OPENAI_API_KEY"
... )
>>> binding.secret_name, binding.env_variable
('openai-prod', 'OPENAI_API_KEY')
"""
__slots__ = ("_secret_name", "_env_variable", "_secret_namespace_path")
def __init__(
self, secret_name: str, env_variable: str, *, secret_namespace_path=None
):
self._secret_name = validate_secret_name(secret_name)
self._env_variable = validate_env_variable(env_variable)
self._secret_namespace_path = validate_namespace_path(secret_namespace_path)
@property
def secret_name(self) -> str:
"""The Secret's database-scoped name."""
return self._secret_name
@property
def env_variable(self) -> str:
"""The environment variable the value is delivered in."""
return self._env_variable
@property
def secret_namespace_path(self):
"""The namespace path the Secret is addressed within, root when empty."""
return list(self._secret_namespace_path)
def __repr__(self) -> str:
path = (
f", secret_namespace_path={list(self._secret_namespace_path)!r}"
if self._secret_namespace_path
else ""
)
return (
f"EnvVarSecret(secret_name={self._secret_name!r}, "
f"env_variable={self._env_variable!r}{path})"
)
def __eq__(self, other: object) -> bool:
return (
isinstance(other, EnvVarSecret)
and other._secret_name == self._secret_name
and other._env_variable == self._env_variable
and other._secret_namespace_path == self._secret_namespace_path
)
def __hash__(self) -> int:
return hash(
(
EnvVarSecret,
self._secret_name,
self._env_variable,
self._secret_namespace_path,
)
)
class SecretInfo:
"""What a database records about a Secret. Never its value.
Returned by
[DBConnection.describe_secret][lancedb.db.DBConnection.describe_secret].
"""
__slots__ = ("_name", "_created_at_millis", "_updated_at_millis")
def __init__(self, name: str, created_at_millis: int, updated_at_millis: int):
self._name = name
self._created_at_millis = created_at_millis
self._updated_at_millis = updated_at_millis
@property
def name(self) -> str:
"""The Secret's database-scoped name."""
return self._name
@property
def created_at_millis(self) -> int:
"""When the Secret was created, in milliseconds since the Unix epoch."""
return self._created_at_millis
@property
def updated_at_millis(self) -> int:
"""When the Secret's value was last rotated, in epoch milliseconds.
The only observable that a rotation landed: no API returns a credential,
so a caller confirms ``alter_secret`` took effect by watching this move.
"""
return self._updated_at_millis
@classmethod
def from_json(cls, value: dict) -> "SecretInfo":
return cls(
name=value["name"],
created_at_millis=value["created_at_millis"],
updated_at_millis=value["updated_at_millis"],
)
def __repr__(self) -> str:
return (
f"SecretInfo(name={self._name!r}, "
f"created_at_millis={self._created_at_millis!r}, "
f"updated_at_millis={self._updated_at_millis!r})"
)
def __eq__(self, other: object) -> bool:
return (
isinstance(other, SecretInfo)
and other._name == self._name
and other._created_at_millis == self._created_at_millis
and other._updated_at_millis == self._updated_at_millis
)
__all__ = [
"EnvVarSecret",
"SecretInfo",
"validate_env_variable",
"validate_secret_name",
]
python/python/lancedb/sql.py
+88
View File
@@ -0,0 +1,88 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Handles to SQL queries running on a remote database."""
from uuid import UUID
import pyarrow as pa
from lancedb.background_loop import LOOP
from . import _lancedb
from .arrow import AsyncRecordBatchReader
QueryDescription = _lancedb.QueryDescription
class AsyncQuery:
"""A handle to a submitted SQL query on an asynchronous connection."""
def __init__(self, inner: "_lancedb.SqlQuery"):
self._inner = inner
@property
def id(self) -> UUID:
"""The stable identifier scoped to the connection that submitted it."""
return self._inner.id
async def describe(self) -> QueryDescription:
"""Get a point-in-time description of the query."""
return await self._inner.describe()
async def reader(self) -> AsyncRecordBatchReader:
"""Wait for the initial result stream and return its Arrow reader.
Results are single-consumer. Calling this method more than once on the
same query raises an error. Later batches are streamed as they become
available without waiting for the full query to finish.
"""
return AsyncRecordBatchReader(await self._inner.reader())
async def cancel(self) -> None:
"""Request cancellation of the query."""
await self._inner.cancel()
class Query:
"""Synchronous counterpart of :class:`AsyncQuery`."""
def __init__(self, inner: AsyncQuery):
self._inner = inner
@property
def id(self) -> UUID:
"""The stable identifier scoped to the connection that submitted it."""
return self._inner.id
def describe(self) -> QueryDescription:
"""Get a point-in-time description of the query."""
return LOOP.run(self._inner.describe())
def reader(self) -> pa.RecordBatchReader:
"""Wait for the initial result stream and return a blocking reader.
Results are single-consumer. Calling this method more than once on the
same query raises an error. Later batches block only until they become
available, without waiting for the full query to finish.
"""
reader = LOOP.run(self._inner.reader())
def next_batch():
try:
return LOOP.run(reader.__anext__())
except StopAsyncIteration:
return None
def batches():
while (batch := next_batch()) is not None:
yield batch
return pa.RecordBatchReader.from_batches(reader.schema, batches())
def cancel(self) -> None:
"""Request cancellation of the query."""
LOOP.run(self._inner.cancel())
__all__ = ["AsyncQuery", "Query", "QueryDescription"]
python/python/lancedb/table.py
+112 -45
View File
@@ -717,26 +717,43 @@ def _align_field_types(
return new_fields
def _align_list_value_field(
value_field: pa.Field, target_value_field: pa.Field
) -> pa.Field:
# A list has exactly one child, so the inferred child name ("item") aligns
# positionally and adopts the table's child name; pa.Table.cast renames it.
return _align_field(value_field, target_value_field).with_name(
target_value_field.name
)
def _align_container_child(child: pa.Field, target_child: pa.Field) -> pa.Field:
# A list has one child, a map one key and one item, so an inferred child name
# ("item") aligns positionally and adopts the table's; pa.Table.cast renames it.
return _align_field(child, target_child).with_name(target_child.name)
def _arrow_json_storage_type(input_type: pa.DataType) -> Optional[pa.DataType]:
"""The storage type arrow.json would use for ``input_type``.
Returns None if the type cannot hold JSON text.
"""
if pa.types.is_string(input_type) or pa.types.is_string_view(input_type):
return pa.string()
if pa.types.is_large_string(input_type):
return pa.large_string()
return None
def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
# 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
# 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"
):
return field
if _field_extension_name(target_field) == "lance.json":
if _field_extension_name(field) == "arrow.json":
return field
# Plain JSON text, which is what pyarrow infers for a column of `str`, only
# needs the arrow.json label.
json_storage = _arrow_json_storage_type(field.type)
if json_storage is not None:
# Labelled through metadata rather than pa.json_(), which only exists on
# newer PyArrow; Lance reads the extension name off the field either way.
return pa.field(
field.name,
json_storage,
field.nullable,
{"ARROW:extension:name": "arrow.json"},
)
if pa.types.is_struct(target_field.type):
if pa.types.is_struct(field.type):
new_type = pa.struct(
@@ -750,7 +767,7 @@ def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
elif pa.types.is_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.list_(
_align_list_value_field(
_align_container_child(
field.type.value_field, target_field.type.value_field
)
)
@@ -759,7 +776,7 @@ def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
elif pa.types.is_large_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.large_list(
_align_list_value_field(
_align_container_child(
field.type.value_field, target_field.type.value_field
)
)
@@ -768,13 +785,28 @@ def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
elif pa.types.is_fixed_size_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.list_(
_align_list_value_field(
_align_container_child(
field.type.value_field, target_field.type.value_field
),
target_field.type.list_size,
)
else:
new_type = target_field.type
elif pa.types.is_map(target_field.type):
if pa.types.is_map(field.type):
# A map has exactly one key and one item field, so like a list's child they
# align positionally and adopt the table's names.
new_type = pa.map_(
_align_container_child(
field.type.key_field, target_field.type.key_field
),
_align_container_child(
field.type.item_field, target_field.type.item_field
),
keys_sorted=target_field.type.keys_sorted,
)
else:
new_type = target_field.type
else:
new_type = target_field.type
return pa.field(field.name, new_type, field.nullable, target_field.metadata)
@@ -1630,7 +1662,9 @@ class Table(ABC):
on: Union[str, Iterable[str]]
A column (or columns) to join on. This is how records from the
source table and target table are matched. Typically this is some
kind of key or id column.
kind of key or id column. Passing several columns matches on the
composite key: a source row updates a target row only when it
agrees on every one of them.
Examples
--------
@@ -1700,7 +1734,7 @@ class Table(ABC):
Parameters
----------
query: list/np.ndarray/str/PIL.Image.Image, default None
The targetted vector to search for.
The targeted vector to search for.
- *default None*.
Acceptable types are: list, np.ndarray, PIL.Image.Image
@@ -1761,9 +1795,9 @@ class Table(ABC):
Offsets are mostly useful for sampling as the set of all valid offsets is easily
known in advance to be [0, len(table)).
No guarantees are made regarding the order in which results are returned. If
you desire an output order that matches the order of the given offsets, you will
need to add the row offset column to the output and align it yourself.
No guarantees are made regarding the order in which results are returned.
Repeated offsets produce repeated rows, which makes this method suitable for
sampling with replacement.
Parameters
----------
@@ -1874,6 +1908,9 @@ class Table(ABC):
The result has the same length and order as ``row_ids``. Null blobs
produce null slots; valid empty blobs produce ``b""``.
``_rowid`` values stay valid after compaction when the table has stable
row ids.
Convenience for small payloads. For large values use
:meth:`fetch_blob_files`.
"""
@@ -1891,6 +1928,9 @@ class Table(ABC):
The result has the same length and order as ``requests``; null blobs
produce null slots and empty ranges on non-null blobs produce ``b""``.
``_rowid`` values stay valid after compaction when the table has stable
row ids.
Row IDs can be obtained from a query with ``with_row_id(True)``. This
API is currently supported only by local tables.
"""
@@ -1906,6 +1946,9 @@ class Table(ABC):
``_rowid`` or a ``_lance_row_id`` field on the blob descriptor. Null
rows are ``None``. Remote tables require LanceDB Cloud server 0.5.0 or
newer.
``_rowid`` values stay valid after compaction when the table has stable
row ids.
"""
@abstractmethod
@@ -2266,10 +2309,10 @@ class Table(ABC):
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.
A refresh also recomputes the rows whose inputs changed since they
were computed, so a mutated input is reflected by the next refresh.
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
@@ -2289,7 +2332,7 @@ class Table(ABC):
>>> table.add_columns(computed={"doubled": "x * 2"})
AddColumnsResult(version=2)
>>> table.refresh_column("doubled")
RefreshColumnResult(rows_filled=2, version=3)
RefreshColumnResult(rows_filled=2, version=4)
>>> table.to_arrow().sort_by("x").to_pandas()
x doubled
0 1 2
@@ -2303,8 +2346,8 @@ class Table(ABC):
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.
by the next one, and rows whose inputs changed since they were computed
are recomputed; everything else is left as it is.
Local tables only: a remote refresh runs as a server job, through
[`refresh_column_async`][lancedb.table.Table.refresh_column_async].
@@ -3919,7 +3962,7 @@ class LanceTable(Table):
Parameters
----------
query: list/np.ndarray/str/PIL.Image.Image, default None
The targetted vector to search for.
The targeted vector to search for.
- *default None*.
Acceptable types are: list, np.ndarray, PIL.Image.Image
@@ -4179,6 +4222,7 @@ class LanceTable(Table):
)
and not self._route_pushdown_to_rust
and self.current_branch() is None
and query.take_offsets is None
):
from lancedb.namespace import _execute_server_side_query
@@ -4402,13 +4446,14 @@ class LanceTable(Table):
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
"""Fill a computed column's unfilled rows and recompute those whose
inputs changed. 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
"""Fill a computed column's unfilled rows and recompute those whose
inputs changed, 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)))
@@ -5722,7 +5767,7 @@ class AsyncTable:
if fill_value is None:
fill_value = 0.0
# _santitize_data is an old code path, but we will use it until the
# _sanitize_data is an old code path, but we will use it until the
# new code path is ready.
if mode == "overwrite":
# For overwrite, apply the same preprocessing as create_table
@@ -5796,7 +5841,9 @@ class AsyncTable:
on: Union[str, Iterable[str]]
A column (or columns) to join on. This is how records from the
source table and target table are matched. Typically this is some
kind of key or id column.
kind of key or id column. Passing several columns matches on the
composite key: a source row updates a target row only when it
agrees on every one of them.
Examples
--------
@@ -5896,7 +5943,7 @@ class AsyncTable:
Parameters
----------
query: list/np.ndarray/str/PIL.Image.Image, default None
The targetted vector to search for.
The targeted vector to search for.
- *default None*.
Acceptable types are: list, np.ndarray, PIL.Image.Image
@@ -6079,7 +6126,23 @@ class AsyncTable:
def _sync_query_to_async(
self, query: Query
) -> AsyncHybridQuery | AsyncFTSQuery | AsyncVectorQuery | AsyncQuery:
) -> (
AsyncHybridQuery
| AsyncFTSQuery
| AsyncVectorQuery
| AsyncQuery
| AsyncTakeQuery
):
if query.take_offsets is not None:
take_query = self.take_offsets(query.take_offsets)
if query.columns:
take_query = take_query.select(query.columns)
if query.use_lsm is not None:
take_query = take_query.use_lsm(query.use_lsm)
if query.with_row_id:
take_query = take_query.with_row_id()
return take_query
async_query = self.query()
if query.limit is not None:
async_query = async_query.limit(query.limit)
@@ -6144,6 +6207,7 @@ class AsyncTable:
self._namespace_client, self._pushdown_operations
)
and not self._route_pushdown_to_rust
and query.take_offsets is None
):
from lancedb.namespace import _execute_server_side_query
@@ -6377,10 +6441,10 @@ class AsyncTable:
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.
A refresh also recomputes the rows whose inputs changed since they
were computed, so a mutated input is reflected by the next refresh.
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``.
@@ -6442,8 +6506,8 @@ class AsyncTable:
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.
by the next one, and rows whose inputs changed since they were computed
are recomputed; everything else is left as it is.
Local tables only: a remote refresh runs as a server job, through
[`refresh_column_async`][lancedb.table.Table.refresh_column_async].
@@ -6641,6 +6705,9 @@ class AsyncTable:
Offsets are mostly useful for sampling as the set of all valid offsets is easily
known in advance to be [0, len(table)).
No guarantees are made regarding the order in which results are returned.
Repeated offsets produce repeated rows.
Parameters
----------
offsets: list[int]
python/python/tests/test_blob.py
+134 -2
View File
@@ -66,6 +66,25 @@ def _row_ids_by_id(table):
return dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
def _assert_missing_blob_row_ids(exc_info):
message = str(exc_info.value)
assert "row ids" in message
assert "rowaddr" not in message
assert "fragment" not in message
def _assert_fetch_apis_reject_missing_row_ids(table, row_ids):
with pytest.raises(ValueError) as exc_info:
table.fetch_blobs("image", row_ids)
_assert_missing_blob_row_ids(exc_info)
with pytest.raises(ValueError) as exc_info:
table.fetch_blob_files("image", row_ids)
_assert_missing_blob_row_ids(exc_info)
with pytest.raises(ValueError) as exc_info:
table.fetch_blob_ranges("image", [(row_id, 0, 1) for row_id in row_ids])
_assert_missing_blob_row_ids(exc_info)
def test_blob_factory_declares_v2_field():
field = lancedb.blob("image")
assert isinstance(field.type, pa.ExtensionType)
@@ -278,7 +297,10 @@ def test_blob_v2_projection_sources_use_typed_column_name():
def _legacy_v1_table(name):
db = lancedb.connect("memory:///")
# Legacy v1 blob columns are only writable at file version <= 2.1.
db = lancedb.connect(
"memory:///", storage_options={"new_table_data_storage_version": "2.1"}
)
schema = pa.schema(
[
pa.field("id", pa.int64()),
@@ -691,6 +713,25 @@ def test_fetch_blobs_accepts_query_result():
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"gamma"}
def test_fetch_blobs_after_compact_with_stable_row_ids(tmp_path):
db = lancedb.connect(
tmp_path, storage_options={"new_table_enable_stable_row_ids": "true"}
)
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("t", schema=schema)
table.add([{"id": 1, "image": b"frag-one"}])
table.add([{"id": 2, "image": b"frag-two"}])
by_id = _row_ids_by_id(table)
ids = [by_id[1], by_id[2]]
table.optimize()
blobs = table.fetch_blobs("image", ids)
assert blobs.to_pylist() == [b"frag-one", b"frag-two"]
ranges = table.fetch_blob_ranges("image", [(ids[0], 5, 3), (ids[1], 5, 3)])
assert ranges.to_pylist() == [b"one", b"two"]
def test_fetch_blobs_preserves_null_and_empty_values():
table = _blob_table(
"nulls",
@@ -710,6 +751,80 @@ def test_fetch_blobs_preserves_null_and_empty_values():
assert blobs[3].as_py() == b"present"
def test_add_all_null_list_to_blob_column():
table = _blob_table("all_null_add", [{"id": 1, "image": None}])
hits = table.search().to_arrow()
blobs = table.fetch_blobs("image", hits)
assert len(blobs) == 1
assert blobs[0].as_py() is None
def test_add_all_null_list_to_blob_column_with_sanitizer():
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("all_null_sanitized_add", schema=schema)
table.add([{"id": 1, "image": None}], on_bad_vectors="fill")
hits = table.search().to_arrow()
blobs = table.fetch_blobs("image", hits)
assert len(blobs) == 1
assert blobs[0].as_py() is None
def test_add_all_null_list_to_nested_blob_column():
db = lancedb.connect("memory:///")
blob_field = lancedb.blob("image")
info_field = pa.field("info", pa.struct([blob_field]))
info = pa.StructArray.from_arrays(
[_blob_array("image", [b"seed"])], fields=[blob_field]
)
seed = pa.Table.from_arrays(
[pa.array([0], type=pa.int64()), info],
schema=pa.schema([pa.field("id", pa.int64()), info_field]),
)
table = db.create_table("nested_null_add", data=seed)
table.add([{"id": 1, "info": {"image": None}}])
table.add([{"id": 2, "info": {"image": None}}], on_bad_vectors="fill")
hits = table.search().where("id > 0").to_arrow()
blobs = table.fetch_blobs("info.image", hits)
assert len(blobs) == 2
assert all(blob.as_py() is None for blob in blobs)
@pytest.mark.parametrize("large_list", [False, True], ids=["list", "large_list"])
def test_add_list_of_dicts_to_blob_list_column(large_list):
db = lancedb.connect("memory:///")
blob_field = lancedb.blob("image")
blob_values = _blob_array("image", [b"seed"])
if large_list:
items_field = pa.field("items", pa.large_list(blob_field))
items = pa.LargeListArray.from_arrays(
pa.array([0, 1], type=pa.int64()), blob_values
)
else:
items_field = pa.field("items", pa.list_(blob_field))
items = pa.ListArray.from_arrays(pa.array([0, 1], type=pa.int32()), blob_values)
seed = pa.Table.from_arrays(
[pa.array([0], type=pa.int64()), items],
schema=pa.schema([pa.field("id", pa.int64()), items_field]),
)
table = db.create_table(f"blob_{large_list}_list_add", data=seed)
table.add([{"id": 1, "items": [None]}])
table.add(
[{"id": 2, "items": [b"a", None]}],
on_bad_vectors="fill",
)
ids = table.search().select(["id"]).to_arrow()["id"].to_pylist()
assert sorted(ids) == [0, 1, 2]
assert pa.types.is_large_list(table.schema.field("items").type) is large_list
def test_fetch_blob_ranges_aligns_repeated_ranges_and_nulls():
table = _blob_table(
"range_alignment",
@@ -739,8 +854,25 @@ def test_fetch_blob_ranges_validates_requests():
with pytest.raises(ValueError, match="offset \\+ length overflowed"):
table.fetch_blob_ranges("image", [(row_id, 2**64 - 1, 1)])
with pytest.raises(ValueError, match="row IDs"):
with pytest.raises(ValueError) as exc_info:
table.fetch_blob_ranges("image", [(2**64 - 1, 0, 1)])
_assert_missing_blob_row_ids(exc_info)
def test_fetch_blob_apis_reject_missing_fragment_row_addr():
table = _blob_table("missing_frag", [{"id": 1, "image": b"x"}])
live = _row_ids_by_id(table)[1]
_assert_fetch_apis_reject_missing_row_ids(table, [1 << 32, live])
def test_fetch_blob_apis_reject_deleted_row_ids():
table = _blob_table(
"deleted_rows",
[{"id": 1, "image": b"one"}, {"id": 2, "image": b"two"}],
)
by_id = _row_ids_by_id(table)
table.delete("id = 2")
_assert_fetch_apis_reject_missing_row_ids(table, [by_id[2], by_id[1]])
def test_fetch_blob_ranges_empty_requests_returns_empty_array():
python/python/tests/test_catalog.py
+140
View File
@@ -0,0 +1,140 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import json
from threading import Thread
import pytest
import lancedb
from lancedb.db import AsyncConnection, DBConnection
from lancedb.remote.errors import HttpError
@pytest.fixture
def catalog_server():
requests = []
responses = []
class Handler(BaseHTTPRequestHandler):
def handle_request(self):
body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
requests.append(
(
self.path,
dict(self.headers.items()),
json.loads(body) if body else None,
)
)
status, response = responses.pop(0)
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.end_headers()
if status != 204:
self.wfile.write(json.dumps(response).encode())
do_GET = handle_request
do_POST = handle_request
with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server:
thread = Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
yield f"http://127.0.0.1:{server.server_port}", requests, responses
finally:
server.shutdown()
thread.join()
def test_catalog_sync_scope_and_serialization(catalog_server):
endpoint, requests, responses = catalog_server
responses.extend(
[
(204, None),
(200, {}),
(200, {"tables": []}),
(200, {"tables": []}),
(200, {"namespaces": ["team/search"], "page_token": "next"}),
(204, None),
]
)
catalog = lancedb.connect_catalog(
endpoint,
api_key="secret",
sql_host_override="invalid://localhost",
client_config={
"extra_headers": {
"X-LanceDB-Database": "wrong",
"X-LanceDB-Database-Prefix": "wrong",
}
},
)
assert isinstance(catalog, lancedb.Catalog)
assert catalog.uri == endpoint
db = catalog.create_database("team/search", exist_ok=True)
assert isinstance(db, DBConnection)
with pytest.raises(ValueError, match="sql_host_override must use"):
db.execute_query_async("SELECT 1")
db = catalog.connect_database("team/search")
assert isinstance(db, DBConnection)
assert db.table_names() == []
restored = lancedb.deserialize_conn(db.serialize())
assert restored.sql_host_override == "invalid://localhost"
for connection in (db, restored):
with pytest.raises(ValueError, match="sql_host_override must use"):
connection.execute_query_async("SELECT 1")
assert restored.table_names() == []
page = catalog.list_databases(limit=1, page_token="a/b")
assert page == lancedb.ListDatabasesResponse(["team/search"], "next")
catalog.drop_database("team/search", ignore_missing=True)
assert requests[0][0] == "/v1/namespace/team%2Fsearch/create"
assert requests[0][2] == {"mode": "ExistOk"}
assert requests[1][0] == "/v1/namespace/team%2Fsearch/describe"
assert requests[4][0] == "/v1/namespace/$/list?limit=1&page_token=a%2Fb"
assert requests[5][2] == {"mode": "Skip", "behavior": "Restrict"}
for i, (_, headers, _) in enumerate(requests):
headers = {key.lower(): value for key, value in headers.items()}
assert headers.get("x-lancedb-database") == (
"team/search" if i in (2, 3) else None
)
assert "x-lancedb-database-prefix" not in headers
assert headers["x-api-key"] == "secret"
@pytest.mark.asyncio
async def test_catalog_async_and_errors(catalog_server):
endpoint, requests, responses = catalog_server
responses.extend(
[
(200, {}),
(200, {"tables": []}),
(404, {"error": "missing"}),
(409, {"error": "exists"}),
(400, {"error": "not empty"}),
(404, {"error": "missing"}),
]
)
catalog = await lancedb.connect_catalog_async(
endpoint, sql_host_override="invalid://localhost"
)
assert isinstance(catalog, lancedb.AsyncCatalog)
db = await catalog.connect_database("analytics")
assert isinstance(db, AsyncConnection)
with pytest.raises(ValueError, match="sql_host_override must use"):
await db.execute_query_async("SELECT 1")
assert await db.table_names() == []
with pytest.raises(ValueError, match="missing"):
await catalog.connect_database("missing")
with pytest.raises(ValueError, match="exists"):
await catalog.create_database("exists")
with pytest.raises(HttpError):
await catalog.drop_database("full")
await catalog.drop_database("missing", ignore_missing=True)
assert requests[4][2] == {"mode": "Fail", "behavior": "Restrict"}
@pytest.mark.parametrize("endpoint", ["/tmp/catalog", "s3://bucket", "db://db"])
def test_catalog_requires_remote_endpoint(endpoint):
with pytest.raises(ValueError, match="endpoint"):
lancedb.connect_catalog(endpoint)
python/python/tests/test_embeddings.py
+5 -5
View File
@@ -327,8 +327,8 @@ def test_embedding_function_with_pandas(tmp_path):
) -> List[np.array]:
return [np.random.randn(self.ndims()).tolist() for _ in range(len(texts))]
registery = get_registry()
func = registery.get("mock-embedding").create()
registry = get_registry()
func = registry.get("mock-embedding").create()
class TestSchema(LanceModel):
text: str = func.SourceField()
@@ -394,9 +394,9 @@ def test_multiple_embeddings_for_pandas(tmp_path):
) -> List[np.array]:
return [np.random.randn(self.ndims()).tolist() for _ in range(len(texts))]
registery = get_registry()
func1 = registery.get("mock-embedding").create()
func2 = registery.get("mock-embedding2").create()
registry = get_registry()
func1 = registry.get("mock-embedding").create()
func2 = registry.get("mock-embedding2").create()
class TestSchema(LanceModel):
text: str = func1.SourceField()
python/python/tests/test_first_class_function_slice1.py
+75 -12
View File
@@ -13,7 +13,9 @@ from lancedb.functions import (
FunctionBinding,
FunctionVersion,
PythonRuntimeSpec,
SecretBinding,
RefreshColumnResult,
SecretReference,
)
from lancedb.table import AsyncTable
@@ -37,6 +39,22 @@ def job_result(name: str) -> dict:
return json.loads(fixture(name))["result"]
def assert_no_secret_values(value):
"""No client value models a resolved credential, at any nesting depth."""
if isinstance(value, dict):
for key, child in value.items():
assert key not in {
"secret_value",
"secret_values",
"resolved_secret",
"resolved_secrets",
}
assert_no_secret_values(child)
elif isinstance(value, list):
for child in value:
assert_no_secret_values(child)
def test_public_function_values_are_in_api_reference():
docs = Path(__file__).parents[3] / "docs" / "src" / "python" / "python.md"
rendered = docs.read_text()
@@ -93,16 +111,27 @@ 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"
assert version.version == "1"
assert version.image.manifest_digest.startswith("sha256:")
assert version.version != version.image.manifest_digest
assert list(version.secret_bindings) == [
SecretBinding(
kind="env", variable="HF_TOKEN", secret_ref=SecretReference(name="hf-prod")
)
]
with pytest.raises((TypeError, ValueError)):
version.version = "fv_changed"
version.version = "1"
with pytest.raises(TypeError, match="immutable"):
version.runtime.env["TOKENIZERS_PARALLELISM"] = "true"
version.image.descriptor["format_version"] = "changed"
changed = dict(value)
changed["version"] = "fv_changed"
changed["version"] = "2"
assert FunctionVersion(**changed) != version
assert FunctionVersion(**changed).image == version.image
for invalid in [version.image.manifest_digest, "0", "01", "-1", str(2**64)]:
with pytest.raises(ValueError):
FunctionVersion(**{**value, "version": invalid})
def test_function_version_binds_named_columns_as_one_immutable_application():
@@ -137,7 +166,7 @@ def test_function_version_binding_validates_names_and_direct_columns():
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["version"] = "1"
value["signature"] = {
"inputs": [
{"name": "title", "arrow_type": "utf8", "nullable": True},
@@ -182,14 +211,14 @@ def test_function_version_keeps_named_struct_outputs_in_one_application():
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["image"]["descriptor"]["future_interface"] = {"kind": "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.image.descriptor["future_interface"] == {"kind": "wasm"}
assert json.loads(version.to_canonical_json())["image"]["descriptor"][
"future_interface"
] == {"kind": "wasm"}
assert version.signature.output.kind == "future_output_shape"
@@ -222,7 +251,7 @@ def test_function_application_uses_rename_columns_only():
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 binding.function.version == "1"
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
@@ -276,6 +305,31 @@ def test_refresh_result_rejects_non_u64_values(field):
RefreshColumnResult.from_json(json.dumps(value))
def test_canonical_client_values_carry_bindings_and_no_credentials():
"""A binding names a Secret; the credential behind it has no client field."""
version = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
)
canonical = json.loads(version.to_canonical_json())
assert canonical["secret_bindings"] == [
{"kind": "env", "variable": "HF_TOKEN", "secret_ref": {"name": "hf-prod"}}
]
assert_no_secret_values(canonical)
def test_a_version_without_bindings_omits_the_field_in_both_directions():
"""A Function that binds nothing carries no ``secret_bindings`` key.
Absent decodes as an empty list, and an empty list serializes back to
absent.
"""
value = job_result("remote_function_job.json")
del value["secret_bindings"]
version = FunctionVersion.from_json(json.dumps(value))
assert list(version.secret_bindings) == []
assert "secret_bindings" not in json.loads(version.to_canonical_json())
class _FunctionDeclarationInner:
def __init__(self):
self.calls = []
@@ -339,7 +393,16 @@ def test_rename_requires_named_struct_and_keeps_partial_mapping_immutable():
scalar = FunctionApplication.from_json(
json.dumps(
{
"function": {"name": "embed", "version": "fv_exact"},
"function": {
"name": "embed",
"version": "1",
"object_id": "fixture",
"location": "memory:///fixture",
"manifest_digest": (
"sha256:"
"7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"
),
},
"inputs": [],
"output": {
"kind": "scalar",
python/python/tests/test_first_class_function_slice2.py
+1241 -39
View File
File diff suppressed because it is too large Load Diff
python/python/tests/test_fts.py
+14 -4
View File
@@ -1011,8 +1011,13 @@ def test_fts_ngram(mem_db: DBConnection):
assert set(r["text"] for r in results) == {"lance database", "lance is cool"}
results = (
table.search("nce", query_type="fts").limit(10).to_list()
) # spellchecker:disable-line
table.search(
"nce", # spellchecker:disable-line
query_type="fts",
)
.limit(10)
.to_list()
)
assert len(results) == 2
assert set(r["text"] for r in results) == {"lance database", "lance is cool"}
@@ -1034,8 +1039,13 @@ def test_fts_ngram(mem_db: DBConnection):
assert set(r["text"] for r in results) == {"lance database", "lance is cool"}
results = (
table.search("nce", query_type="fts").limit(10).to_list()
) # spellchecker:disable-line
table.search(
"nce", # spellchecker:disable-line
query_type="fts",
)
.limit(10)
.to_list()
)
assert len(results) == 0
results = table.search("la", query_type="fts").limit(10).to_list()
python/python/tests/test_header_provider.py
+24 -6
View File
@@ -54,7 +54,10 @@ class TestOAuthProvider:
provider = OAuthProvider(fetcher)
headers = provider.get_headers()
assert headers == {"Authorization": "Bearer token123"}
assert headers == {
"Authorization": "Bearer token123",
"x-lancedb-credential-type": "oidc",
}
assert provider._current_token == "token123"
assert provider._token_expires_at is not None
@@ -73,14 +76,20 @@ class TestOAuthProvider:
# First call
headers1 = provider.get_headers()
assert headers1 == {"Authorization": "Bearer token1"}
assert headers1 == {
"Authorization": "Bearer token1",
"x-lancedb-credential-type": "oidc",
}
# Wait for token to expire
time.sleep(1.1)
# Second call should refresh
headers2 = provider.get_headers()
assert headers2 == {"Authorization": "Bearer token2"}
assert headers2 == {
"Authorization": "Bearer token2",
"x-lancedb-credential-type": "oidc",
}
assert call_count == 2
def test_no_expiry_info(self):
@@ -92,12 +101,18 @@ class TestOAuthProvider:
provider = OAuthProvider(fetcher)
headers = provider.get_headers()
assert headers == {"Authorization": "Bearer permanent_token"}
assert headers == {
"Authorization": "Bearer permanent_token",
"x-lancedb-credential-type": "oidc",
}
assert provider._token_expires_at is None
# Should not refresh on second call
headers2 = provider.get_headers()
assert headers2 == {"Authorization": "Bearer permanent_token"}
assert headers2 == {
"Authorization": "Bearer permanent_token",
"x-lancedb-credential-type": "oidc",
}
def test_missing_access_token(self):
"""Test error handling when access_token is missing."""
@@ -121,7 +136,10 @@ class TestOAuthProvider:
provider = OAuthProvider(fetcher)
headers = provider.get_headers()
assert headers == {"Authorization": "Bearer sync_token"}
assert headers == {
"Authorization": "Bearer sync_token",
"x-lancedb-credential-type": "oidc",
}
class TestClientConfigIntegration:
python/python/tests/test_hybrid_query.py
+87
View File
@@ -203,6 +203,93 @@ async def test_async_hybrid_query_default_limit(table: AsyncTable):
assert texts.count("a") == 1
@pytest.mark.asyncio
async def test_async_hybrid_query_offset(table: AsyncTable):
# The offset window of a hybrid query must be a suffix of the same query
# run without an offset. Skipping the first rows of each sub-query instead
# of the first rows of the fused result silently changes which rows land in
# the window.
full = await (
table.query()
.nearest_to([0.0, 0.4])
.nearest_to_text("dog")
.limit(4)
.with_row_id()
.to_arrow()
)
assert len(full) == 4
second_page = await (
table.query()
.nearest_to([0.0, 0.4])
.nearest_to_text("dog")
.offset(2)
.limit(2)
.with_row_id()
.to_arrow()
)
assert second_page["_rowid"].to_pylist() == full["_rowid"].to_pylist()[2:]
first_page = await (
table.query()
.nearest_to([0.0, 0.4])
.nearest_to_text("dog")
.limit(2)
.with_row_id()
.to_arrow()
)
# Paging through the result must visit every row exactly once: no row
# repeated from the previous page and none dropped between the two.
paged = first_page["_rowid"].to_pylist() + second_page["_rowid"].to_pylist()
assert sorted(paged) == sorted(full["_rowid"].to_pylist())
@pytest.mark.asyncio
async def test_async_hybrid_query_fts_first_default_limit(table: AsyncTable):
# nearest_to() and nearest_to_text() build their new sibling sub-query from
# scratch, and that is the sub-query the default limit ends up on. So the
# side that carries the limit depends on the order the hybrid query was
# built in, and looking at only one side loses the limit for half the ways
# a hybrid query can be written. Without a limit the combined results are
# not truncated at all and the whole union of both candidate lists is
# returned.
await table.add([{"text": "dog", "vector": [50.0 + i, 50.0]} for i in range(10)])
result = await (
table.query().nearest_to_text("dog").nearest_to([0.1, 0.1]).to_arrow()
)
assert len(result) == 10
offset_result = await (
table.query().nearest_to_text("dog").nearest_to([0.1, 0.1]).offset(2).to_arrow()
)
assert len(offset_result) == 10
@pytest.mark.asyncio
async def test_async_hybrid_query_explain_plan_matches_execution(table: AsyncTable):
# Paging rewrites the sub-queries: each one fetches limit + offset rows with
# no offset of its own, and the window is sliced out after fusion. The plans
# have to be built from those rewritten sub-queries, otherwise explain_plan
# and analyze_plan describe a query that is never run.
query = (
table.query().nearest_to([0.0, 0.4]).nearest_to_text("dog").offset(2).limit(2)
)
await query.to_arrow()
plan = await query.explain_plan()
assert [
line.strip() for line in plan.splitlines() if "GlobalLimitExec" in line
] == [
"GlobalLimitExec: skip=0, fetch=4",
"GlobalLimitExec: skip=0, fetch=4",
]
analyzed = await query.analyze_plan()
assert analyzed.count("skip=0, fetch=4") == 2
assert "skip=2" not in analyzed
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.
python/python/tests/test_materialized_views.py
+275 -24
View File
@@ -1,9 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import contextlib
import http.server
import json
import threading
import lancedb
import pytest
from lancedb.materialized_view import MaterializedViewDefinition
from lancedb.remote.db import RemoteDBConnection
STABLE_ROW_IDS = {"new_table_enable_stable_row_ids": "true"}
@@ -22,6 +28,152 @@ def make_db(tmp_path):
return db
@contextlib.contextmanager
def mock_remote_materialized_views():
requests = []
class Handler(http.server.BaseHTTPRequestHandler):
def log_message(self, *args):
pass
def do_GET(self):
requests.append(self.path)
encoded = json.dumps({"views": ["daily_sales"]}).encode()
self.send_response(200)
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]}", requests
finally:
server.shutdown()
thread.join()
@contextlib.contextmanager
def mock_remote_materialized_view_create():
requests = []
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"{}")
requests.append((self.path, body))
job_id = "mv-drop-123" if self.path.endswith("/drop") else "mv-create-123"
encoded = json.dumps({"job_id": job_id}).encode()
self.send_response(202)
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]}", requests
finally:
server.shutdown()
thread.join()
def test_remote_list_uses_namespace_route():
with mock_remote_materialized_views() as (host, requests):
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
assert db.list_materialized_views() == ["daily_sales"]
assert requests == ["/v1/namespace/$/materialized_view/list"]
def test_remote_create_async_returns_server_job():
with mock_remote_materialized_view_create() as (host, requests):
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
job = db.create_materialized_view_async("adults", "people", where="age >= 18")
assert job.id == "mv-create-123"
assert requests == [
(
"/v1/materialized_view/adults/create",
{"query": 'SELECT * FROM "people" WHERE age >= 18', "with_no_data": False},
)
]
def test_remote_drop_async_returns_server_job():
with mock_remote_materialized_view_create() as (host, requests):
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
job = db.drop_materialized_view_async("adults")
assert job.id == "mv-drop-123"
assert requests == [("/v1/materialized_view/adults/drop", {})]
def test_sync_remote_create_uses_public_async_connection():
calls = []
class StubAsyncTable:
name = "adults"
class StubAsyncMaterializedView:
table = StubAsyncTable()
class StrictAsyncConnection:
async def create_materialized_view(
self,
name,
source,
*,
select=None,
where=None,
limit=None,
with_no_data=False,
):
calls.append((name, source, select, where, limit, with_no_data))
return StubAsyncMaterializedView()
async def drop_materialized_view(self, name, *, namespace_path=None):
calls.append(("drop", name, namespace_path))
db = RemoteDBConnection.__new__(RemoteDBConnection)
db._conn = StrictAsyncConnection()
db.db_name = "example"
db.serialize = lambda: "{}"
view = db.create_materialized_view(
"adults",
"people",
select=["name"],
where="age >= 18",
limit=10,
with_no_data=True,
)
assert view.name == "adults"
assert calls == [("adults", "people", ["name"], "age >= 18", 10, True)]
db.drop_materialized_view("adults", namespace_path=["analytics"])
assert calls[-1] == ("drop", "adults", ["analytics"])
def test_create_refresh_and_query(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view(
@@ -31,16 +183,34 @@ def test_create_refresh_and_query(tmp_path):
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
assert view.table.count_rows() == 2
rows = view.table.search().to_list()
assert sorted(row["shout"] for row in rows) == ["ADA", "GRACE"]
def test_create_and_refresh_jobs(tmp_path):
db = make_db(tmp_path)
create_job = db.create_materialized_view_async(
"adults", "people", where="age >= 18", with_no_data=True
)
assert create_job.id is None
assert create_job.wait() is None
view = db.open_materialized_view("adults")
refresh_job = view.refresh_async()
assert refresh_job.id is None
result = refresh_job.wait()
assert result.mode == "rebuild"
assert result.rows_written == 2
assert view.table.count_rows() == 2
drop_job = db.drop_materialized_view_async("adults")
assert drop_job.id is None
assert drop_job.wait() is None
assert "adults" not in db.list_materialized_views()
def test_definition_round_trips(tmp_path):
db = make_db(tmp_path)
db.create_materialized_view("adults", "people", where="age >= 18")
@@ -56,7 +226,7 @@ def test_definition_round_trips(tmp_path):
def test_incremental_refresh_after_append(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view("copy", "people")
view = db.create_materialized_view("copy", "people", with_no_data=True)
view.refresh()
db.open_table("people").add([{"name": "alan", "age": 41}])
@@ -70,7 +240,7 @@ def test_incremental_refresh_after_append(tmp_path):
def test_incremental_refresh_after_update(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view("copy", "people")
view = db.create_materialized_view("copy", "people", with_no_data=True)
view.refresh()
db.open_table("people").update(where="name = 'kid'", values={"age": 8})
@@ -87,7 +257,7 @@ def test_legacy_storage_source_update_rebuilds(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 = db.create_materialized_view("copy", "people", with_no_data=True)
view.refresh()
db.open_table("people").update(where="name = 'kid'", values={"age": 8})
@@ -104,6 +274,11 @@ def test_list_and_not_a_view(tmp_path):
assert db.list_materialized_views() == ["adults"]
with pytest.raises(ValueError, match="not a materialized view"):
db.open_materialized_view("people")
with pytest.raises(ValueError, match="not a materialized view"):
db.drop_materialized_view("people")
db.drop_materialized_view("adults")
assert db.list_materialized_views() == []
def test_invalid_expression_fails_at_create(tmp_path):
@@ -119,7 +294,10 @@ async def test_async_create_refresh_and_open(tmp_path):
await db.create_table("people", [{"name": "ada", "age": 36}])
view = await db.create_materialized_view(
"shouts", "people", select=[("shout", "upper(name)")]
"shouts",
"people",
select=[("shout", "upper(name)")],
with_no_data=True,
)
result = await view.refresh()
assert result.mode == "rebuild"
@@ -131,11 +309,35 @@ async def test_async_create_refresh_and_open(tmp_path):
assert await db.list_materialized_views() == ["shouts"]
@pytest.mark.asyncio
async def test_async_create_and_refresh_jobs(tmp_path):
db = await lancedb.connect_async(tmp_path, storage_options=STABLE_ROW_IDS)
await db.create_table("people", [{"name": "ada", "age": 36}])
create_job = await db.create_materialized_view_async(
"adults", "people", with_no_data=True
)
assert create_job.id is None
assert await create_job.wait() is None
view = await db.open_materialized_view("adults")
refresh_job = await view.refresh_async()
assert refresh_job.id is None
result = await refresh_job.wait()
assert result.mode == "rebuild"
assert result.rows_written == 1
drop_job = await db.drop_materialized_view_async("adults")
assert drop_job.id is None
assert await drop_job.wait() is None
assert "adults" not in await db.list_materialized_views()
@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")
view = await db.create_materialized_view("copy", "people", with_no_data=True)
await view.refresh()
table = await db.open_table("people")
@@ -157,7 +359,10 @@ def test_bare_select_names_are_quoted(tmp_path):
db.create_table("odd_names", [{"order item": "widget", "select": 2}])
view = db.create_materialized_view(
"quoted", "odd_names", select=["order item", "select"]
"quoted",
"odd_names",
select=["order item", "select"],
with_no_data=True,
)
result = view.refresh()
assert result.rows_written == 1
@@ -166,19 +371,6 @@ def test_bare_select_names_are_quoted(tmp_path):
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")
@@ -235,6 +427,17 @@ def test_namespace_connection_materialized_views(tmp_path):
with pytest.raises(ValueError, match="not a materialized view"):
db.open_materialized_view("people")
create_job = db.create_materialized_view_async(
"job_view", "people", with_no_data=True
)
assert create_job.wait() is None
refresh_job = db.open_materialized_view("job_view").refresh_async()
assert refresh_job.wait().rows_written == 2
assert db.drop_materialized_view_async("job_view").wait() is None
db.drop_materialized_view("adults")
assert db.list_materialized_views() == []
@pytest.mark.asyncio
async def test_async_namespace_connection_materialized_views(tmp_path):
@@ -266,3 +469,51 @@ async def test_async_namespace_connection_materialized_views(tmp_path):
handle._route_pushdown_to_rust == through_namespace._route_pushdown_to_rust
)
assert handle._namespace_path == through_namespace._namespace_path
create_job = await db.create_materialized_view_async(
"job_view", "people", with_no_data=True
)
assert await create_job.wait() is None
job_view = await db.open_materialized_view("job_view")
refresh_job = await job_view.refresh_async()
assert (await refresh_job.wait()).rows_written == 2
drop_job = await db.drop_materialized_view_async("job_view")
assert await drop_job.wait() is None
await db.drop_materialized_view("adults")
assert await db.list_materialized_views() == []
def test_namespaced_select_kind_is_read_and_unknown_kinds_are_refused():
import json
import pyarrow as pa
from lancedb.materialized_view import _definition_from_schema
def schema_with(definition: dict) -> pa.Schema:
return pa.schema([pa.field("id", pa.int32())]).with_metadata(
{b"mv.definition": json.dumps(definition).encode()}
)
# "namespaced_select" is the namespaced form of "select": same shape,
# a separate kind so readers that predate it refuse instead of
# resolving the source at the root.
definition = _definition_from_schema(
schema_with(
{
"kind": "namespaced_select",
"source_table": "people",
"source_namespace": ["ns"],
"projections": [{"output": "name", "expression": "name"}],
}
),
"v",
)
assert definition.source_table == "people"
assert definition.source_namespace == ["ns"]
with pytest.raises(NotImplementedError, match="cannot refresh"):
_definition_from_schema(
schema_with({"kind": "select_v3", "source_table": "people"}), "v"
)
python/python/tests/test_namespace.py
+7 -1
View File
@@ -193,7 +193,13 @@ class TestNamespaceConnection:
),
)
table = db.create_table("blob_table", data, namespace_path=["test_ns"])
# Legacy v1 blob columns are only writable at file version <= 2.1.
table = db.create_table(
"blob_table",
data,
namespace_path=["test_ns"],
storage_options={"new_table_data_storage_version": "2.1"},
)
df = table.to_pandas(blob_mode="lazy").sort_values("id")
blob = df["blob"].iloc[0]
python/python/tests/test_namespace_no_pylance.py
+149
View File
@@ -0,0 +1,149 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Namespace operations must not require the optional ``pylance`` dependency.
The sync ``lancedb.connect()`` connection used to route namespace operations
through the Python ``lance_namespace`` client, whose ``dir`` implementation
lives in ``lance.namespace`` (shipped by the optional ``pylance`` extra). On an
install without that extra, even ``db.list_namespaces()`` failed, while the
async API worked because it goes straight to the native Rust connection.
The ``without_pylance`` fixture below simulates a missing ``pylance`` so these
tests fail on any environment when the native routing regresses. The ground
truth remains the "Test without pylance or pandas" CI job, which runs this file
with ``pylance`` and ``pandas`` actually uninstalled -- so nothing here may
import either at module scope.
"""
import sys
from importlib import import_module
from importlib.abc import MetaPathFinder
import lancedb
import pyarrow as pa
import pytest
def _is_lance(module_name: str) -> bool:
# "lance_namespace" is a separate, non-optional package -- leave it alone.
return module_name == "lance" or module_name.startswith("lance.")
class _BlockLanceImports(MetaPathFinder):
def find_spec(self, fullname, path=None, target=None):
if _is_lance(fullname):
raise ModuleNotFoundError(f"No module named {fullname!r}", name=fullname)
return None
@pytest.fixture
def without_pylance(monkeypatch):
"""Make ``lance`` unimportable, as on an install without the pylance extra.
Uninstalled means absent from ``sys.modules`` too, not merely unimportable:
LanceDB has code that branches on ``"lance" in sys.modules``, so a fixture
that poisons the entry instead of removing it would fail tests that a real
install passes.
"""
for name in list(sys.modules):
if _is_lance(name):
monkeypatch.delitem(sys.modules, name)
monkeypatch.setattr(sys, "meta_path", [_BlockLanceImports(), *sys.meta_path])
def _schema() -> pa.Schema:
return pa.schema([pa.field("id", pa.int64())])
def test_fixture_matches_an_uninstalled_pylance(without_pylance):
"""Guard the guard: the other tests are meaningless if lance stays importable."""
assert "lance" not in sys.modules
with pytest.raises(ModuleNotFoundError):
import_module("lance.namespace")
def test_list_namespaces_on_sync_connection(tmp_path, without_pylance):
"""The original reproducer: this alone used to raise on a sync connection."""
db = lancedb.connect(tmp_path)
assert db.list_namespaces().namespaces == []
def test_sync_namespace_lifecycle(tmp_path, without_pylance):
db = lancedb.connect(tmp_path)
db.create_namespace(["child"])
assert db.list_namespaces().namespaces == ["child"]
assert db.list_namespaces(namespace_path=["child"]).namespaces == []
db.describe_namespace(["child"])
db.create_namespace(["child", "grandchild"])
assert db.list_namespaces(namespace_path=["child"]).namespaces == ["grandchild"]
db.drop_namespace(["child", "grandchild"])
db.drop_namespace(["child"])
assert db.list_namespaces().namespaces == []
def test_sync_namespaced_table_lifecycle(tmp_path, without_pylance):
db = lancedb.connect(tmp_path)
db.create_namespace(["child"])
table = db.create_table("tbl", schema=_schema(), namespace_path=["child"])
assert table.namespace == ["child"]
table.add([{"id": 1}])
assert db.list_tables(namespace_path=["child"]).tables == ["tbl"]
assert db.list_tables().tables == []
opened = db.open_table("tbl", namespace_path=["child"])
assert opened.namespace == ["child"]
assert opened.count_rows() == 1
assert opened.search().limit(5).to_arrow().num_rows == 1
db.drop_table("tbl", namespace_path=["child"])
assert db.list_tables(namespace_path=["child"]).tables == []
db.drop_namespace(["child"])
def test_sync_root_table_lifecycle(tmp_path, without_pylance):
"""Root-namespace tables share the namespace plumbing, so cover them too."""
db = lancedb.connect(tmp_path)
table = db.create_table("tbl", schema=_schema())
table.add([{"id": 1}])
assert db.table_names() == ["tbl"]
assert "tbl" in db
assert db["tbl"].count_rows() == 1
db.drop_table("tbl")
assert db.table_names() == []
@pytest.mark.asyncio
async def test_async_namespace_lifecycle(tmp_path, without_pylance):
db = await lancedb.connect_async(tmp_path)
await db.create_namespace(["child"])
assert (await db.list_namespaces()).namespaces == ["child"]
table = await db.create_table("tbl", schema=_schema(), namespace_path=["child"])
await table.add([{"id": 1}])
assert (await db.list_tables(namespace_path=["child"])).tables == ["tbl"]
assert await table.count_rows() == 1
await db.drop_table("tbl", namespace_path=["child"])
await db.drop_namespace(["child"])
assert (await db.list_namespaces()).namespaces == []
def test_namespace_client_still_requires_pylance(tmp_path, without_pylance):
"""``namespace_client()`` is the one namespace API that opts into pylance.
It hands out a Python ``LanceNamespace``, so it cannot be served natively.
Pinned here so the boundary stays explicit and the error stays actionable.
"""
db = lancedb.connect(tmp_path)
with pytest.raises(ValueError, match="lance.namespace.DirectoryNamespace"):
db.namespace_client()
python/python/tests/test_package_metadata.py
+9
View File
@@ -9,6 +9,15 @@ from pathlib import Path
import pytest
@pytest.mark.parametrize("name", ["JobInfo", "JobDescription", "JobFailureInfo"])
def test_job_metadata_types_have_resolvable_modules(name):
"""Documentation tools resolve re-exports through each type's module."""
public_type = getattr(importlib.import_module("lancedb.job"), name)
defining_module = importlib.import_module(public_type.__module__)
assert getattr(defining_module, public_type.__name__, None) is public_type
def test_pyo3_abi_matches_minimum_supported_python():
project_dir = Path(__file__).parents[2]
pyproject = (project_dir / "pyproject.toml").read_text()
python/python/tests/test_query.py
+53 -10
View File
@@ -40,6 +40,10 @@ from utils import exception_output
from importlib.util import find_spec
# Legacy v1 blob columns are only writable at file version <= 2.1.
LEGACY_BLOB_STORAGE_OPTIONS = {"new_table_data_storage_version": "2.1"}
def _blob_query_data():
return pa.table(
{
@@ -119,13 +123,17 @@ def _assert_blob_bytes_projection(df):
def _blob_query_table(db, name, blob_schema):
if blob_schema == "v1":
return db.create_table(name, _blob_query_data())
return db.create_table(
name, _blob_query_data(), storage_options=LEGACY_BLOB_STORAGE_OPTIONS
)
return _create_blob_v2_query_table(db, name)
async def _blob_query_table_async(db, name, blob_schema):
if blob_schema == "v1":
return await db.create_table(name, _blob_query_data())
return await db.create_table(
name, _blob_query_data(), storage_options=LEGACY_BLOB_STORAGE_OPTIONS
)
return await _create_blob_v2_query_table_async(db, name)
@@ -275,7 +283,9 @@ async def test_query_to_pandas_kwargs(table, table_async):
def test_plain_scan_query_to_pandas_blob_modes(tmp_db, blob_mode):
pytest.importorskip("lance")
table = tmp_db.create_table(
f"test_query_to_pandas_blob_{blob_mode}", _blob_query_data()
f"test_query_to_pandas_blob_{blob_mode}",
_blob_query_data(),
storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
)
df = (
@@ -322,7 +332,9 @@ def test_plain_scan_query_to_pandas_blob_mode_does_not_collect_arrow(
):
pytest.importorskip("lance")
table = tmp_db.create_table(
"test_query_to_pandas_blob_no_arrow_collect", _blob_query_data()
"test_query_to_pandas_blob_no_arrow_collect",
_blob_query_data(),
storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
)
query = table.search().where("id = 1").select(["id", "blob"])
@@ -347,7 +359,9 @@ def test_plain_scan_query_to_pandas_blob_descriptions_flatten_uses_scanner(
):
pytest.importorskip("lance")
table = tmp_db.create_table(
"test_query_to_pandas_blob_desc_flatten", _blob_query_data()
"test_query_to_pandas_blob_desc_flatten",
_blob_query_data(),
storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
)
query = table.search().where("id = 1").select(["id", "blob"])
@@ -365,7 +379,11 @@ def test_plain_scan_query_to_pandas_blob_descriptions_flatten_uses_scanner(
def test_plain_scan_query_to_pandas_scanner_state(tmp_db):
pytest.importorskip("lance")
data = _blob_query_data()
table = tmp_db.create_table("test_query_to_pandas_scanner_state", data.slice(0, 2))
table = tmp_db.create_table(
"test_query_to_pandas_scanner_state",
data.slice(0, 2),
storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
)
table.add(data.slice(2, 2))
fragments = table.to_lance().get_fragments()
@@ -400,7 +418,9 @@ def test_plain_scan_query_to_pandas_scanner_state(tmp_db):
async def test_async_plain_scan_query_to_pandas_blob_projection(tmp_db_async):
pytest.importorskip("lance")
table = await tmp_db_async.create_table(
"test_async_query_to_pandas_blob_projection", _blob_query_data()
"test_async_query_to_pandas_blob_projection",
_blob_query_data(),
storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
)
lazy_df = await (
@@ -452,7 +472,9 @@ async def test_async_plain_scan_query_to_pandas_blob_mode_does_not_collect_arrow
):
pytest.importorskip("lance")
table = await tmp_db_async.create_table(
"test_async_query_to_pandas_blob_no_arrow_collect", _blob_query_data()
"test_async_query_to_pandas_blob_no_arrow_collect",
_blob_query_data(),
storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
)
query = table.query().where("id = 1").select(["id", "blob"])
@@ -474,7 +496,11 @@ async def test_async_plain_scan_query_to_pandas_blob_mode_does_not_collect_arrow
def test_vector_query_to_pandas_blob_mode_requires_native_path(tmp_db):
pytest.importorskip("lance")
table = tmp_db.create_table("test_vector_query_blob_mode", _blob_query_data())
table = tmp_db.create_table(
"test_vector_query_blob_mode",
_blob_query_data(),
storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
)
with pytest.raises(RuntimeError, match="Lance native pandas conversion"):
table.search([1.0, 0.0]).select(["blob", "vector"]).limit(1).to_pandas(
@@ -485,7 +511,9 @@ def test_vector_query_to_pandas_blob_mode_requires_native_path(tmp_db):
def test_vector_query_to_pandas_blob_descriptions_requires_plain_scan(tmp_db):
pytest.importorskip("lance")
table = tmp_db.create_table(
"test_vector_query_blob_descriptions", _blob_query_data()
"test_vector_query_blob_descriptions",
_blob_query_data(),
storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
)
with pytest.raises(RuntimeError, match="plain scan query"):
@@ -1923,6 +1951,21 @@ def test_take_queries(tmp_path):
17,
]
# Duplicate offsets are occurrences, not set members. Ordering is unspecified.
assert sorted(table.take_offsets([5, 2, 5, 17]).to_pandas()["idx"].to_list()) == [
2,
5,
5,
17,
]
# Converting a take builder to its serializable query representation must
# retain occurrence metadata and execute with the same multiplicity.
query = table.take_offsets([5, 2, 5, 17]).select(["idx"]).to_query_object()
assert query.take_offsets == [5, 2, 5, 17]
converted = table._execute_query(query).read_all()
assert sorted(converted["idx"].to_pylist()) == [2, 5, 5, 17]
# Take by row id
assert list(
sorted(table.take_row_ids([5, 2, 17]).to_pandas()["idx"].to_list())
python/python/tests/test_remote_db.py
+156 -21
View File
@@ -479,24 +479,49 @@ def test_remote_permutation_is_picklable():
match = re.search(
r"_rowoffset\s+in\s+\((.*?)\)", body["filter"], re.IGNORECASE
)
offsets = [int(o.strip()) for o in match.group(1).split(",")]
offsets = list(
dict.fromkeys(int(o.strip()) for o in match.group(1).split(","))
)
else:
offsets = list(range(len(rows)))
table = pa.table({"a": [rows[offset] for offset in offsets]})
columns = body.get("columns") or ["a"]
table = pa.table(
{
column: (
[rows[offset] for offset in offsets]
if column == "a"
else offsets
)
for column in columns
}
)
request.send_response(200)
request.send_header("Content-Type", "application/vnd.apache.arrow.file")
request.end_headers()
with pa.ipc.new_file(request.wfile, schema=table.schema) as writer:
writer.write_table(table)
writer.write_table(table, max_chunksize=2)
else:
request.send_response(404)
request.end_headers()
with mock_lancedb_connection(handler) as db:
permutation = Permutation.identity(db.open_table("test"))
table = db.open_table("test")
assert table.take_offsets([0, 2, 0, 4]).to_list() == [
{"a": 0},
{"a": 0},
{"a": 2},
{"a": 4},
]
permutation = Permutation.identity(table)
restored = pickle.loads(pickle.dumps(permutation))
assert restored.__getitems__([0, 2, 4]) == [{"a": 0}, {"a": 2}, {"a": 4}]
assert restored.__getitems__([0, 2, 0, 4]) == [
{"a": 0},
{"a": 2},
{"a": 0},
{"a": 4},
]
def test_create_table_exist_ok():
@@ -795,11 +820,13 @@ def test_table_create_indices():
scalar_req = received_requests[0]
assert "name" in scalar_req
assert scalar_req["name"] == "custom_scalar_idx"
assert scalar_req["replace"] is False
# Check FTS index request has custom name
fts_req = received_requests[1]
assert "name" in fts_req
assert fts_req["name"] == "custom_fts_idx"
assert fts_req["replace"] is False
assert fts_req["block_size"] == 256
assert fts_req["custom_stop_words"] == ["cloud"]
@@ -807,6 +834,7 @@ def test_table_create_indices():
vector_req = received_requests[2]
assert "name" in vector_req
assert vector_req["name"] == "custom_vector_idx"
assert "replace" not in vector_req
table.wait_for_index(["custom_scalar_idx"], timedelta(seconds=2))
table.wait_for_index(
@@ -1079,6 +1107,9 @@ def test_remote_create_index_new_api():
table.create_index("text", config=FTS(block_size=256))
# IvfRq via new API
table.create_index("vector", config=IvfRq(distance_type="l2"))
table.create_index(
"vector", config=IvfPq(distance_type="l2"), replace=False
)
# Legacy index_type="IVF_RQ" routes to IvfRq config under the hood.
with pytest.warns(DeprecationWarning, match="create_index"):
@@ -1088,15 +1119,17 @@ def test_remote_create_index_new_api():
num_partitions=8,
)
assert len(received_requests) == 5
assert len(received_requests) == 6
assert [req["column"] for req in received_requests] == [
"vector",
"category",
"text",
"vector",
"vector",
"vector",
]
assert received_requests[2]["block_size"] == 256
assert received_requests[4]["replace"] is False
def test_table_wait_for_index_timeout():
@@ -2434,7 +2467,7 @@ def test_remote_blob_byte_apis_not_supported_on_old_server():
def test_remote_connection_jobs_surface():
from lancedb.exceptions import JobFailedError
from lancedb.exceptions import JobFailedError, JobNotFoundError
schema = pa.schema([("state", pa.string())])
batch = pa.record_batch([pa.array(["created", "done"])], schema=schema)
@@ -2442,6 +2475,7 @@ def test_remote_connection_jobs_surface():
with pa.ipc.new_stream(sink, schema) as writer:
writer.write_batch(batch)
events_body = sink.getvalue().to_pybytes()
query_events_payloads = []
def handler(request):
content_len = int(request.headers.get("Content-Length", 0))
@@ -2479,6 +2513,22 @@ def test_remote_connection_jobs_surface():
request.end_headers()
request.wfile.write(json.dumps(rsp).encode())
elif request.path == "/v1/jobs/describe":
if payload["job_id"] == "job-2":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(
json.dumps(
dict(
job_id="job-2",
job_type="refresh_column",
job_state="DONE",
creation_ms=2000,
result=dict(rows_assigned=1000000, rows_failed=0),
)
).encode()
)
return
if payload["job_id"] != "job-1":
request.send_response(404)
request.end_headers()
@@ -2510,7 +2560,7 @@ def test_remote_connection_jobs_surface():
request.end_headers()
request.wfile.write(b'{"job_id": "job-1"}')
elif request.path == "/v1/jobs/query_events":
assert payload["job_id"] == "job-1"
query_events_payloads.append(payload)
request.send_response(200)
request.send_header("Content-Type", "application/vnd.apache.arrow.stream")
request.end_headers()
@@ -2526,24 +2576,109 @@ def test_remote_connection_jobs_surface():
assert jobs[0].table == "t1"
assert jobs[1].state == "finished"
description = db.get_job("job-1")
assert description.job_type == "create_index"
assert description.state == "failed"
assert json.loads(description.spec_json) == {"column": "vec"}
assert description.failure.message == "worker died"
assert description.failure.retryable is True
assert db.get_job("missing") is None
assert db.cancel_job("job-1") is True
assert db.cancel_job("missing") is False
batches = db.job_history("job-1")
assert len(batches) == 1
assert batches[0].num_rows == 2
assert batches[0].column("state").to_pylist() == ["created", "done"]
# Opening a job hands back a populated handle; a missing one fails.
with pytest.raises(JobNotFoundError, match="missing"):
db.open_job("missing")
finished = db.open_job("job-2")
assert finished.state == "finished"
assert finished.result == {"rows_assigned": 1000000, "rows_failed": 0}
job = db.job("job-1")
job = db.open_job("job-1")
assert job.id == "job-1"
# Opening already populated the handle.
assert job.state == "failed"
assert job.spec == {"column": "vec"}
assert job.failure.message == "worker died"
assert job.status() == "failed"
with pytest.raises(JobFailedError, match="worker died"):
job.wait(timeout=timedelta(seconds=5))
def test_remote_job_handle_reports_its_own_detail():
schema = pa.schema([("state", pa.string())])
batch = pa.record_batch([pa.array(["claim_complete"])], schema=schema)
sink = pa.BufferOutputStream()
with pa.ipc.new_stream(sink, schema) as writer:
writer.write_batch(batch)
events_body = sink.getvalue().to_pybytes()
event_payloads = []
def handler(request):
content_len = int(request.headers.get("Content-Length", 0))
body = request.rfile.read(content_len) if content_len > 0 else b""
payload = json.loads(body) if body else {}
if request.path == "/v1/jobs/describe":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(
json.dumps(
dict(
job_id="job-1",
job_type="refresh_column",
job_state="DONE",
creation_ms=2000,
spec=dict(column="vec"),
result=dict(rows_assigned=1000000),
)
).encode()
)
elif request.path == "/v1/jobs/query_events":
event_payloads.append(payload)
request.send_response(200)
request.send_header("Content-Type", "application/vnd.apache.arrow.stream")
request.end_headers()
request.wfile.write(events_body)
else:
request.send_response(404)
request.end_headers()
with mock_lancedb_connection(handler) as db:
job = db.open_job("job-1")
# Opening populates the handle in the same round trip.
assert job.state == "finished"
job.refresh()
assert job.job_type == "refresh_column"
assert job.creation_ms == 2000
assert job.spec == {"column": "vec"}
assert job.result == {"rows_assigned": 1000000}
assert job.failure is None
# The JSON payloads stay reachable, but as internal APIs.
assert json.loads(job._spec_json) == {"column": "vec"}
assert json.loads(job._result_json) == {"rows_assigned": 1000000}
# print() shows everything the handle knows and nothing it does not.
# print() lays every known field out on its own line, with the JSON
# payloads indented rather than crammed onto one line.
assert repr(job) == "\n".join(
[
"Job(",
" id='job-1',",
" state='finished',",
" job_type='refresh_column',",
" creation_ms=2000,",
" spec={",
' "column": "vec"',
" },",
" result={",
' "rows_assigned": 1000000',
" },",
")",
]
)
# Nothing it does not know shows up.
assert "failure" not in repr(job)
events = job.events(filter="state = 'claim_complete'", limit=500)
assert isinstance(events, pa.Table)
assert events.column("state").to_pylist() == ["claim_complete"]
# The handle supplies job_id; the caller only narrows the query.
assert event_payloads[-1] == {
"job_id": "job-1",
"limit": 500,
"filter": "state = 'claim_complete'",
}
python/python/tests/test_rerankers.py
+1 -1
View File
@@ -81,7 +81,7 @@ def get_test_table(tmp_path):
"but his son was mortal",
"there hasn't been a good battlefield game since 2142",
"I wish they would make another one",
"campains are not as good as they used to be",
"campaigns are not as good as they used to be",
"Multiplayer and open world games have destroyed the single player experience",
"Maybe the future is console games",
"I don't know",
python/python/tests/test_sql.py
+162
View File
@@ -0,0 +1,162 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from uuid import UUID
import pytest
import pyarrow as pa
import lancedb
from lancedb import _lancedb
from lancedb.arrow import AsyncRecordBatchReader
from lancedb.db import AsyncConnection
from lancedb.remote.db import RemoteDBConnection
from lancedb.sql import AsyncQuery, Query
NIL_QUERY_ID = UUID(int=0)
class FakeNativeQuery:
id = UUID("0198f1b2-c3d4-7e5f-8123-456789abcdef")
async def reader(self):
return pa.table({"value": [1, 2]})
class FakeNativeConnection:
async def execute_query_async(self, query, *, default_namespace_path=None):
return FakeNativeQuery()
class FakeAsyncConnection:
async def execute_query_async(self, query, *, default_namespace_path=None):
return AsyncQuery(FakeNativeQuery())
def remote_connection(sql_host_override=None):
return lancedb.connect(
"db://analytics",
api_key="test-key",
host_override="http://localhost:10024",
sql_host_override=sql_host_override,
)
def test_sql_is_connection_scoped():
assert hasattr(lancedb, "sql")
assert not callable(lancedb.sql)
assert not hasattr(_lancedb, "sql")
assert not hasattr(remote_connection(), "sql")
assert hasattr(remote_connection(), "execute_query")
assert hasattr(remote_connection(), "execute_query_async")
assert hasattr(remote_connection(), "describe_query")
def test_query_id_is_uuid():
query = AsyncQuery(FakeNativeQuery())
assert isinstance(query.id, UUID)
assert Query(query).id == query.id
def test_connection_serializes_sql_host_override():
endpoint = "grpc+tls://sql.example.com:10026"
restored = lancedb.deserialize_conn(
remote_connection(sql_host_override=endpoint).serialize()
)
assert restored.sql_host_override == endpoint
@pytest.mark.asyncio
async def test_async_sql_reader_is_record_batch_stream():
reader = await AsyncQuery(FakeNativeQuery()).reader()
assert isinstance(reader, AsyncRecordBatchReader)
assert (await reader.read_all())[0].column(0).to_pylist() == [1, 2]
def test_sync_sql_reader_is_record_batch_reader():
reader = Query(AsyncQuery(FakeNativeQuery())).reader()
assert isinstance(reader, pa.RecordBatchReader)
assert reader.read_all().column(0).to_pylist() == [1, 2]
def test_execute_query_returns_blocking_reader():
connection = RemoteDBConnection.__new__(RemoteDBConnection)
connection._conn = FakeAsyncConnection()
reader = connection.execute_query("SELECT 1")
assert isinstance(reader, pa.RecordBatchReader)
assert reader.read_all().column(0).to_pylist() == [1, 2]
@pytest.mark.asyncio
async def test_async_execute_query_returns_async_reader():
connection = AsyncConnection(FakeNativeConnection())
reader = await connection.execute_query("SELECT 1")
assert isinstance(reader, AsyncRecordBatchReader)
assert (await reader.read_all())[0].column(0).to_pylist() == [1, 2]
def test_local_connection_rejects_sql(tmp_path):
connection = lancedb.connect(tmp_path)
with pytest.raises(NotImplementedError, match="SQL"):
connection.execute_query("SELECT 1")
with pytest.raises(NotImplementedError, match="SQL"):
connection.execute_query_async("SELECT 1")
with pytest.raises(NotImplementedError, match="SQL"):
connection.describe_query(NIL_QUERY_ID)
@pytest.mark.asyncio
async def test_local_async_connection_rejects_sql(tmp_path):
connection = await lancedb.connect_async(tmp_path)
with pytest.raises(NotImplementedError, match="SQL"):
await connection.execute_query("SELECT 1")
with pytest.raises(NotImplementedError, match="SQL"):
await connection.execute_query_async("SELECT 1")
with pytest.raises(NotImplementedError, match="SQL"):
await connection.describe_query(NIL_QUERY_ID)
@pytest.mark.asyncio
async def test_async_namespace_connection_rejects_sql(tmp_path):
connection = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)})
with pytest.raises(NotImplementedError, match="SQL"):
await connection.execute_query("SELECT 1")
with pytest.raises(NotImplementedError, match="SQL"):
await connection.execute_query_async("SELECT 1")
with pytest.raises(NotImplementedError, match="SQL"):
await connection.describe_query(NIL_QUERY_ID)
def test_describe_query_requires_uuid():
with pytest.raises(TypeError, match="UUID"):
remote_connection().describe_query(str(NIL_QUERY_ID))
@pytest.mark.parametrize(
"default_namespace_path",
["public", ("public",), [1]],
)
def test_execute_query_async_requires_namespace_path_list(default_namespace_path):
with pytest.raises(ValueError, match="default_namespace_path"):
remote_connection().execute_query_async(
"SELECT 1", default_namespace_path=default_namespace_path
)
def test_execute_query_async_rejects_invalid_endpoint():
connection = remote_connection(sql_host_override="invalid://localhost")
with pytest.raises(ValueError, match="sql_host_override"):
connection.execute_query_async("SELECT 1")
@pytest.mark.parametrize(
"default_namespace_path",
[[""], ["café"], ["pub\tlic"], ["events$raw"]],
)
def test_execute_query_async_rejects_invalid_namespace_components(
default_namespace_path,
):
with pytest.raises(ValueError, match="default_namespace_path"):
remote_connection().execute_query_async(
"SELECT 1", default_namespace_path=default_namespace_path
)
python/python/tests/test_table.py
+235 -28
View File
@@ -64,15 +64,23 @@ async def _blob_v2_table_async(db: AsyncConnection, name: str):
return table
# Legacy v1 blob columns are only writable at file version <= 2.1.
LEGACY_BLOB_STORAGE_OPTIONS = {"new_table_data_storage_version": "2.1"}
def _blob_table(db: DBConnection, name: str, blob_schema: str):
if blob_schema == "v1":
return db.create_table(name, data=_blob_test_data())
return db.create_table(
name, data=_blob_test_data(), storage_options=LEGACY_BLOB_STORAGE_OPTIONS
)
return _blob_v2_table(db, name)
async def _blob_table_async(db: AsyncConnection, name: str, blob_schema: str):
if blob_schema == "v1":
return await db.create_table(name, data=_blob_test_data())
return await db.create_table(
name, data=_blob_test_data(), storage_options=LEGACY_BLOB_STORAGE_OPTIONS
)
return await _blob_v2_table_async(db, name)
@@ -147,7 +155,11 @@ def test_table_to_pandas_invalid_blob_mode_non_blob_table(tmp_db: DBConnection):
@pytest.mark.parametrize("blob_mode", ["lazy", "bytes", "descriptions"])
def test_table_to_pandas_blob_modes(tmp_db: DBConnection, blob_mode):
pytest.importorskip("lance")
table = tmp_db.create_table(f"test_to_pandas_blob_{blob_mode}", _blob_test_data())
table = tmp_db.create_table(
f"test_to_pandas_blob_{blob_mode}",
_blob_test_data(),
storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
)
df = table.to_pandas(blob_mode=blob_mode)
@@ -2682,6 +2694,43 @@ def test_merge_insert(mem_db: DBConnection):
)
def test_merge_insert_composite_key(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
data=pa.table(
{
"shard": ["a", "a", "b"],
"id": [1, 2, 1],
"val": ["x", "y", "z"],
}
),
)
# ("a", 1) matches an existing row and updates it. ("b", 2) agrees with an
# existing row on each key column separately but on neither pair, so it is
# an insert.
new_data = pa.table({"shard": ["a", "b"], "id": [1, 2], "val": ["X", "W"]})
res = (
table.merge_insert(["shard", "id"])
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(new_data)
)
assert res.num_updated_rows == 1
assert res.num_inserted_rows == 1
expected = pa.table(
{
"shard": ["a", "a", "b", "b"],
"id": [1, 2, 1, 2],
"val": ["X", "y", "z", "W"],
}
)
assert table.to_arrow().sort_by([("shard", "ascending"), ("id", "ascending")]) == (
expected
)
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")
@@ -2930,28 +2979,29 @@ async def test_merge_insert_async(mem_db_async: AsyncConnection):
assert (await table.to_arrow()).sort_by("a") == expected
def _json_arrow_table(schema, rows):
json_type = schema.field("j").type
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
)
@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
)
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", pa.json_())])
table = await mem_db_async.create_table("json_merge", schema=schema)
await table.add(json_table([("a", '{"k": 1}'), ("b", '{"k": 9}')]))
await table.add(_json_arrow_table(schema, [("a", '{"k": 1}'), ("b", '{"k": 9}')]))
await (
table.merge_insert("id")
.when_matched_update_all()
.execute(json_table([("a", '{"k": 2}')]))
.execute(_json_arrow_table(schema, [("a", '{"k": 2}')]))
)
rows = sorted(await table.query().to_list(), key=lambda row: row["id"])
@@ -2966,20 +3016,176 @@ async def test_merge_insert_encodes_json(mem_db_async: AsyncConnection):
@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)
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", pa.json_())])
table = await mem_db_async.create_table("json_add", schema=schema)
await table.add(data, on_bad_vectors="fill")
await table.add(
_json_arrow_table(schema, [("c", '{"k": 3}')]), on_bad_vectors="fill"
)
rows = await table.query().where("json_extract(j, '$.k') = '3'").to_list()
assert rows == [{"id": "c", "j": '{"k":3}'}]
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
@pytest.mark.asyncio
async def test_add_all_null_json_batch(mem_db_async: AsyncConnection):
"""A batch of dicts whose json values are all None infers as pa.null(), which used
to fail with a `json` vs `large_binary` schema mismatch. A row-at-a-time insert of
an optional json column always looks like this."""
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", pa.json_())])
table = await mem_db_async.create_table("json_nulls", schema=schema)
await table.add([{"id": "a", "j": None}])
assert await table.count_rows() == 1
# ... and again once real JSON has been written.
await table.add(_json_arrow_table(schema, [("b", '{"k": 9}')]))
await table.add([{"id": "c", "j": None}])
rows = sorted(await table.query().to_list(), key=lambda row: row["id"])
assert rows == [
{"id": "a", "j": None},
{"id": "b", "j": '{"k":9}'},
{"id": "c", "j": None},
]
# The nulls must not disturb reads of the column.
filtered = await table.query().where("json_extract(j, '$.k') = '9'").to_list()
assert filtered == [{"id": "b", "j": '{"k":9}'}]
assert len(await table.query().where("j IS NULL").to_list()) == 2
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
def test_add_all_null_json_batch_sync(mem_db: DBConnection):
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", pa.json_())])
table = mem_db.create_table("json_nulls_sync", schema=schema)
table.add([{"id": "a", "j": None}])
assert table.count_rows() == 1
assert table.to_arrow()["j"].to_pylist() == [None]
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
@pytest.mark.asyncio
@pytest.mark.parametrize(
("values", "expected"),
[
([None], [None]),
([None, '{"k": 1}'], [None, '{"k":1}']),
(['{"k": 2}'], ['{"k":2}']),
],
)
async def test_add_list_of_dicts_to_json_column(
mem_db_async: AsyncConnection, values, expected
):
schema = pa.schema([pa.field("id", pa.int64()), pa.field("value", pa.json_())])
table = await mem_db_async.create_table("json_list_add", schema=schema)
await table.add([{"id": idx, "value": value} for idx, value in enumerate(values)])
rows = (await table.to_arrow()).sort_by("id").to_pylist()
assert [row["value"] for row in rows] == expected
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
@pytest.mark.asyncio
async def test_add_list_of_dicts_to_nested_json_column(
mem_db_async: AsyncConnection,
):
json_field = pa.field("value", pa.json_())
info_field = pa.field("info", pa.struct([json_field]))
info = pa.StructArray.from_arrays(
[pa.array(['{"seed": 0}'], type=pa.json_())], fields=[json_field]
)
seed = pa.Table.from_arrays(
[pa.array([0], type=pa.int64()), info],
schema=pa.schema([pa.field("id", pa.int64()), info_field]),
)
table = await mem_db_async.create_table("nested_json_list_add", data=seed)
await table.add([{"id": 1, "info": {"value": '{"k": 1}'}}])
await table.add([{"id": 2, "info": {"value": '{"k": 2}'}}], on_bad_vectors="fill")
rows = (await table.to_arrow()).sort_by("id").to_pylist()
assert rows == [
{"id": 0, "info": {"value": '{"seed":0}'}},
{"id": 1, "info": {"value": '{"k":1}'}},
{"id": 2, "info": {"value": '{"k":2}'}},
]
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
@pytest.mark.asyncio
async def test_add_list_of_dicts_to_json_list_column(mem_db_async: AsyncConnection):
"""JSON inside a list must be JSONB-encoded, not stored as the raw text.
Storing raw text appends without error but leaves the column unreadable, so the
round trip is checked with a filter as well as by value.
"""
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("docs", pa.list_(pa.field("item", pa.json_()))),
]
)
table = await mem_db_async.create_table("json_list_add", schema=schema)
await table.add([{"id": 1, "docs": ['{"k": 1}', '{"k": 2}']}])
await table.add([{"id": 2, "docs": ['{"k": 3}']}], on_bad_vectors="fill")
rows = (await table.to_arrow()).sort_by("id").to_pylist()
assert rows == [
{"id": 1, "docs": ['{"k":1}', '{"k":2}']},
{"id": 2, "docs": ['{"k":3}']},
]
matched = await table.query().where("json_extract(docs[1], '$.k') = 3").to_arrow()
assert matched.column("id").to_pylist() == [2]
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
@pytest.mark.asyncio
async def test_add_map_of_json_values(mem_db_async: AsyncConnection):
"""JSON in a map's values needs the same encoding a list's items do."""
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("m", pa.map_(pa.string(), pa.json_())),
]
)
table = await mem_db_async.create_table(
"json_map_add",
schema=schema,
storage_options={"new_table_data_storage_version": "2.2"},
)
def batch(row_id: int, text: str) -> pa.Table:
return pa.table(
{
"id": pa.array([row_id], type=pa.int64()),
"m": pa.array([[("k", text)]], type=pa.map_(pa.string(), pa.string())),
}
)
await table.add(batch(1, '{"x": 1}'))
await table.add(batch(2, '{"x": 2}'), on_bad_vectors="fill")
rows = (await table.to_arrow()).sort_by("id").to_pylist()
assert rows == [
{"id": 1, "m": [("k", '{"x":1}')]},
{"id": 2, "m": [("k", '{"x":2}')]},
]
matched = (
await table.query()
.where("json_extract(element_at(m, 'k')[1], '$.x') = '2'")
.to_arrow()
)
assert matched.column("id").to_pylist() == [2]
def test_create_with_embedding_function(mem_db: DBConnection):
class MyTable(LanceModel):
text: str
@@ -3305,7 +3511,7 @@ def test_empty_query(mem_db: DBConnection):
# None is the same as default
df = table.search().select(["id"]).limit(None).to_arrow()
assert df.num_rows == 100
# invalid limist is the same as None, wihch is the same as default
# invalid limist is the same as None, which is the same as default
df = table.search().select(["id"]).limit(-1).to_arrow()
assert df.num_rows == 100
# valid limit should work
@@ -4016,7 +4222,7 @@ def test_stats(mem_db: DBConnection):
print(f"{stats=}")
assert stats == {
# Full on-disk size of the data file, footer and metadata included.
"total_bytes": 633,
"total_bytes": 637,
"num_rows": 2,
"num_indices": 0,
"fragment_stats": {
@@ -4228,13 +4434,14 @@ def test_refresh_column_async_returns_job(tmp_path):
assert result.rows_failed == 0
assert result.rows_remaining == 0
assert result.source_version == 2
assert result.published_version == 3
# The fill lands at 3; the stamp recording its inputs is published at 4.
assert result.published_version == 4
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.source_version == 4
assert no_op.published_version is None
# Bad input raises at the call, not through the job.
@@ -4253,6 +4460,6 @@ async def test_refresh_column_async_job_async_table(tmp_path):
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.rows_assigned == 1
assert result.source_version == 2
assert result.published_version == 3
assert result.published_version == 4
assert await job.status() == "finished"
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
python/src/catalog.rs
+125
View File
@@ -0,0 +1,125 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::time::Duration;
use lancedb::catalog::{
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
};
use pyo3::exceptions::PyValueError;
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pyfunction, pymethods};
use crate::connection::{Connection, PyClientConfig};
use crate::error::PythonErrorExt;
use crate::runtime::future_into_py;
#[pyclass]
pub struct Catalog {
inner: CatalogConnection,
}
#[pymethods]
impl Catalog {
#[getter]
fn uri(&self) -> &str {
self.inner.uri()
}
#[pyo3(signature = (name, *, exist_ok=false))]
fn create_database<'py>(
self_: PyRef<'py, Self>,
name: String,
exist_ok: bool,
) -> PyResult<Bound<'py, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner
.create_database(CreateDatabaseRequest::new(name).exist_ok(exist_ok))
.await
.map(Connection::new)
.infer_error()
})
}
fn connect_database<'py>(self_: PyRef<'py, Self>, name: String) -> PyResult<Bound<'py, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner
.connect_database(name)
.await
.map(Connection::new)
.infer_error()
})
}
#[pyo3(signature = (name, *, ignore_missing=false))]
fn drop_database<'py>(
self_: PyRef<'py, Self>,
name: String,
ignore_missing: bool,
) -> PyResult<Bound<'py, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner
.drop_database(DropDatabaseRequest::new(name).ignore_missing(ignore_missing))
.await
.infer_error()
})
}
#[pyo3(signature = (*, limit=None, page_token=None))]
fn list_databases<'py>(
self_: PyRef<'py, Self>,
limit: Option<u32>,
page_token: Option<String>,
) -> PyResult<Bound<'py, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let mut request = ListDatabasesRequest::default();
request.limit = limit;
request.page_token = page_token;
let response = inner.list_databases(request).await.infer_error()?;
Ok((response.databases, response.page_token))
})
}
}
#[pyfunction]
#[pyo3(signature = (endpoint, *, api_key=None, client_config=None, sql_host_override=None, read_consistency_interval=None, oauth_config=None))]
pub fn connect_catalog(
py: Python<'_>,
endpoint: String,
api_key: Option<String>,
client_config: Option<PyClientConfig>,
sql_host_override: Option<String>,
read_consistency_interval: Option<f64>,
oauth_config: Option<crate::oauth::PyOAuthConfig>,
) -> PyResult<Bound<'_, PyAny>> {
let interval = read_consistency_interval
.map(Duration::try_from_secs_f64)
.transpose()
.map_err(|err| {
PyValueError::new_err(format!("Invalid read consistency interval: {err}"))
})?;
future_into_py(py, async move {
let mut builder = lancedb::connect_catalog(endpoint);
if let Some(api_key) = api_key {
builder = builder.api_key(api_key);
}
if let Some(config) = client_config {
builder = builder.client_config(config.into());
}
if let Some(endpoint) = sql_host_override {
builder = builder.sql_host_override(endpoint);
}
if let Some(interval) = interval {
builder = builder.read_consistency_interval(interval);
}
if let Some(config) = oauth_config {
builder = builder.oauth_config(config.try_into().infer_error()?);
}
Ok(Catalog {
inner: builder.execute().await.infer_error()?,
})
})
}
python/src/connection.rs
+238 -39
View File
@@ -13,11 +13,7 @@ use crate::{
runtime::future_into_py,
table::Table,
};
use arrow::{
datatypes::Schema,
ffi_stream::ArrowArrayStreamReader,
pyarrow::{FromPyArrow, ToPyArrow},
};
use arrow::{datatypes::Schema, ffi_stream::ArrowArrayStreamReader, pyarrow::FromPyArrow};
use lancedb::{
connection::Connection as LanceConnection,
connection::NamespaceClientPushdownOperation,
@@ -28,7 +24,7 @@ use pyo3::{
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
exceptions::{PyRuntimeError, PyValueError},
pyclass, pyfunction, pymethods,
types::{PyDict, PyDictMethods, PyList, PyListMethods},
types::{PyAnyMethods, PyDict, PyDictMethods, PyList},
};
#[pyclass]
@@ -86,6 +82,24 @@ impl Connection {
}
}
fn parse_default_namespace_path(path: Option<Bound<'_, PyAny>>) -> PyResult<Vec<String>> {
match path {
Some(path) => {
if !path.is_instance_of::<PyList>() {
return Err(PyValueError::new_err(
"Connection.execute_query_async default_namespace_path must be a list",
));
}
path.extract::<Vec<String>>().map_err(|_| {
PyValueError::new_err(
"Connection.execute_query_async default_namespace_path components must be strings",
)
})
}
None => Ok(vec!["public".to_string()]),
}
}
#[pymethods]
impl Connection {
fn __repr__(&self) -> String {
@@ -108,6 +122,40 @@ impl Connection {
self.get_inner().map(|inner| inner.uri().to_string())
}
#[pyo3(signature = (query, *, default_namespace_path=None))]
pub fn execute_query_async<'a>(
self_: PyRef<'a, Self>,
query: String,
default_namespace_path: Option<Bound<'_, PyAny>>,
) -> PyResult<Bound<'a, PyAny>> {
let inner = self_.get_inner()?.clone();
let default_namespace_path = parse_default_namespace_path(default_namespace_path)?;
future_into_py(self_.py(), async move {
let operation = inner
.execute_query_async(query)
.default_namespace_path(default_namespace_path);
operation
.execute()
.await
.map(crate::sql::Query::new)
.infer_error()
})
}
pub fn describe_query<'a>(
self_: PyRef<'a, Self>,
query_id: uuid::Uuid,
) -> PyResult<Bound<'a, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
inner
.describe_query(query_id)
.await
.map(crate::sql::QueryDescription::from)
.infer_error()
})
}
#[pyo3(signature = ())]
pub fn get_read_consistency_interval(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
@@ -333,7 +381,7 @@ impl Connection {
})
}
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None))]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
pub fn create_materialized_view(
self_: PyRef<'_, Self>,
name: String,
@@ -341,6 +389,7 @@ impl Connection {
projections: Option<Vec<(String, String)>>,
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
@@ -354,16 +403,79 @@ impl Connection {
if let Some(limit) = limit {
builder = builder.limit(limit);
}
let view = builder.execute().await.infer_error()?;
builder = builder.with_no_data(with_no_data);
let view = Box::pin(builder.execute()).await.infer_error()?;
Ok(Table::new(view.table().clone()))
})
}
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
pub fn create_materialized_view_async(
self_: PyRef<'_, Self>,
name: String,
source: String,
projections: Option<Vec<(String, String)>>,
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
) -> 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 job = Box::pin(builder.with_no_data(with_no_data).execute_async())
.await
.infer_error()?;
Ok(crate::job::Job::new(job))
})
}
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<_>>())
Ok(views)
})
}
#[pyo3(signature = (name, namespace_path=None))]
pub fn drop_materialized_view(
self_: PyRef<'_, Self>,
name: String,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner
.drop_materialized_view(name, &namespace_path)
.await
.infer_error()
})
}
#[pyo3(signature = (name, namespace_path=None))]
pub fn drop_materialized_view_async(
self_: PyRef<'_, Self>,
name: String,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner
.drop_materialized_view_async(name, &namespace_path)
.await
.infer_error()
.map(crate::job::Job::new)
})
}
@@ -592,9 +704,12 @@ impl Connection {
})
}
pub fn job(&self, job_id: String) -> PyResult<crate::job::Job> {
let inner = self.get_inner()?.clone();
Ok(crate::job::Job::new(inner.job(job_id).infer_error()?))
pub fn open_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let job = inner.open_job(&job_id).await.infer_error()?;
Ok(crate::job::Job::new(job))
})
}
pub fn create_function_async(
@@ -629,6 +744,109 @@ impl Connection {
})
}
pub fn list_functions(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
inner
.list_functions()
.await
.infer_error()?
.into_iter()
.map(|function| function.to_canonical_json().infer_error())
.collect::<PyResult<Vec<_>>>()
})
}
pub fn drop_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.drop_function(name, version).await.infer_error()
})
}
#[pyo3(signature = (name, value, namespace_path=None))]
pub fn create_secret(
self_: PyRef<'_, Self>,
name: String,
value: String,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner
.create_secret(name, value, &namespace_path)
.await
.infer_error()
})
}
#[pyo3(signature = (name, value, namespace_path=None))]
pub fn alter_secret(
self_: PyRef<'_, Self>,
name: String,
value: String,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner
.alter_secret(name, value, &namespace_path)
.await
.infer_error()
})
}
#[pyo3(signature = (namespace_path=None))]
pub fn list_secrets(
self_: PyRef<'_, Self>,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner.list_secrets(&namespace_path).await.infer_error()
})
}
#[pyo3(signature = (name, namespace_path=None))]
pub fn drop_secret(
self_: PyRef<'_, Self>,
name: String,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
inner.drop_secret(name, &namespace_path).await.infer_error()
})
}
/// Name and timestamps as a plain tuple. `SecretInfo` carries no value, so
/// there is none to filter out here. Timestamps stay integers rather than
/// going through a string, so the caller can compare two without parsing.
#[pyo3(signature = (name, namespace_path=None))]
pub fn describe_secret(
self_: PyRef<'_, Self>,
name: String,
namespace_path: Option<Vec<String>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
let namespace_path = namespace_path.unwrap_or_default();
future_into_py(self_.py(), async move {
let info = inner
.describe_secret(name, &namespace_path)
.await
.infer_error()?;
Ok((info.name, info.created_at_millis, info.updated_at_millis))
})
}
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
@@ -640,42 +858,16 @@ impl Connection {
})
}
pub fn get_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let description = inner.get_job(&job_id).await.infer_error()?;
Ok(description.map(crate::job::JobDescription::from))
})
}
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
inner.cancel_job(&job_id).await.infer_error()
})
}
#[pyo3(signature = (job_id=None))]
pub fn job_history(
self_: PyRef<'_, Self>,
job_id: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let batches = inner.job_history(job_id.as_deref()).await.infer_error()?;
Python::attach(|py| {
let list = PyList::empty(py);
for batch in batches {
list.append(batch.to_pyarrow(py)?)?;
}
Ok(list.unbind())
})
})
}
}
#[pyfunction]
#[pyo3(signature = (uri, api_key=None, region=None, host_override=None, read_consistency_interval=None, client_config=None, storage_options=None, session=None, manifest_enabled=false, namespace_client_properties=None, oauth_config=None))]
#[pyo3(signature = (uri, api_key=None, region=None, host_override=None, sql_host_override=None, read_consistency_interval=None, client_config=None, storage_options=None, session=None, manifest_enabled=false, namespace_client_properties=None, oauth_config=None))]
#[allow(clippy::too_many_arguments)]
pub fn connect(
py: Python<'_>,
@@ -683,6 +875,7 @@ pub fn connect(
api_key: Option<String>,
region: Option<String>,
host_override: Option<String>,
sql_host_override: Option<String>,
read_consistency_interval: Option<f64>,
client_config: Option<PyClientConfig>,
storage_options: Option<HashMap<String, String>>,
@@ -702,6 +895,12 @@ pub fn connect(
if let Some(host_override) = host_override {
builder = builder.host_override(&host_override);
}
#[cfg(feature = "remote")]
if let Some(sql_host_override) = sql_host_override {
builder = builder.sql_host_override(&sql_host_override);
}
#[cfg(not(feature = "remote"))]
let _ = sql_host_override;
if let Some(read_consistency_interval) = read_consistency_interval {
let read_consistency_interval = Duration::from_secs_f64(read_consistency_interval);
builder = builder.read_consistency_interval(read_consistency_interval);
python/src/error.rs
+9
View File
@@ -29,7 +29,10 @@ impl<T> PythonErrorExt<T> for std::result::Result<T, LanceError> {
LanceError::InvalidInput { .. }
| LanceError::InvalidTableName { .. }
| LanceError::TableNotFound { .. }
| LanceError::NotAMaterializedView { .. }
| LanceError::Schema { .. }
| LanceError::DatabaseNotFound { .. }
| LanceError::DatabaseAlreadyExists { .. }
| LanceError::TableAlreadyExists { .. } => self.value_error(),
LanceError::CreateDir { .. } => self.os_error(),
LanceError::ObjectStore { .. } => Err(PyIOError::new_err(err.to_string())),
@@ -114,6 +117,12 @@ impl<T> PythonErrorExt<T> for std::result::Result<T, LanceError> {
.getattr(intern!(py, "JobCancelledError"))?;
Err(PyErr::from_value(cls.call1((err.to_string(),))?))
}),
LanceError::JobNotFound { .. } => Python::attach(|py| {
let cls = py
.import(intern!(py, "lancedb.exceptions"))?
.getattr(intern!(py, "JobNotFoundError"))?;
Err(PyErr::from_value(cls.call1((err.to_string(),))?))
}),
_ => self.runtime_error(),
},
}
python/src/job.rs
+129 -12
View File
@@ -4,11 +4,50 @@
use std::sync::Arc;
use crate::runtime::future_into_py;
use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods};
use arrow::{
datatypes::Schema,
pyarrow::{IntoPyArrow, Table as PyArrowTable},
};
use lancedb::job::JobEventsRequest;
use pyo3::{
Bound, PyAny, PyRef, PyResult, Python,
exceptions::PyValueError,
pyclass, pymethods,
types::{PyAnyMethods, PyDict, PyDictMethods},
};
use serde::Serialize;
use crate::error::PythonErrorExt;
const REPR_INDENT: &str = " ";
/// Parse a stored JSON payload into Python data. The bindings carry these as
/// strings because that is what crosses the boundary cheaply; the public
/// Python surface is the parsed form.
fn parse_json_payload<'py>(
py: Python<'py>,
raw: Option<&str>,
) -> PyResult<Option<Bound<'py, PyAny>>> {
match raw {
None => Ok(None),
Some(raw) => Ok(Some(py.import("json")?.call_method1("loads", (raw,))?)),
}
}
/// A payload rendered as indented JSON, aligned under the field that holds it.
fn pretty_json_payload(py: Python<'_>, raw: Option<&str>) -> PyResult<Option<String>> {
let Some(parsed) = parse_json_payload(py, raw)? else {
return Ok(None);
};
let kwargs = PyDict::new(py);
kwargs.set_item("indent", 4)?;
let rendered: String = py
.import("json")?
.call_method("dumps", (parsed,), Some(&kwargs))?
.extract()?;
Ok(Some(rendered.replace('\n', &format!("\n{REPR_INDENT}"))))
}
#[pyclass]
pub struct Job {
inner: Arc<lancedb::Job<std::result::Result<Option<String>, String>>>,
@@ -67,10 +106,52 @@ impl Job {
Ok(())
})
}
pub fn refresh(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner.refresh().await.infer_error()?;
Ok(())
})
}
/// The last observed lifecycle state, without contacting the backend.
#[getter]
pub fn _state(&self) -> Option<String> {
self.inner.state()
}
/// The last observed server-side record. `None` for an in-process job.
#[getter]
pub fn _description(&self) -> Option<JobDescription> {
self.inner.description().map(JobDescription::from)
}
#[pyo3(signature = (*, limit=None, filter=None))]
pub fn events(
self_: PyRef<'_, Self>,
limit: Option<u32>,
filter: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
let request = JobEventsRequest { limit, filter };
future_into_py(self_.py(), async move {
let batches = inner.events(request).await.infer_error()?;
Python::attach(|py| {
let schema = batches
.first()
.map(|batch| batch.schema())
.unwrap_or_else(|| Arc::new(Schema::empty()));
let table = PyArrowTable::try_new(batches, schema)
.map_err(|err| PyValueError::new_err(err.to_string()))?;
table.into_pyarrow(py).map(|table| table.unbind())
})
})
}
}
/// A row from `Connection.list_jobs`: one server-side job.
#[pyclass(get_all, skip_from_py_object)]
#[pyclass(module = "lancedb._lancedb", get_all, skip_from_py_object)]
#[derive(Clone)]
pub struct JobInfo {
job_id: String,
@@ -103,7 +184,7 @@ impl From<lancedb::database::JobInfo> for JobInfo {
}
/// The server's account of why a job failed.
#[pyclass(get_all, skip_from_py_object)]
#[pyclass(module = "lancedb._lancedb", get_all, skip_from_py_object)]
#[derive(Clone)]
pub struct JobFailureInfo {
phase: Option<String>,
@@ -121,25 +202,57 @@ impl JobFailureInfo {
}
}
/// A described job from `Connection.get_job`.
#[pyclass(get_all, skip_from_py_object)]
/// The server-side record behind a `Job` handle.
#[pyclass(module = "lancedb._lancedb", get_all, skip_from_py_object)]
#[derive(Clone)]
pub struct JobDescription {
job_id: String,
job_type: String,
state: String,
creation_ms: i64,
spec_json: Option<String>,
/// Internal: the wire form behind the `spec` property.
_spec_json: Option<String>,
/// Internal: the wire form behind the `result` property.
_result_json: Option<String>,
failure: Option<JobFailureInfo>,
}
#[pymethods]
impl JobDescription {
fn __repr__(&self) -> String {
format!(
"JobDescription(job_id={:?}, job_type={:?}, state={:?}, creation_ms={})",
self.job_id, self.job_type, self.state, self.creation_ms
)
/// The job-type-specific specification it was submitted with.
#[getter]
fn spec<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyAny>>> {
parse_json_payload(py, self._spec_json.as_deref())
}
/// The job-type-specific terminal result. `None` until the job succeeds.
#[getter]
fn result<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyAny>>> {
parse_json_payload(py, self._result_json.as_deref())
}
fn __repr__(&self, py: Python<'_>) -> PyResult<String> {
let mut fields = vec![
format!("job_id={:?}", self.job_id),
format!("job_type={:?}", self.job_type),
format!("state={:?}", self.state),
format!("creation_ms={}", self.creation_ms),
];
// Lay the payloads out as indented JSON, the same way the `Job` repr
// does, so the two agree on how the same data looks.
for (name, payload) in [("spec", &self._spec_json), ("result", &self._result_json)] {
if let Some(rendered) = pretty_json_payload(py, payload.as_deref())? {
fields.push(format!("{name}={rendered}"));
}
}
if let Some(failure) = &self.failure {
fields.push(format!("failure={}", failure.__repr__()));
}
let body = fields
.iter()
.map(|field| format!("\n{REPR_INDENT}{field},"))
.collect::<String>();
Ok(format!("JobDescription({body}\n)"))
}
}
@@ -150,7 +263,11 @@ impl From<lancedb::database::JobDescription> for JobDescription {
job_type: description.job_type,
state: description.state,
creation_ms: description.creation_ms,
spec_json: (!description.spec.is_null()).then(|| description.spec.to_string()),
_spec_json: (!description.spec.is_null()).then(|| description.spec.to_string()),
_result_json: description
.result
.filter(|result| !result.is_null())
.map(|result| result.to_string()),
failure: description.failure.map(|failure| JobFailureInfo {
phase: failure.phase,
message: failure.message,
python/src/lib.rs
+32 -3
View File
@@ -8,8 +8,8 @@ use expr::{PyExpr, expr_col, expr_func, expr_lit};
use index::IndexConfig;
use permutation::{PyAsyncPermutationBuilder, PyPermutationReader};
use pyo3::{
Bound, PyResult, Python, pymodule,
types::{PyModule, PyModuleMethods},
Bound, PyResult, Python, pyfunction, pymodule,
types::{PyAnyMethods, PyModule, PyModuleMethods},
wrap_pyfunction,
};
use query::{FTSQuery, HybridQuery, Query, VectorQuery};
@@ -21,6 +21,7 @@ use table::{
};
pub mod arrow;
pub mod catalog;
pub mod connection;
pub mod error;
pub mod expr;
@@ -34,22 +35,46 @@ pub mod permutation;
pub mod query;
pub mod runtime;
pub mod session;
pub mod sql;
pub mod table;
pub mod util;
/// Shut down the shared Tokio runtime (see `runtime::shutdown`).
///
/// Registered below as a Python `atexit` callback rather than called
/// directly: `atexit` runs while the interpreter is still fully valid,
/// which is the coordinated, bounded exit the runtime otherwise never gets.
///
/// Runs the actual wait with the GIL released (`Python::detach`): shutdown
/// blocks the calling thread waiting on the runtime's own worker threads,
/// and if any in-flight task needs the GIL to finish (e.g. one that calls
/// back into Python), holding it here while waiting on that same task would
/// deadlock rather than time out.
#[pyfunction]
fn shutdown_runtime(py: Python<'_>) {
py.detach(|| runtime::shutdown(std::time::Duration::from_secs(5)));
}
#[pymodule]
pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
let env = Env::new()
.filter_or("LANCEDB_LOG", "warn")
.write_style("LANCEDB_LOG_STYLE");
env_logger::init_from_env(env);
m.add_class::<Connection>()?;
m.add_class::<catalog::Catalog>()?;
m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?;
m.add_class::<Session>()?;
m.add_class::<Table>()?;
m.add_class::<crate::oauth::PyOAuthSession>()?;
m.add_class::<crate::oauth::PySessionStatus>()?;
m.add_class::<crate::oauth::PySessionLogout>()?;
m.add_class::<crate::job::Job>()?;
m.add_class::<crate::job::JobInfo>()?;
m.add_class::<crate::job::JobDescription>()?;
m.add_class::<crate::job::JobFailureInfo>()?;
m.add_class::<crate::sql::Query>()?;
m.add_class::<crate::sql::QueryDescription>()?;
m.add_class::<PyBlobFile>()?;
m.add_class::<IndexConfig>()?;
m.add_class::<Query>()?;
@@ -92,5 +117,9 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(expr_lit, m)?)?;
m.add_function(wrap_pyfunction!(expr_func, m)?)?;
m.add("__version__", env!("CARGO_PKG_VERSION"))?;
// Give the shared runtime a coordinated, bounded shutdown at normal
// process exit -- see `shutdown_runtime` and `runtime::shutdown` for why.
py.import("atexit")?
.call_method1("register", (wrap_pyfunction!(shutdown_runtime, m)?,))?;
Ok(())
}
python/src/oauth.rs
+319 -6
View File
@@ -1,10 +1,33 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use pyo3::FromPyObject;
use std::path::PathBuf;
use std::sync::Arc;
use pyo3::{FromPyObject, PyResult, Python, pyclass, pymethods};
use crate::error::PythonErrorExt;
use crate::runtime::future_into_py;
use lancedb::error::Error;
use lancedb::remote::oauth::{OAuthConfig, OAuthFlow};
use lancedb::remote::oauth::{AuthorizationCodeOptions, ClientAuthMethod, OAuthConfig, OAuthFlow};
use lancedb::remote::{OAuthSession, SessionLogout, SessionStatus, TokenCacheOptions};
/// Python-side persistent token cache options, extracted via FromPyObject.
/// Maps to `lancedb.remote.oauth.TokenCacheOptions` Python dataclass.
#[derive(FromPyObject, Default)]
pub struct PyTokenCacheOptions {
pub cache_dir: Option<String>,
pub lock_timeout_secs: Option<u64>,
}
impl From<PyTokenCacheOptions> for TokenCacheOptions {
fn from(py: PyTokenCacheOptions) -> Self {
Self {
cache_dir: py.cache_dir.map(PathBuf::from),
lock_timeout_secs: py.lock_timeout_secs,
}
}
}
/// Python-side OAuth configuration, extracted via FromPyObject.
/// Maps to `lancedb.remote.oauth.OAuthConfig` Python dataclass.
@@ -13,10 +36,19 @@ pub struct PyOAuthConfig {
pub issuer_url: String,
pub client_id: String,
pub scopes: Vec<String>,
/// Optional resource indicator for authorization and token requests.
pub resource: Option<String>,
/// Optional provider-specific audience for authorization and token requests.
pub audience: Option<String>,
pub flow: String,
pub client_secret: Option<String>,
pub client_auth_method: Option<String>,
pub redirect_uri: Option<String>,
pub callback_port: Option<u16>,
pub use_pkce: bool,
pub managed_identity_client_id: Option<String>,
pub refresh_buffer_secs: Option<u64>,
pub token_cache: Option<PyTokenCacheOptions>,
}
impl TryFrom<PyOAuthConfig> for OAuthConfig {
@@ -25,6 +57,17 @@ impl TryFrom<PyOAuthConfig> for OAuthConfig {
fn try_from(py: PyOAuthConfig) -> Result<Self, Self::Error> {
let flow = match py.flow.as_str() {
"client_credentials" => OAuthFlow::ClientCredentials,
"authorization_code" => {
let mut options = AuthorizationCodeOptions::new().use_pkce(py.use_pkce);
if let Some(redirect_uri) = py.redirect_uri {
options = options.redirect_uri(redirect_uri);
}
if let Some(callback_port) = py.callback_port {
options = options.callback_port(callback_port);
}
OAuthFlow::AuthorizationCode(options)
}
"device_code" => OAuthFlow::DeviceCode,
"azure_managed_identity" => OAuthFlow::AzureManagedIdentity {
client_id: py.managed_identity_client_id,
},
@@ -35,13 +78,182 @@ impl TryFrom<PyOAuthConfig> for OAuthConfig {
}
};
let client_auth_method = match py.client_auth_method.as_deref() {
Some("none") => Some(ClientAuthMethod::None),
Some("client_secret_basic") => Some(ClientAuthMethod::ClientSecretBasic),
Some("client_secret_post") => Some(ClientAuthMethod::ClientSecretPost),
None => None,
Some(other) => {
return Err(Error::InvalidInput {
message: format!("Unknown OAuth client auth method: {other}"),
});
}
};
Ok(Self {
issuer_url: py.issuer_url,
client_id: py.client_id,
client_secret: py.client_secret,
client_auth_method,
scopes: py.scopes,
resource: py.resource,
audience: py.audience,
flow,
refresh_buffer_secs: py.refresh_buffer_secs,
token_cache: py.token_cache.map(TokenCacheOptions::from),
})
}
}
/// Wrapper around [`lancedb::remote::SessionStatus`] exposing safe metadata.
#[pyclass(name = "SessionStatus", skip_from_py_object)]
#[derive(Clone)]
pub struct PySessionStatus {
inner: SessionStatus,
}
#[pymethods]
impl PySessionStatus {
/// Whether a cached session exists that can obtain tokens without
/// interactive authentication.
#[getter]
pub fn refreshable(&self) -> bool {
self.inner.refreshable
}
/// Canonical issuer URL of the cached session.
#[getter]
pub fn issuer_url(&self) -> String {
self.inner.issuer_url.clone()
}
/// Client ID of the cached session.
#[getter]
pub fn client_id(&self) -> String {
self.inner.client_id.clone()
}
/// Canonical (sorted, de-duplicated) scopes of the cached session.
#[getter]
pub fn scopes(&self) -> Vec<String> {
self.inner.scopes.clone()
}
/// Resource indicator used to obtain the cached session.
#[getter]
pub fn resource(&self) -> Option<String> {
self.inner.resource.clone()
}
/// Provider-specific audience used to obtain the cached session.
#[getter]
pub fn audience(&self) -> Option<String> {
self.inner.audience.clone()
}
/// Flow that produced the cached session.
#[getter]
pub fn flow(&self) -> String {
self.inner.flow.clone()
}
/// When the cached session was obtained, as Unix seconds.
#[getter]
pub fn obtained_at(&self) -> Option<u64> {
self.inner.obtained_at
}
pub fn __repr__(&self) -> String {
format!(
"SessionStatus(refreshable={}, issuer_url='{}', client_id='{}', flow='{}')",
self.inner.refreshable, self.inner.issuer_url, self.inner.client_id, self.inner.flow
)
}
}
impl From<SessionStatus> for PySessionStatus {
fn from(inner: SessionStatus) -> Self {
Self { inner }
}
}
/// Wrapper around [`lancedb::remote::SessionLogout`].
#[pyclass(name = "SessionLogout", skip_from_py_object)]
#[derive(Clone)]
pub struct PySessionLogout {
inner: SessionLogout,
}
#[pymethods]
impl PySessionLogout {
/// Whether a cached credential was removed.
#[getter]
pub fn removed(&self) -> bool {
self.inner.removed
}
pub fn __repr__(&self) -> String {
format!("SessionLogout(removed={})", self.inner.removed)
}
}
impl From<SessionLogout> for PySessionLogout {
fn from(inner: SessionLogout) -> Self {
Self { inner }
}
}
/// Wrapper around [`lancedb::remote::OAuthSession`].
#[pyclass(name = "OAuthSession", skip_from_py_object)]
#[derive(Clone)]
pub struct PyOAuthSession {
inner: Arc<OAuthSession>,
}
#[pymethods]
impl PyOAuthSession {
/// Create a session manager for the given OAuth configuration.
///
/// The configuration must set ``token_cache`` options and use a flow that
/// supports persistent sessions (authorization code or device code).
#[new]
pub fn new(config: PyOAuthConfig) -> PyResult<Self> {
let config: OAuthConfig = config.try_into().infer_error()?;
let inner = OAuthSession::new(config).infer_error()?;
Ok(Self {
inner: Arc::new(inner),
})
}
/// Eagerly run the configured authentication flow and store the session.
pub fn login<'py>(&self, py: Python<'py>) -> PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
let inner = Arc::clone(&self.inner);
future_into_py(py, async move {
inner.login().await.map(PySessionStatus::from).infer_error()
})
}
/// Report whether a matching cached session exists, with safe metadata.
pub fn status<'py>(&self, py: Python<'py>) -> PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
let inner = Arc::clone(&self.inner);
future_into_py(py, async move {
inner
.status()
.await
.map(PySessionStatus::from)
.infer_error()
})
}
/// Remove the matching local cached credential.
pub fn logout<'py>(&self, py: Python<'py>) -> PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
let inner = Arc::clone(&self.inner);
future_into_py(py, async move {
inner
.logout()
.await
.map(PySessionLogout::from)
.infer_error()
})
}
}
@@ -50,16 +262,30 @@ impl TryFrom<PyOAuthConfig> for OAuthConfig {
mod tests {
use super::*;
#[test]
fn test_unknown_oauth_flow_returns_invalid_input() {
let config = PyOAuthConfig {
fn base_config() -> PyOAuthConfig {
PyOAuthConfig {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
scopes: vec!["scope".to_string()],
flow: "typo".to_string(),
flow: "device_code".to_string(),
client_secret: None,
client_auth_method: None,
redirect_uri: None,
callback_port: None,
use_pkce: true,
managed_identity_client_id: None,
refresh_buffer_secs: None,
resource: None,
audience: None,
token_cache: None,
}
}
#[test]
fn test_unknown_oauth_flow_returns_invalid_input() {
let config = PyOAuthConfig {
flow: "typo".to_string(),
..base_config()
};
let err = OAuthConfig::try_from(config).unwrap_err();
@@ -69,4 +295,91 @@ mod tests {
if message == "Unknown OAuth flow type: typo"
));
}
#[test]
fn test_authorization_code_conversion_preserves_options() {
let config = PyOAuthConfig {
flow: "authorization_code".to_string(),
client_secret: Some("secret".to_string()),
redirect_uri: Some("http://127.0.0.1:9000/callback".to_string()),
callback_port: Some(9000),
use_pkce: false,
resource: Some("urn:resource".into()),
audience: Some("audience".into()),
..base_config()
};
let converted = OAuthConfig::try_from(config).unwrap();
let OAuthFlow::AuthorizationCode(options) = converted.flow else {
panic!("expected authorization code flow");
};
assert_eq!(
options.redirect_uri.as_deref(),
Some("http://127.0.0.1:9000/callback")
);
assert_eq!(options.callback_port, Some(9000));
assert!(!options.use_pkce);
assert_eq!(converted.resource.as_deref(), Some("urn:resource"));
assert_eq!(converted.audience.as_deref(), Some("audience"));
}
#[test]
fn test_device_code_conversion() {
let config = base_config();
let converted = OAuthConfig::try_from(config).unwrap();
assert!(matches!(converted.flow, OAuthFlow::DeviceCode));
}
#[test]
fn test_client_auth_method_conversion() {
for (value, expected) in [
("none", ClientAuthMethod::None),
("client_secret_basic", ClientAuthMethod::ClientSecretBasic),
("client_secret_post", ClientAuthMethod::ClientSecretPost),
] {
let config = PyOAuthConfig {
client_auth_method: Some(value.to_string()),
..base_config()
};
let converted = OAuthConfig::try_from(config).unwrap();
assert_eq!(converted.client_auth_method, Some(expected));
}
}
#[test]
fn test_unknown_client_auth_method_returns_invalid_input() {
let config = PyOAuthConfig {
client_auth_method: Some("typo".to_string()),
..base_config()
};
let err = OAuthConfig::try_from(config).unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message == "Unknown OAuth client auth method: typo"
));
}
#[test]
fn test_token_cache_conversion() {
let config = PyOAuthConfig {
resource: None,
audience: None,
token_cache: Some(PyTokenCacheOptions {
cache_dir: Some("/tmp/oauth-cache".to_string()),
lock_timeout_secs: Some(5),
}),
..base_config()
};
let converted = OAuthConfig::try_from(config).unwrap();
let cache = converted.token_cache.expect("token cache options");
assert_eq!(
cache.cache_dir.as_deref(),
Some(std::path::Path::new("/tmp/oauth-cache"))
);
assert_eq!(cache.lock_timeout_secs, Some(5));
}
}
python/src/query.rs
+4 -1
View File
@@ -323,6 +323,7 @@ impl<'py> IntoPyObject<'py> for PyQueryVectors {
pub struct PyQueryRequest {
pub limit: Option<usize>,
pub offset: Option<usize>,
pub take_offsets: Option<Vec<u64>>,
pub filter: Option<PyQueryFilter>,
pub full_text_search: Option<PyLanceDB<FtsQuery>>,
pub select: PySelect,
@@ -333,7 +334,7 @@ pub struct PyQueryRequest {
pub column: Option<String>,
pub query_vector: Option<PyQueryVectors>,
pub minimum_nprobes: Option<usize>,
// None means user did not set it and default shoud be used (currenty 20)
// None means user did not set it and default should be used (currently 20)
// Some(0) means user set it to None and there is no limit
pub maximum_nprobes: Option<usize>,
pub lower_bound: Option<f32>,
@@ -353,6 +354,7 @@ impl From<AnyQuery> for PyQueryRequest {
AnyQuery::Query(query_request) => Self {
limit: query_request.limit,
offset: query_request.offset,
take_offsets: query_request.take_offsets,
filter: query_request.filter.map(PyQueryFilter),
full_text_search: query_request
.full_text_search
@@ -381,6 +383,7 @@ impl From<AnyQuery> for PyQueryRequest {
AnyQuery::VectorQuery(vector_query) => Self {
limit: vector_query.base.limit,
offset: vector_query.base.offset,
take_offsets: vector_query.base.take_offsets,
filter: vector_query.base.filter.map(PyQueryFilter),
full_text_search: None,
select_source_columns: PySelect::source_columns(&vector_query.base.select),
python/src/runtime.rs
+460 -30
View File
@@ -4,20 +4,84 @@
//! Fork-safe wrapper around tokio + pyo3-async-runtimes.
//!
//! `pyo3_async_runtimes::tokio` keeps its multi-threaded runtime in a
//! `OnceLock` that can never be replaced. Tokio's worker threads do not
//! `OnceLock` that can never be replaced. Tokio's worker threads do not
//! survive `fork()`, so once a child inherits a "frozen" runtime, every
//! `future_into_py` call hangs forever.
//! `future_into_py` call hangs forever. Normal (non-fork) process exit has
//! its own gap: nothing tells the runtime to shut down, so its worker
//! threads keep running, uncoordinated with the interpreter, right up until
//! the process ends. If one of them is mid-task exactly as `Py_Finalize`
//! starts tearing down interpreter state, it can panic on state that's
//! already gone -- and since that happens on a background thread with no
//! PyO3-wrapped call frame to catch it, Rust aborts the whole process
//! rather than failing that one call. [`shutdown`], registered as a Python
//! `atexit` callback, closes that gap by giving the runtime a coordinated,
//! bounded exit while the interpreter is still fully valid.
//!
//! We sidestep the global by routing every future through our own
//! [`LanceRuntime`] (a [`pyo3_async_runtimes::generic::Runtime`] impl) backed
//! by an [`AtomicPtr`] to a tokio runtime that we own. A `pthread_atfork`
//! child handler nulls the pointer; the next `spawn` rebuilds the runtime in
//! the child. This mirrors the pattern used in the Lance Python bindings.
//! Getting both of these right at once took a few tries; the design here
//! rests on three separate mechanisms, each solving one problem the others
//! cannot:
//!
//! **`OUTSTANDING`, not `Arc::strong_count`, decides when the runtime is
//! idle.** Early versions tried to infer "is anything still using this
//! runtime" from how many `Arc<Runtime>` clones existed. That signal is
//! wrong in both directions: a clone taken only for the instant a task is
//! *submitted* says nothing about whether that task has actually finished
//! running (`Runtime::shutdown_timeout` gives spawned, non-blocking tasks
//! no grace period at all -- a task "keeps running until it yields, then is
//! dropped" -- so a reclaim landing right after submission would silently
//! abandon it before it ever got to run); and a clone held for a task's
//! *whole* lifetime can end up making that task the final owner of the
//! `Runtime`, so completing it drops the `Runtime` from inside one of its
//! own worker threads, which tokio itself forbids ("cannot drop a runtime
//! in a context where blocking is not allowed") and panics. `OUTSTANDING`
//! is an explicit courtesy counter instead: every top-level `spawn`,
//! `spawn_blocking`, or `block_on` call increments it before it starts and
//! decrements it (via [`OutstandingGuard`]) when it is truly done, entirely
//! decoupled from how many `Arc` clones exist at any instant. `shutdown`
//! waits for it to reach zero before ever touching the runtime, which also
//! closes a narrower race: because the counter is incremented *before*
//! `get_runtime()` is even called, an install already in progress when
//! `shutdown` runs is never invisible to it the way an empty slot would be.
//! If the counter never reaches zero within the bound, `shutdown` stops
//! waiting and forces the retirement attempt anyway -- silently returning
//! with the runtime and its workers still fully alive would just recreate
//! the exact race this function exists to close, for any call slower than
//! the grace period.
//!
//! **Tasks never hold an `Arc<Runtime>`.** Because `OUTSTANDING` (not
//! reference counting) is what `shutdown` waits on, a top-level task only
//! needs to carry an `OutstandingGuard` -- a token whose `Drop` is a plain
//! atomic decrement -- not a clone of the runtime itself. That is what
//! makes it impossible for a task's completion to become the final,
//! worker-thread-side drop of the `Runtime`: nothing a task holds ever
//! *is* the `Runtime`.
//!
//! **`atfork_child` touches nothing but a plain counter.** `future_into_py`
//! spawns a task that, once running, spawns a second one to do the real
//! work and awaits its `JoinHandle`; if a nested call re-resolved "the
//! current runtime" independently, a reclaim landing between the two calls
//! could bind them to different instances. `spawn`/`spawn_blocking` close
//! that with `Handle::try_current`: a call already running on one of our
//! worker threads stays pinned to that instance, so only the first,
//! outermost call of a chain ever consults [`get_runtime`]. That leaves
//! fork as the other place identity can change, and it has to be handled
//! without ever calling into `ArcSwapOption` from the child handler itself
//! -- `swap`/`compare_and_swap` reconcile reader "debts" internally (via
//! thread-local state, and potentially an allocation), none of which is
//! safe to run in a forked child that may have inherited another thread's
//! lock mid-acquisition. `atfork_child` therefore does nothing but bump a
//! bare `GENERATION` counter; [`get_runtime`] compares the generation its
//! installed runtime was built in against the live counter on every call,
//! from ordinary (non-signal) context, and treats a mismatch as "stale,
//! rebuild" -- exactly the check `atfork_child` used to perform directly.
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::{Duration, Instant};
use arc_swap::ArcSwapOption;
use pyo3::{Bound, PyAny, PyResult, Python, conversion::IntoPyObject};
use pyo3_async_runtimes::{
TaskLocals,
@@ -25,8 +89,31 @@ use pyo3_async_runtimes::{
};
use tokio::{runtime, task};
static RUNTIME: AtomicPtr<runtime::Runtime> = AtomicPtr::new(std::ptr::null_mut());
static RUNTIME_INSTALLING: AtomicBool = AtomicBool::new(false);
/// A runtime tagged with the fork generation it was built in, so a stale
/// (post-fork, dead-worker-threads) instance can be told apart from a live
/// one without `atfork_child` ever having to touch it directly.
struct Tagged {
runtime: runtime::Runtime,
generation: u64,
}
impl std::ops::Deref for Tagged {
type Target = runtime::Runtime;
fn deref(&self) -> &runtime::Runtime {
&self.runtime
}
}
static RUNTIME: ArcSwapOption<Tagged> = ArcSwapOption::const_empty();
/// Bumped only by `atfork_child`, and only ever read elsewhere. This is the
/// entire fork-safety mechanism: no lock, no allocation, no thread-local
/// access -- just one atomic add, which is all a `pthread_atfork` child
/// handler is ever safe to do.
static GENERATION: AtomicU64 = AtomicU64::new(0);
/// Count of top-level `spawn`/`spawn_blocking`/`block_on` calls that have
/// started but not yet finished. See the module docs for why this, and not
/// `Arc::strong_count`, is what `shutdown` waits on.
static OUTSTANDING: AtomicU64 = AtomicU64::new(0);
static ATFORK_INSTALLED: AtomicBool = AtomicBool::new(false);
fn create_runtime() -> runtime::Runtime {
@@ -37,23 +124,67 @@ fn create_runtime() -> runtime::Runtime {
.expect("Failed to build tokio runtime")
}
fn get_runtime() -> &'static runtime::Runtime {
/// Get a live, owned handle to the shared runtime, rebuilding it if the
/// installed one predates the most recent `fork()`.
fn get_runtime() -> Arc<Tagged> {
let current_gen = GENERATION.load(Ordering::SeqCst);
loop {
let ptr = RUNTIME.load(Ordering::SeqCst);
if !ptr.is_null() {
return unsafe { &*ptr };
let existing = RUNTIME.load_full();
if let Some(existing) = &existing
&& existing.generation == current_gen
{
return Arc::clone(existing);
}
if !RUNTIME_INSTALLING.fetch_or(true, Ordering::SeqCst) {
break;
if !ATFORK_INSTALLED.fetch_or(true, Ordering::SeqCst) {
install_atfork();
}
std::thread::yield_now();
// Built optimistically, outside any lock: on the rare race where two
// threads both find the slot empty (or stale), one candidate wins
// the compare-and-swap below and the other is simply dropped here,
// tearing down its own (never shared, never used) worker pool the
// ordinary way.
let candidate = Arc::new(Tagged {
runtime: create_runtime(),
generation: current_gen,
});
let previous = RUNTIME.compare_and_swap(&existing, Some(Arc::clone(&candidate)));
let won = match (&*previous, &existing) {
(None, None) => true,
(Some(prev), Some(exist)) => Arc::ptr_eq(prev, exist),
_ => false,
};
if won {
if let Some(stale) = existing {
// A prior generation's runtime: its worker threads are dead
// in this process (they do not survive fork), so dropping it
// normally would try to join them and hang. Leak it instead.
std::mem::forget(stale);
}
return candidate;
}
// Someone else's candidate (or a concurrent shutdown) won; go around
// and reload.
}
if !ATFORK_INSTALLED.fetch_or(true, Ordering::SeqCst) {
install_atfork();
}
/// RAII token tracked by [`OUTSTANDING`]. Held for the duration of a
/// top-level `block_on` call, or moved into a top-level `spawn`/
/// `spawn_blocking` task so it decrements only once that task's *entire*
/// body -- including anything it nested-spawns and awaits -- has run to
/// completion or been dropped without completing.
struct OutstandingGuard;
impl OutstandingGuard {
fn new() -> Self {
OUTSTANDING.fetch_add(1, Ordering::SeqCst);
Self
}
}
impl Drop for OutstandingGuard {
fn drop(&mut self) {
OUTSTANDING.fetch_sub(1, Ordering::SeqCst);
}
let new_ptr = Box::into_raw(Box::new(create_runtime()));
RUNTIME.store(new_ptr, Ordering::SeqCst);
unsafe { &*new_ptr }
}
/// Block the current thread on a future using the shared runtime.
@@ -62,16 +193,78 @@ fn get_runtime() -> &'static runtime::Runtime {
/// building a namespace client). Must not be called from within the runtime's
/// own worker threads.
pub fn block_on<F: std::future::Future>(fut: F) -> F::Output {
let _guard = OutstandingGuard::new();
get_runtime().block_on(fut)
}
/// Runs in async-signal context after `fork()` in the child. We can only
/// touch atomics here; we deliberately leak the previous runtime because
/// dropping a tokio `Runtime` would try to join its (now-dead) worker
/// threads and hang.
/// Gracefully quiesce the shared runtime, meant to run at normal process exit.
///
/// Waits (bounded by `timeout`) for [`OUTSTANDING`] to reach zero -- i.e.
/// for every top-level call already under way to actually finish, not just
/// for `Arc::strong_count` to look low -- before ever touching the runtime.
/// If the bound elapses first, it stops waiting and attempts retirement
/// anyway: leaving the runtime and its worker threads untouched would just
/// recreate the exact race this function exists to close, for any call
/// slower than the grace period.
///
/// Retirement itself removes the runtime from the slot and calls
/// `shutdown_timeout` rather than a bare `drop`: dropping a tokio `Runtime`
/// waits (in the worst case indefinitely) for its worker threads to join,
/// whereas `shutdown_timeout` gives real in-flight work -- a connection
/// pool's keep-alive, a graceful close -- a bounded chance to finish first,
/// then forcibly ends whatever has not. Reclaiming can only proceed once
/// `Arc::try_unwrap` proves no other reference remains; if some transient
/// `get_runtime()` caller is, at that exact instant, still between loading
/// the slot and finishing its own call, this abandons the runtime instead
/// of forcing the issue -- the same trade `atfork_child` already makes.
///
/// Neither of the two ways this can fail to cleanly retire the runtime --
/// the wait timing out, or `try_unwrap` losing that race -- has any other
/// signal to report through (`shutdown_timeout` itself returns nothing),
/// so both log a warning instead of failing silently.
pub fn shutdown(timeout: Duration) {
let deadline = Instant::now() + timeout;
loop {
let outstanding = OUTSTANDING.load(Ordering::SeqCst);
if outstanding == 0 {
break;
}
if Instant::now() >= deadline {
log::warn!(
"lancedb: runtime shutdown timed out with {outstanding} call(s) still in flight; forcing shutdown anyway, some in-flight work may be abandoned"
);
break;
}
std::thread::sleep(Duration::from_millis(1));
}
let Some(current) = RUNTIME.load_full() else {
return;
};
RUNTIME.compare_and_swap(&Some(Arc::clone(&current)), None);
match Arc::try_unwrap(current) {
Ok(tagged) => {
tagged
.runtime
.shutdown_timeout(deadline.saturating_duration_since(Instant::now()));
}
Err(_) => {
// Some transient `get_runtime()` caller is, at this exact
// instant, still between loading the slot and finishing its own
// call: we have no owned handle to call `shutdown_timeout` on,
// and no way to force one (tokio has no shutdown API over a
// shared reference). Nothing more to do but say so.
log::warn!(
"lancedb: runtime shutdown could not obtain exclusive ownership; the shared runtime was left running"
);
}
}
}
/// Runs in async-signal context after `fork()` in the child. Touches
/// nothing but a plain atomic add -- see the module docs for why even
/// `ArcSwapOption::swap` is not safe to call here.
extern "C" fn atfork_child() {
RUNTIME.store(std::ptr::null_mut(), Ordering::SeqCst);
RUNTIME_INSTALLING.store(false, Ordering::SeqCst);
GENERATION.fetch_add(1, Ordering::SeqCst);
}
#[cfg(not(windows))]
@@ -102,11 +295,26 @@ impl Runtime for LanceRuntime {
type JoinError = LanceJoinError;
type JoinHandle = Pin<Box<dyn Future<Output = Result<(), Self::JoinError>> + Send>>;
/// `pyo3_async_runtimes::generic::future_into_py` spawns a task that,
/// once it starts running, spawns a second one to do the real work and
/// awaits its `JoinHandle`. `Handle::try_current` pins that nested call
/// to whatever runtime is already executing it, so only the first,
/// outermost call of a chain -- one running on a thread outside any
/// runtime -- ever consults [`get_runtime`] or [`OUTSTANDING`].
fn spawn<F>(fut: F) -> Self::JoinHandle
where
F: Future<Output = ()> + Send + 'static,
{
let handle = get_runtime().spawn(fut);
let handle = match tokio::runtime::Handle::try_current() {
Ok(handle) => handle.spawn(fut),
Err(_) => {
let guard = OutstandingGuard::new();
get_runtime().spawn(async move {
let _guard = guard;
fut.await;
})
}
};
Box::pin(async move { handle.await.map_err(LanceJoinError) })
}
@@ -114,7 +322,16 @@ impl Runtime for LanceRuntime {
where
F: FnOnce() + Send + 'static,
{
let handle = get_runtime().spawn_blocking(f);
let handle = match tokio::runtime::Handle::try_current() {
Ok(handle) => handle.spawn_blocking(f),
Err(_) => {
let guard = OutstandingGuard::new();
get_runtime().spawn_blocking(move || {
let _guard = guard;
f();
})
}
};
Box::pin(async move { handle.await.map_err(LanceJoinError) })
}
}
@@ -149,3 +366,216 @@ where
{
pyo3_async_runtimes::generic::future_into_py::<LanceRuntime, _, T>(py, fut)
}
#[cfg(test)]
mod tests {
use super::*;
// RUNTIME, GENERATION, OUTSTANDING, and ATFORK_INSTALLED are process-wide
// statics, and Rust's test harness runs tests in parallel by default,
// so separate test functions below would otherwise race each other
// through this shared state (this reproduced in CI: one test's
// in-flight task got reclaimed by a *different* test's concurrent
// `shutdown()` call, and another observed `OUTSTANDING` left non-zero
// by a still-running sibling). Every test takes this lock first so
// only one of them touches the shared runtime state at a time; a
// poisoned lock (a previous test's genuine failure) is still honored
// rather than cascading into every later test as an unrelated panic.
static TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
fn lock_runtime_state_for_test() -> std::sync::MutexGuard<'static, ()> {
TEST_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
// A task's own completion must never be the final drop of the shared
// `Runtime`: tasks carry only an `OutstandingGuard` (a plain counter
// token), never an `Arc<Runtime>`, specifically so this can't happen.
// Getting this wrong panics ("cannot drop a runtime in a context where
// blocking is not allowed") -- this reproduced unprompted, twice, in a
// single run of `test_nested_spawn_survives_concurrent_shutdown` under
// an earlier design that held the `Arc` for a task's whole lifetime.
#[test]
#[allow(unused_must_use)] // fire-and-forget spawn, same as future_into_py itself
fn test_top_level_task_survives_concurrent_shutdown_reclaim() {
use std::sync::mpsc;
let _lock = lock_runtime_state_for_test();
for _ in 0..50 {
let (tx, rx) = mpsc::channel::<()>();
LanceRuntime::spawn(async move {
for _ in 0..5 {
task::yield_now().await;
}
let _ = tx.send(());
});
shutdown(Duration::from_secs(2));
rx.recv_timeout(Duration::from_secs(5))
.expect("top-level task was abandoned by a concurrent shutdown reclaim");
}
}
#[test]
fn test_shutdown_stops_and_the_runtime_rebuilds_lazily_after() {
let _lock = lock_runtime_state_for_test();
// No runtime created yet in this process: shutdown must be a no-op,
// not a null-pointer dereference.
shutdown(Duration::from_secs(1));
// Force the runtime into existence, then shut it down. Bounded by
// the timeout, so a hang here means shutdown itself is broken, not
// that the test is slow.
assert_eq!(block_on(async { 1 + 1 }), 2);
shutdown(Duration::from_secs(5));
// A caller after shutdown -- e.g. a stray call racing with the
// atexit callback -- must get a fresh, working runtime rather than
// a dangling reference into the one just torn down.
assert_eq!(block_on(async { 2 + 2 }), 4);
// Shutting down twice in a row (e.g. atexit firing more than once)
// must not panic or double-free.
shutdown(Duration::from_secs(5));
shutdown(Duration::from_secs(5));
}
// Adapted from the concurrent reproducer that found the original bug:
// many threads hammering the runtime while shutdown races them, not
// just the sequential rebuild path above. Before every reader held its
// own `Arc`, this could dereference a runtime `shutdown` had already
// freed, or hang forever. Repeated, since a race is not guaranteed to
// show up on any single attempt.
#[test]
fn test_shutdown_is_safe_concurrently_with_live_callers() {
use std::sync::Barrier;
use std::sync::atomic::AtomicBool as StopFlag;
let _lock = lock_runtime_state_for_test();
for _ in 0..50 {
let barrier = Arc::new(Barrier::new(9));
let stop = Arc::new(StopFlag::new(false));
let workers: Vec<_> = (0..8)
.map(|_| {
let barrier = Arc::clone(&barrier);
let stop = Arc::clone(&stop);
std::thread::spawn(move || {
barrier.wait();
while !stop.load(Ordering::Relaxed) {
assert_eq!(block_on(async { 1 + 1 }), 2);
}
})
})
.collect();
barrier.wait();
shutdown(Duration::from_millis(50));
stop.store(true, Ordering::Relaxed);
for worker in workers {
worker.join().unwrap();
}
}
}
// Reproduces the actual bug mechanism, not just concurrent `block_on`
// traffic: `future_into_py` spawns an outer task that, once running,
// spawns a second (inner) one for the real work and awaits its
// `JoinHandle` (see `pyo3_async_runtimes::generic::future_into_py_with_locals`).
// Before `spawn`/`spawn_blocking` pinned nested calls to whatever
// runtime is already executing them, a reclaim landing between the
// outer and inner spawn could bind them to two different runtime
// instances -- and if the outer task's own runtime was the one torn
// down while it awaited the inner task, it never resumed. Every wait
// here is bounded so a reintroduced bug fails this test instead of
// hanging the suite.
#[test]
#[allow(unused_must_use)] // fire-and-forget outer spawn, same as future_into_py itself
fn test_nested_spawn_survives_concurrent_shutdown() {
use std::sync::mpsc;
let _lock = lock_runtime_state_for_test();
for _ in 0..50 {
let stop = Arc::new(AtomicBool::new(false));
let (done_tx, done_rx) = mpsc::channel::<()>();
let workers: Vec<_> = (0..8)
.map(|_| {
let stop = Arc::clone(&stop);
let done_tx = done_tx.clone();
std::thread::spawn(move || {
while !stop.load(Ordering::Relaxed) {
let (tx, rx) = mpsc::sync_channel::<()>(1);
LanceRuntime::spawn(async move {
let inner = LanceRuntime::spawn(async {
let _ = 1 + 1;
});
let _ = inner.await;
let _ = tx.send(());
});
// A hang here across a concurrent shutdown is
// exactly the bug this test exists to catch.
let _ = rx.recv_timeout(Duration::from_secs(2));
}
let _ = done_tx.send(());
})
})
.collect();
std::thread::sleep(Duration::from_millis(10));
shutdown(Duration::from_millis(50));
stop.store(true, Ordering::Relaxed);
for _ in 0..8 {
done_rx
.recv_timeout(Duration::from_secs(10))
.expect("worker hung after shutdown raced a nested spawn");
}
for worker in workers {
worker.join().unwrap();
}
}
}
// Forces the exact interleaving the install-race finding described: an
// installer registers as outstanding (as `spawn`'s outermost branch
// does, before it ever calls `get_runtime()`) and then pauses, while a
// concurrent `shutdown()` must not decide "nothing here" and return
// before that install actually completes and is retired in turn.
#[test]
fn test_shutdown_waits_for_a_racing_install() {
use std::sync::mpsc;
let _lock = lock_runtime_state_for_test();
// Clean slate: no runtime installed, OUTSTANDING at zero.
shutdown(Duration::from_secs(5));
let (installer_ready_tx, installer_ready_rx) = mpsc::channel::<()>();
let (proceed_tx, proceed_rx) = mpsc::channel::<()>();
let installer = std::thread::spawn(move || {
let _guard = OutstandingGuard::new();
installer_ready_tx.send(()).unwrap();
proceed_rx.recv().unwrap();
assert_eq!(get_runtime().block_on(async { 1 + 1 }), 2);
});
installer_ready_rx.recv().unwrap();
let shutdown_thread = std::thread::spawn(|| shutdown(Duration::from_secs(5)));
// Give shutdown's polling loop several chances to (wrongly) observe
// an idle runtime before the installer is allowed to proceed.
std::thread::sleep(Duration::from_millis(50));
proceed_tx.send(()).unwrap();
installer.join().unwrap();
shutdown_thread.join().unwrap();
// shutdown must have waited for the install to finish and then
// retired it, not returned early and left it stranded.
assert!(RUNTIME.load_full().is_none());
assert_eq!(OUTSTANDING.load(Ordering::SeqCst), 0);
}
}
python/src/sql.rs
+90
View File
@@ -0,0 +1,90 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::sync::Arc;
use chrono::{DateTime, Utc};
use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods};
use uuid::Uuid;
use crate::arrow::RecordBatchStream;
use crate::error::PythonErrorExt;
use crate::runtime::future_into_py;
#[pyclass(name = "SqlQuery")]
pub struct Query {
inner: Arc<lancedb::sql::Query>,
}
impl Query {
pub(crate) fn new(inner: lancedb::sql::Query) -> Self {
Self {
inner: Arc::new(inner),
}
}
}
#[pymethods]
impl Query {
#[getter]
pub fn id(&self) -> Uuid {
self.inner.id()
}
pub fn describe(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner
.describe()
.await
.map(QueryDescription::from)
.infer_error()
})
}
pub fn reader(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let stream = inner.reader().await.infer_error()?;
Ok(RecordBatchStream::new(stream))
})
}
pub fn cancel(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner.cancel().await.infer_error()?;
Ok(())
})
}
}
#[pyclass(get_all, skip_from_py_object)]
#[derive(Clone)]
pub struct QueryDescription {
id: Uuid,
status: String,
progress: Option<f64>,
expires_at: Option<DateTime<Utc>>,
}
#[pymethods]
impl QueryDescription {
fn __repr__(&self) -> String {
format!(
"QueryDescription(id={:?}, status={:?}, progress={:?}, expires_at={:?})",
self.id, self.status, self.progress, self.expires_at
)
}
}
impl From<lancedb::sql::QueryDescription> for QueryDescription {
fn from(description: lancedb::sql::QueryDescription) -> Self {
Self {
id: description.id,
status: description.status.to_string(),
progress: description.progress,
expires_at: description.expires_at,
}
}
}
python/src/table.rs
+45
View File
@@ -575,6 +575,17 @@ pub struct RefreshMaterializedViewResult {
#[pymethods]
impl RefreshMaterializedViewResult {
#[staticmethod]
pub fn from_json(value: &str) -> PyResult<Self> {
let result: lancedb::RefreshMaterializedViewResult =
serde_json::from_str(value).map_err(|err| {
PyValueError::new_err(format!(
"failed to decode materialized-view refresh result: {err}"
))
})?;
Ok(Self::from(result))
}
pub fn __repr__(&self) -> String {
format!(
"RefreshMaterializedViewResult(mode={}, rows_written={}, source_version={}, version={})",
@@ -1772,6 +1783,40 @@ impl Table {
})
}
#[pyo3(signature = (full=false, source_version=None))]
pub fn refresh_materialized_view_async(
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 job = builder.execute_async().await.infer_error()?;
Ok(crate::job::Job::new_typed(job))
})
}
pub fn materialized_view_definition(self_: PyRef<'_, Self>) -> 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()?;
serde_json::to_string(view.definition()).map_err(|err| {
PyRuntimeError::new_err(format!(
"failed to serialize materialized-view definition: {err}"
))
})
})
}
pub fn add_columns_with_schema(
self_: PyRef<'_, Self>,
schema: PyArrowType<Schema>,
python/tests/test_oauth.py
+347
View File
@@ -1,10 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import asyncio
import importlib.util
import json
import os
import subprocess
import sys
import threading
import urllib.parse
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
import pytest
def _load_oauth_module():
oauth_path = (
@@ -31,3 +40,341 @@ def test_oauth_config_repr_redacts_client_secret():
rendered = repr(config)
assert "super-secret" not in rendered
assert "client_secret" not in rendered
def test_authorization_code_uses_pkce_by_default():
oauth = _load_oauth_module()
config = oauth.OAuthConfig(
issuer_url="https://issuer.example.com",
client_id="client-id",
scopes=["openid"],
flow=oauth.OAuthFlowType.AUTHORIZATION_CODE,
)
assert config.use_pkce is True
assert config.redirect_uri is None
assert config.callback_port is None
def test_device_code_flow_value():
oauth = _load_oauth_module()
assert oauth.OAuthFlowType.DEVICE_CODE.value == "device_code"
def test_client_auth_method_values():
oauth = _load_oauth_module()
assert oauth.ClientAuthMethod.NONE.value == "none"
assert oauth.ClientAuthMethod.CLIENT_SECRET_BASIC.value == "client_secret_basic"
assert oauth.ClientAuthMethod.CLIENT_SECRET_POST.value == "client_secret_post"
def test_client_auth_method_defaults_to_none():
oauth = _load_oauth_module()
config = oauth.OAuthConfig(
issuer_url="https://issuer.example.com",
client_id="client-id",
scopes=["openid"],
client_auth_method=oauth.ClientAuthMethod.CLIENT_SECRET_POST,
)
assert config.client_auth_method is oauth.ClientAuthMethod.CLIENT_SECRET_POST
default_config = oauth.OAuthConfig(
issuer_url="https://issuer.example.com",
client_id="client-id",
scopes=["openid"],
)
assert default_config.client_auth_method is None
def test_token_cache_options_default_to_memory_only():
oauth = _load_oauth_module()
config = oauth.OAuthConfig(
issuer_url="https://issuer.example.com",
client_id="client-id",
scopes=["openid"],
)
assert config.token_cache is None
assert config.resource is None
assert config.audience is None
options = oauth.TokenCacheOptions()
assert options.cache_dir is None
assert options.lock_timeout_secs is None
def _remote_oauth():
pytest.importorskip("lancedb")
from lancedb.remote import oauth as remote_oauth
return remote_oauth
def _device_config(remote_oauth, issuer_url, cache_dir):
return remote_oauth.OAuthConfig(
issuer_url=issuer_url,
client_id="client-id",
scopes=["openid"],
flow=remote_oauth.OAuthFlowType.DEVICE_CODE,
token_cache=remote_oauth.TokenCacheOptions(cache_dir=str(cache_dir)),
)
def test_oauth_session_status_and_logout_without_cache_entry(tmp_path):
remote_oauth = _remote_oauth()
config = _device_config(remote_oauth, "https://issuer.example.com", tmp_path)
session = remote_oauth.OAuthSession(config)
status = asyncio.run(session.status())
assert status.refreshable is False
assert status.issuer_url == "https://issuer.example.com"
assert status.client_id == "client-id"
assert status.scopes == ["openid"]
assert status.flow == "device_code"
assert status.obtained_at is None
logout = asyncio.run(session.logout())
assert logout.removed is False
class _MockIdpState:
def __init__(self, port):
self.port = port
self.lock = threading.Lock()
self.device_authorizations = 0
self.refresh_grants = 0
self.invalid_grant_rejections = 0
self.access_tokens_issued = 0
self.current_refresh = None
self.requests = []
class _MockIdpHandler(BaseHTTPRequestHandler):
@property
def state(self) -> _MockIdpState:
return self.server.state
def log_message(self, fmt, *args):
pass
def _respond(self, status, payload):
body = json.dumps(payload).encode()
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def do_GET(self):
if self.path.startswith("/.well-known/openid-configuration"):
base = f"http://127.0.0.1:{self.state.port}"
self._respond(
200,
{
"token_endpoint": f"{base}/token",
"device_authorization_endpoint": f"{base}/device",
},
)
else:
self._respond(404, {})
def do_POST(self):
length = int(self.headers.get("Content-Length", 0))
body = self.rfile.read(length).decode()
params = urllib.parse.parse_qs(body)
with self.state.lock:
self.state.requests.append(params)
if self.path == "/device":
with self.state.lock:
self.state.device_authorizations += 1
base = f"http://127.0.0.1:{self.state.port}"
self._respond(
200,
{
"device_code": "device-code",
"user_code": "ABCD-EFGH",
"verification_uri": f"{base}/verify",
"expires_in": 60,
"interval": 1,
},
)
return
if self.path == "/token":
grant_type = params.get("grant_type", [""])[0]
with self.state.lock:
if grant_type == "refresh_token":
self.state.refresh_grants += 1
offered = params.get("refresh_token", [""])[0]
if offered != self.state.current_refresh:
self.state.invalid_grant_rejections += 1
self._respond(400, {"error": "invalid_grant"})
return
elif "device_code" not in grant_type:
self._respond(400, {"error": "unsupported_grant_type"})
return
self.state.access_tokens_issued += 1
number = self.state.access_tokens_issued
refresh = f"refresh-{number}"
self.state.current_refresh = refresh
self._respond(
200,
{
"access_token": f"access-{number}",
"refresh_token": refresh,
"expires_in": 3600,
},
)
return
self._respond(404, {})
def _start_mock_idp() -> tuple[_MockIdpState, HTTPServer]:
server = HTTPServer(("127.0.0.1", 0), _MockIdpHandler)
state = _MockIdpState(server.server_address[1])
server.state = state
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return state, server
def _run_subprocess(script: Path, issuer_url: str, cache_dir: Path, target: dict):
env = dict(os.environ)
env["LANCEDB_OAUTH_BROWSER"] = "/usr/bin/true"
result = subprocess.run(
[sys.executable, str(script), issuer_url, str(cache_dir), json.dumps(target)],
capture_output=True,
text=True,
timeout=120,
env=env,
)
assert result.returncode == 0, (
f"subprocess failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
)
return result
LOGIN_SCRIPT = """
import asyncio
import json
import sys
from lancedb.remote import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions
issuer_url, cache_dir = sys.argv[1], sys.argv[2]
config = OAuthConfig(
issuer_url=issuer_url,
client_id="client-id",
scopes=["openid"],
flow=OAuthFlowType.DEVICE_CODE,
token_cache=TokenCacheOptions(cache_dir=cache_dir),
**json.loads(sys.argv[3]),
)
session = OAuthSession(config)
status = asyncio.run(session.login())
assert status.refreshable, "login must cache a refresh token"
assert status.resource == config.resource
assert status.audience == config.audience
print("LOGIN-OK")
"""
REUSE_SCRIPT = """
import asyncio
import json
import sys
import lancedb
from lancedb.remote import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions
issuer_url, cache_dir = sys.argv[1], sys.argv[2]
config = OAuthConfig(
issuer_url=issuer_url,
client_id="client-id",
scopes=["openid"],
flow=OAuthFlowType.DEVICE_CODE,
token_cache=TokenCacheOptions(cache_dir=cache_dir),
**json.loads(sys.argv[3]),
)
session = OAuthSession(config)
status = asyncio.run(session.status())
assert status.refreshable, "second process must see the cached session"
async def main():
# Point the database endpoint at a dead port. OAuth headers are fetched
# before the request is sent, so a successful refresh proves the second
# process reused the cached session; only the database call fails.
db = await lancedb.connect_async(
"db://e2e",
host_override="http://127.0.0.1:1",
client_config={"retry_config": {"retries": 0}},
oauth_config=config,
)
try:
await db.table_names()
except Exception:
print("DATABASE-UNREACHABLE-AS-EXPECTED")
else:
raise AssertionError("expected the database request to fail")
asyncio.run(main())
print("REUSE-OK")
"""
@pytest.mark.parametrize(
"target",
[
{},
{
"resource": "https://api.example.com/a?x=1&y=two",
"audience": "audience + & / ü",
},
],
)
def test_cross_process_session_reuse_without_new_prompt(tmp_path, target):
pytest.importorskip("lancedb")
state, server = _start_mock_idp()
try:
issuer_url = f"http://127.0.0.1:{state.port}"
login_script = tmp_path / "login.py"
login_script.write_text(LOGIN_SCRIPT)
reuse_script = tmp_path / "reuse.py"
reuse_script.write_text(REUSE_SCRIPT)
cache_dir = tmp_path / "oauth-cache"
result = _run_subprocess(login_script, issuer_url, cache_dir, target)
assert "LOGIN-OK" in result.stdout
assert state.device_authorizations == 1
result = _run_subprocess(reuse_script, issuer_url, cache_dir, target)
assert "REUSE-OK" in result.stdout
assert "DATABASE-UNREACHABLE-AS-EXPECTED" in result.stdout
# The second process refreshed exactly once and never started a new
# interactive device flow.
assert state.refresh_grants == 1
assert state.device_authorizations == 1
assert state.invalid_grant_rejections == 0
assert len(state.requests) == 3
for params in state.requests:
for key in ("resource", "audience"):
assert params.get(key) == ([target[key]] if key in target else None)
config = _device_config(_remote_oauth(), issuer_url, cache_dir)
config.resource = target.get("resource")
config.audience = target.get("audience")
logout = asyncio.run(_remote_oauth().OAuthSession(config).logout())
assert logout.removed is True
finally:
server.shutdown()
server.server_close()
python/uv.lock
Generated
+8 -8
View File
@@ -2005,7 +2005,7 @@ requires-dist = [
{ name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" },
{ name = "pydantic", specifier = ">=2.7.4,<3" },
{ name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" },
{ name = "pylance", marker = "extra == 'tests'", specifier = "==9.0.0rc1" },
{ name = "pylance", marker = "extra == 'tests'", specifier = "==9.0.0" },
{ name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" },
{ name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" },
{ name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" },
@@ -3912,7 +3912,7 @@ crypto = [
[[package]]
name = "pylance"
version = "7.0.0"
version = "9.0.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "lance-namespace" },
@@ -3922,12 +3922,12 @@ dependencies = [
{ name = "pyarrow" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/ac/ad/2f64921bf346e7075aef24a72595db44821724a3d89a9a92dd24e79632aa/pylance-7.0.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:98422021975be76e72b1572f41b8c9abb3bee5bdc9bfa5e9ce731110a65ed4d1", size = 62134146, upload-time = "2026-05-27T21:59:37.459Z" },
{ url = "https://files.pythonhosted.org/packages/73/1c/c5a01bee0160b55d9a98895cbd33091d038f0a0995b121ab72e629008d02/pylance-7.0.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4bec86ee5b6fbd8bfc493e653f0a1fba0303cfe5492b9b46fc25ab908edc7183", size = 65373684, upload-time = "2026-05-27T22:04:01.584Z" },
{ url = "https://files.pythonhosted.org/packages/eb/da/1fe8b8f7dbfe734d76af76acc994fc360a0d0c79a4874ef69f5a72a58fe3/pylance-7.0.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:881491432c53184e52f8d1db8d5f872f39a03f36fb104bec77b33d379519d8b5", size = 69458555, upload-time = "2026-05-27T22:16:50.567Z" },
{ url = "https://files.pythonhosted.org/packages/76/f0/dd505cf3fd0226ab9d94759acd713125af1d3bfacfd80bbd52e3b9f89509/pylance-7.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:18453999e7fff4f76b16d6b7882c9df0628bd142ff95e2461bd7dd5ee3fe0af3", size = 65394430, upload-time = "2026-05-27T22:05:30.923Z" },
{ url = "https://files.pythonhosted.org/packages/17/ba/2357b81034f28eb00790e258ed140289a6a887a7468ca9df6349fd186b27/pylance-7.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:04a58051d408c60fe76d41a220dcaf8fea8fb6d1aa0ca78a709b60bc3cc8d19a", size = 69473470, upload-time = "2026-05-27T22:17:18.935Z" },
{ url = "https://files.pythonhosted.org/packages/1f/ec/5c00b6303a67d787f9475141832cbdc513d674ac3dcaeef8a7b169905e65/pylance-7.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:467d4864af047eaab4e1370e2f1e88e2c6f507c079874421116cb41d78bc3629", size = 74792863, upload-time = "2026-05-27T22:19:23.875Z" },
{ url = "https://files.pythonhosted.org/packages/e9/be/45733acd64801991852aac8e658601fd8fc12f76ceb81e57fca690896b90/pylance-9.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:8257213501d3298c5b6a344d60938e4bbe4de9f00cd3265371a56d1dc3dd15ca", size = 68377982, upload-time = "2026-07-24T16:53:45.247Z" },
{ url = "https://files.pythonhosted.org/packages/d0/3e/1ef707cb215cc7268c63ad84a91344ad6313d3984b343eeb19a9b708698e/pylance-9.0.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:804eedfa1fda2e703cca8580c76f0b44a1b849b725a8908c06f8acfca811732f", size = 71844362, upload-time = "2026-07-24T16:56:07.6Z" },
{ url = "https://files.pythonhosted.org/packages/c8/fb/a499e5c53ddb75c7de44100fd2bb1f7cc735966200d20292eed7af9ef552/pylance-9.0.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0a0b75595e3766c1d5f4c90abdc52337b4de60f1a52d123a4f8c4e5bcbbbfa8f", size = 75663088, upload-time = "2026-07-24T17:10:31.283Z" },
{ url = "https://files.pythonhosted.org/packages/e9/80/0714e09f64a68dbdf62558955e737a7436df763b9834a6aba506861b5352/pylance-9.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:d2f69c5c390ae3710c35a429905fe15f769951777fafb1b72a359e035ec121f7", size = 71866858, upload-time = "2026-07-24T16:56:35.937Z" },
{ url = "https://files.pythonhosted.org/packages/4b/3c/78d3a6d6ca0d843b7c3ac0c30d9cd2cf4635b0c2cabe6ec66583d5bfc1a1/pylance-9.0.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:836268a7832d62f3d5ccbe1c3fca239621971d90297bcfff14a70b3cb6842aa8", size = 75642656, upload-time = "2026-07-24T17:13:11.879Z" },
{ url = "https://files.pythonhosted.org/packages/cf/c1/dc9c9a31e171530ec0add024d922488c046437d32a7087c29e254a7eacc7/pylance-9.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:96441d27a5ed3805300388ccf8f31835cd280c98751f80b6a1a11dcd6808fc43", size = 81668288, upload-time = "2026-07-24T17:05:38.707Z" },
]
[[package]]