rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-16 00:32:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into refactor/wal-naming

Browse Source 
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Daniel Rammer
2026-10-12 15:30:00 -05:00
co-authored by Claude Opus 5.5
parent ab74aa620c 25fcfc696a
commit d5544cd5d9
636 changed files with 118151 additions and 22369 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/__init__.py
+55 -4
View File
@@ -6,12 +6,15 @@ import importlib.metadata
import os
from concurrent.futures import ThreadPoolExecutor
from datetime import timedelta
from typing import Dict, Optional, Union, Any, List, Iterable
from typing import Dict, Optional, Union, Any, List, Iterable, TYPE_CHECKING
__version__ = importlib.metadata.version("lancedb")
from ._lancedb import connect as lancedb_connect
from ._lancedb import FtsToken
from ._lancedb import FunctionErrorFragment as FunctionErrorFragment
from ._lancedb import FunctionErrorRecord as FunctionErrorRecord
from ._lancedb import FunctionErrors as FunctionErrors
from ._lancedb import LsmWriteSpec
from ._lancedb import tokenize as _tokenize
from .common import URI, sanitize_uri
@@ -20,9 +23,13 @@ from .db import AsyncConnection, DBConnection, LanceDBConnection
from .remote import ClientConfig
from .remote.db import RemoteDBConnection
from .expr import Expr, col, lit, func
from .schema import blob, vector, BlobType
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,11 +40,14 @@ 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,
MaterializedViewDefinition,
)
from .view import ViewDescription as ViewDescription
from .table import AsyncTable, Table
from .types import BaseTokenizerType
from ._lancedb import Session
@@ -48,6 +58,26 @@ from .namespace import (
AsyncLanceNamespaceDBConnection,
)
from .catalog import (
AsyncCatalog,
Catalog,
connect_catalog,
connect_catalog_async,
)
if TYPE_CHECKING:
from lance.blob import BlobType as BlobType
def __getattr__(name: str):
if name == "BlobType":
from .schema import BlobType
globals()["BlobType"] = BlobType
return BlobType
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
def _check_s3_bucket_with_dots(
uri: str, storage_options: Optional[Dict[str, str]]
@@ -88,6 +118,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,
@@ -116,6 +147,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
@@ -257,6 +291,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,
@@ -286,8 +321,8 @@ def tokenize(
language: str = "English",
max_token_length: Optional[int] = 40,
lower_case: bool = True,
stem: bool = True,
remove_stop_words: bool = True,
stem: Optional[bool] = None,
remove_stop_words: Optional[bool] = None,
custom_stop_words: Optional[List[str]] = None,
ascii_folding: bool = True,
ngram_min_length: int = 3,
@@ -298,6 +333,8 @@ def tokenize(
This does not require an FTS index. The tokenizer options match
:class:`lancedb.index.FTS`. ``custom_stop_words`` accepts a list of strings.
When ``stem`` or ``remove_stop_words`` is None, the base tokenizer's default
applies: False for ``code`` and ``ngram``, True otherwise.
"""
return _tokenize(
@@ -399,6 +436,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,
)
@@ -412,6 +450,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,
@@ -434,6 +473,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
@@ -521,6 +563,7 @@ async def connect_async(
api_key,
region,
host_override,
sql_host_override,
read_consistency_interval_secs,
client_config,
storage_options,
@@ -533,9 +576,14 @@ async def connect_async(
__all__ = [
"Catalog",
"AsyncCatalog",
"connect_catalog",
"connect_catalog_async",
"AsyncMaterializedView",
"MaterializedView",
"MaterializedViewDefinition",
"ViewDescription",
"connect",
"connect_async",
"tokenize",
@@ -543,6 +591,7 @@ __all__ = [
"connect_namespace_async",
"AsyncConnection",
"AsyncJob",
"AsyncSqlQuery",
"AsyncLanceNamespaceDBConnection",
"AsyncTable",
"FtsToken",
@@ -557,6 +606,8 @@ __all__ = [
"vector",
"DBConnection",
"Job",
"QueryDescription",
"SqlQuery",
"LanceDBConnection",
"LanceNamespaceDBConnection",
"LsmWriteSpec",
python/python/lancedb/_blob.py
+5 -3
View File
@@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Optional, Union
import pyarrow as pa
from .expr import Expr
from .schema import blob_v2_column_paths
from .schema import row_addressable_blob_v2_paths
from .types import BlobMode, QueryProjection, QueryProjectionSpec
if TYPE_CHECKING:
@@ -119,7 +119,7 @@ def blob_v2_projection_sources(
schema: pa.Schema,
projection: QueryProjection,
) -> dict[str, str]:
blob_columns = blob_v2_column_paths(schema)
blob_columns = row_addressable_blob_v2_paths(schema)
if not blob_columns:
return {}
columns = set(blob_columns)
@@ -140,7 +140,9 @@ def v2_projection_needs_row_id(
) -> bool:
if with_row_id:
return False
return projection_includes_blob_column(projection, blob_v2_column_paths(schema))
return projection_includes_blob_column(
projection, row_addressable_blob_v2_paths(schema)
)
def blob_auto_row_id_for_scan(
python/python/lancedb/_lancedb.pyi
+253 -13
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
@@ -57,8 +58,8 @@ def tokenize(
language: str = "English",
max_token_length: Optional[int] = 40,
lower_case: bool = True,
stem: bool = True,
remove_stop_words: bool = True,
stem: Optional[bool] = None,
remove_stop_words: Optional[bool] = None,
custom_stop_words: Optional[List[str]] = None,
ascii_folding: bool = True,
ngram_min_length: int = 3,
@@ -147,15 +148,78 @@ 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 create_function_async(self, request_json: str) -> Job: ...
async def get_function(self, name: str, version: str) -> str: ...
async def list_jobs(self) -> List[JobInfo]: ...
async def get_job(self, job_id: str) -> Optional[JobDescription]: ...
async def open_job(self, job_id: str) -> Job: ...
async def create_function_async(
self, request_json: str, namespace_path: Optional[List[str]] = None
) -> Job: ...
async def get_function(
self, name: str, version: str, namespace_path: Optional[List[str]] = None
) -> str: ...
def list_functions(
self,
namespace_path: Optional[List[str]] = None,
*,
page_token: Optional[str] = None,
page_limit: Optional[int] = None,
) -> FunctionListing: ...
async def drop_function(
self, name: str, version: str, namespace_path: Optional[List[str]] = None
) -> bool: ...
async def drop_function_async(
self, name: str, version: str, namespace_path: Optional[List[str]] = None
) -> Tuple[bool, Job]: ...
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: ...
def list_secrets(
self,
namespace_path: Optional[List[str]] = None,
*,
page_token: Optional[str] = None,
page_limit: Optional[int] = None,
) -> NameListing: ...
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 create_view(
self, name: str, query: str, namespace_path: Optional[List[str]] = None
) -> Tuple[str, List[str], str, str, List[str], pa.Schema]: ...
async def describe_view(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Tuple[str, List[str], str, str, List[str], pa.Schema]: ...
async def drop_view(
self, name: str, namespace_path: Optional[List[str]] = None
) -> None: ...
async def drop_view_async(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Job: ...
def list_views(
self,
namespace_path: Optional[List[str]] = None,
*,
page_token: Optional[str] = None,
page_limit: Optional[int] = None,
) -> NameListing: ...
def list_jobs(
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> JobListing: ...
async def cancel_job(self, job_id: str) -> bool: ...
async def job_history(
self, job_id: Optional[str] = None
) -> List[pa.RecordBatch]: ...
async def pause_job(self, job_id: str) -> str: ...
async def resume_job(self, job_id: str) -> str: ...
async def execute_query_async(
self,
query: str,
*,
default_namespace_path: Optional[List[str]] = None,
parameters: Optional[pa.RecordBatch] = None,
) -> SqlQuery: ...
async def describe_query(self, query_id: UUID) -> QueryDescription: ...
async def create_table(
self,
name: str,
@@ -205,8 +269,26 @@ 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 list_materialized_views(self) -> List[str]: ...
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: ...
def list_materialized_views(
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> NameListing: ...
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 +316,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 +343,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,14 +392,38 @@ 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: ...
def is_open(self) -> bool: ...
def _is_native(self) -> bool: ...
def close(self) -> None: ...
async def schema(self) -> pa.Schema: ...
async def add(
@@ -359,9 +507,19 @@ class Table:
) -> AddColumnsResult: ...
async def refresh_column(self, column: str) -> RefreshColumnResult: ...
async def refresh_column_async(self, column: str) -> Job: ...
async def function_errors(
self,
job_id: Optional[str] = None,
column: Optional[str] = None,
limit: Optional[int] = None,
) -> FunctionErrors: ...
async def refresh_materialized_view(
self, full: bool = False, source_version: Optional[int] = None
self, source_version: Optional[int] = None
) -> RefreshMaterializedViewResult: ...
async def refresh_materialized_view_async(
self, 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]]
@@ -382,6 +540,9 @@ class Table:
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
async def unset_lsm_write_spec(self) -> None: ...
async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ...
# WAL-PK-FUSION: delete both.
def hybrid_pk_fusion_learned(self) -> bool: ...
def note_hybrid_pk_fusion(self) -> None: ...
async def checkpoint_lsm(self) -> None: ...
async def flush_lsm(self) -> None: ...
async def compact_lsm(self) -> None: ...
@@ -450,6 +611,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]],
@@ -565,6 +727,7 @@ class VectorQuery:
def use_lsm(self, enable: bool): ...
def refine_factor(self, refine_factor: int): ...
def nprobes(self, nprobes: int): ...
def search_effort(self, search_effort: Optional[float]): ...
def minimum_nprobes(self, minimum_nprobes: int): ...
def maximum_nprobes(self, maximum_nprobes: int): ...
def bypass_vector_index(self): ...
@@ -590,6 +753,7 @@ class HybridQuery:
def distance_type(self, distance_type: str): ...
def refine_factor(self, refine_factor: int): ...
def nprobes(self, nprobes: int): ...
def search_effort(self, search_effort: Optional[float]): ...
def minimum_nprobes(self, minimum_nprobes: int): ...
def maximum_nprobes(self, maximum_nprobes: int): ...
def bypass_vector_index(self): ...
@@ -606,6 +770,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]]]
@@ -615,8 +780,10 @@ class PyQueryRequest:
use_lsm: Optional[bool]
column: Optional[str]
query_vector: Optional[List[pa.Array]]
nprobes: Optional[int]
minimum_nprobes: Optional[int]
maximum_nprobes: Optional[int]
search_effort: Optional[float]
lower_bound: Optional[float]
upper_bound: Optional[float]
ef: Optional[int]
@@ -712,7 +879,32 @@ class RefreshColumnResult:
rows_filled: int
version: int
class FunctionErrorRecord:
job_id: str
fragment_id: int
row_offset: Optional[int]
column: str
function: str
function_version: str
table_version: int
error_type: str
error_message: str
created_at_millis: int
class FunctionErrorFragment:
job_id: str
fragment_id: int
rows_skipped: int
rows_recorded: int
class FunctionErrors:
records: list[FunctionErrorRecord]
fragments: list[FunctionErrorFragment]
truncated: bool
class RefreshMaterializedViewResult:
@staticmethod
def from_json(value: str) -> RefreshMaterializedViewResult: ...
mode: str
rows_written: int
source_version: int
@@ -761,3 +953,51 @@ 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: ...
def list_databases(
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> DatabaseNames: ...
class DatabaseNames:
def num_page_results(self) -> int: ...
def page_token(self) -> Optional[str]: ...
def __aiter__(self) -> DatabaseNames: ...
async def __anext__(self) -> 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: ...
class NameListing:
def num_page_results(self) -> int: ...
def page_token(self) -> Optional[str]: ...
def __aiter__(self) -> NameListing: ...
async def __anext__(self) -> str: ...
class FunctionListing:
def num_page_results(self) -> int: ...
def page_token(self) -> Optional[str]: ...
def __aiter__(self) -> FunctionListing: ...
async def __anext__(self) -> str: ...
class JobListing:
def num_page_results(self) -> int: ...
def page_token(self) -> Optional[str]: ...
def __aiter__(self) -> JobListing: ...
async def __anext__(self) -> JobInfo: ...
python/python/lancedb/_wal_hybrid.py
+282
View File
@@ -0,0 +1,282 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""WAL-PK-FUSION: hybrid search's primary-key fallback for MemWAL tables.
Hybrid search joins its vector and full-text legs on `_rowid`, which a MemWAL
table cannot supply: the fresh tier has no stable row id, so the server rejects
`with_row_id` before it plans. When that happens the legs are joined on a
surrogate `_rowid` built from the table's primary key instead.
This is a stopgap until MemWAL supports `_rowid`, and everything it needs lives
here so it can be removed in one pass:
1. Delete this file and `tests/test_wal_hybrid.py`.
2. Run `grep -rn WAL-PK-FUSION python/ rust/` and follow the note at each
marker; the Rust side has its own copy in
`rust/lancedb/src/query/wal_fusion.rs`.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Awaitable, Callable, List, Optional, Tuple, TypeVar
import pyarrow as pa
from ._blob import blob_v2_projection_sources
if TYPE_CHECKING:
from .types import QueryProjection
T = TypeVar("T")
# Lance records the unenforced primary key as field metadata, which survives the
# `describe` round trip — so a remote table can resolve the key without opening
# the dataset. Keys are bytes on a pyarrow field.
_PK_MARKER = b"lance-schema:unenforced-primary-key"
_PK_POSITION_MARKER = b"lance-schema:unenforced-primary-key:position"
_TRUTHY = {"1", "true", "on", "yes", "y"}
def _pk_columns(schema: pa.Schema) -> List[str]:
"""The unenforced primary key columns of `schema`, in key order.
Mirrors Lance's ordering: fields with an explicit 1-based position first, in
position order, then fields carrying only the legacy boolean marker, in
declaration order.
"""
keyed = []
for index, field in enumerate(schema):
metadata = field.metadata or {}
position = metadata.get(_PK_POSITION_MARKER)
if position is not None:
try:
position = int(position)
except ValueError:
position = None
if position is None:
legacy = metadata.get(_PK_MARKER)
if legacy is None or legacy.decode().lower() not in _TRUTHY:
continue
position = 0
keyed.append((position == 0, position, index, field.name))
return [name for *_, name in sorted(keyed)]
# Both layers raise this same sentence verbatim -- sophon's
# `LSM_WITH_ROW_ID_UNSUPPORTED` says so explicitly, and the OSS LSM scanner
# builds it from the same words -- so this substring survives either path and
# the wrapping each adds. It is a coupling to a message rather than a code,
# because the refusal arrives as a generic 400; if the server ever rewords it
# the fallback stops firing and hybrid on a MemWAL table errors visibly again,
# which is loud rather than silently wrong.
_WAL_ROW_ID_REFUSAL = "does not support with_row_id"
WAL_ROW_ID_UNSUPPORTED = (
"hybrid search on a MemWAL table cannot return _rowid: the fresh tier has "
"no stable row id, and the ids the fusion joins on are synthesized from "
"the primary key. Set use_lsm(False) to read the base table only (results "
"will exclude un-compacted MemWAL data)"
)
def _is_wal_row_id_refusal(error: BaseException) -> bool:
"""Whether `error` is the server declining `_rowid` because the table is
MemWAL-backed, as opposed to any other failure the query might hit."""
return _WAL_ROW_ID_REFUSAL in str(error)
def _with_pk_columns(columns, pk_columns: List[str]):
"""Add the key columns a projection is missing.
Returns the projection to use and the columns added, which the caller drops
from the output again so an injected key never reaches a caller who did not
select it.
"""
if not columns:
# Every non-system column, the key among them.
return columns, []
# `PyQueryRequest.select` reports a dynamic projection as (name, sql) pairs.
if isinstance(columns, (list, tuple)) and isinstance(columns[0], (list, tuple)):
columns = dict(columns)
missing = [c for c in pk_columns if c not in columns]
if isinstance(columns, dict):
return {**columns, **{c: c for c in missing}}, missing
return list(columns) + missing, missing
def _stamp_surrogate_row_ids(
tables: List[pa.Table], pk_columns: List[str]
) -> List[pa.Table]:
"""Give each table a `_rowid` column derived from its primary key.
The fusion only ever compares that column for equality — to deduplicate, to
group by in a reranker, and to restore the original scores — so on a MemWAL
table, where the server has no stable row id to hand out, any injective
mapping of the primary key does the same job. Ids are assigned in first-seen
order across `tables`, so passing them in fusion order (vector, then FTS)
preserves "first occurrence wins" and its tie break.
Writing it under the `_rowid` name is what lets every reranker, third-party
implementations included, run unmodified. The column never reaches the
caller.
"""
if not pk_columns:
raise ValueError(
"hybrid search on a MemWAL table needs an unenforced primary key to "
"deduplicate on, and this table declares none"
)
ids = {}
stamped = []
for table in tables:
if table.num_rows == 0:
# Both legs empty: the combined schema carries no key column, and no
# row that would need one.
row_ids = pa.array([], type=pa.uint64())
else:
missing = [c for c in pk_columns if c not in table.column_names]
if missing:
raise ValueError(
"hybrid search could not deduplicate: primary key "
f"column(s) {missing} missing from a result set with "
f"{table.num_rows} rows"
)
keys = zip(*(table.column(c).to_pylist() for c in pk_columns))
row_ids = pa.array(
[ids.setdefault(key, len(ids)) for key in keys], type=pa.uint64()
)
stamped.append(table.append_column("_rowid", row_ids))
return stamped
def pk_fusion_learned(table, use_lsm: Optional[bool]) -> bool:
"""Whether a previous hybrid query on `table` was refused `_rowid`, so this
one should go straight to the primary key.
Learned from a refusal, never probed, so asking is free.
"""
if table is None or use_lsm is False:
return False
# Only an unambiguous True forks to the key: a table type that answers the
# question loosely keeps the `_rowid` path.
return table._hybrid_pk_fusion_learned() is True
class PkFusion:
"""How the legs of a hybrid query on a MemWAL table are joined: on a
surrogate `_rowid` built from the primary key."""
def __init__(self, pk_columns: List[str]):
self.pk_columns = pk_columns
# Key columns the caller did not select, added to the projection so the
# surrogate can be built and dropped again before returning.
self.injected: List[str] = []
@classmethod
def for_query(
cls,
schema: pa.Schema,
projection: QueryProjection,
*,
caller_requested_row_id: bool,
) -> PkFusion:
"""Resolve the key for a hybrid query, refusing the shapes that need a
real `_rowid`."""
if caller_requested_row_id:
raise NotImplementedError(WAL_ROW_ID_UNSUPPORTED)
# A projected blob needs a real row id to fetch through, whatever else
# has turned `_with_row_id` on — including `return_score="all"`.
if blob_v2_projection_sources(schema, projection):
raise NotImplementedError(
"hybrid search cannot project a blob column on a MemWAL table: "
"fetching blobs needs a real _rowid, and the fresh tier has no "
"stable one. Set use_lsm(False) to read the base table only "
"(results will exclude un-compacted MemWAL data)"
)
return cls(_pk_columns(schema))
def inject(self, columns):
"""`columns` with the key columns it is missing, to project in place of
the `_rowid` a MemWAL table would refuse."""
columns, self.injected = _with_pk_columns(columns, self.pk_columns)
return columns
def stamp(
self, vector_results: pa.Table, fts_results: pa.Table
) -> Tuple[pa.Table, pa.Table]:
"""Give both legs a surrogate `_rowid` built from the key."""
# Vector first: `merge_results` concatenates in that order and keeps
# the first occurrence, so first-seen ids preserve its tie break.
vector_results, fts_results = _stamp_surrogate_row_ids(
[vector_results, fts_results], self.pk_columns
)
return vector_results, fts_results
def strip(self, results: pa.Table) -> pa.Table:
"""Drop the surrogate and any injected key columns from the reranked
results.
The surrogate is an internal join key, never an answer, and a caller
who asked for `_rowid` was refused in `for_query`.
"""
drop = [c for c in ["_rowid", *self.injected] if c in results.column_names]
return results.drop(drop) if drop else results
def _check_refusal(error: Exception, *, fused: bool, caller_requested_row_id: bool):
"""Re-raise `error` unless it is a MemWAL `_rowid` refusal worth retrying on
the primary key."""
if fused or not _is_wal_row_id_refusal(error):
raise error
# `return_score="all"` turns row ids on for the reranker, not the caller, so
# it still falls back; only a caller who asked for them is refused, and with
# the reason rather than a bare 400.
if caller_requested_row_id:
raise NotImplementedError(WAL_ROW_ID_UNSUPPORTED) from error
def with_pk_fallback(
run: Callable[[], T],
table,
*,
fused: Callable[[], bool],
caller_requested_row_id: bool,
) -> T:
"""Run a hybrid query, retrying on the primary key if the table refuses
`_rowid`.
`_rowid` is asked for first rather than paying a round trip up front to
find out: only a MemWAL table refuses, and the refusal has to be handled
anyway -- a write spec installed elsewhere can arrive between any two
queries. It costs nothing server-side: it is raised before the query is
planned. `fused` reports whether the attempt already used the key.
"""
try:
return run()
except Exception as e:
_check_refusal(
e, fused=fused(), caller_requested_row_id=caller_requested_row_id
)
# Learned, so this retry and later queries go straight to the key.
table._note_hybrid_pk_fusion()
return run()
async def with_pk_fallback_async(
run: Callable[[], Awaitable[T]],
table,
*,
fused: Callable[[], bool],
caller_requested_row_id: bool,
) -> T:
"""The async form of `with_pk_fallback`."""
try:
return await run()
except Exception as e:
_check_refusal(
e, fused=fused(), caller_requested_row_id=caller_requested_row_id
)
table._note_hybrid_pk_fusion()
return await run()
python/python/lancedb/catalog.py
+268
View File
@@ -0,0 +1,268 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Remote catalogs manage databases through a server's root namespace."""
from datetime import timedelta
from typing import Any, AsyncIterator, Iterator, 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
class AsyncDatabaseNames(AsyncIterator[str]):
"""An async iterator of database names with pagination state.
Returned by [list_databases][lancedb.catalog.AsyncCatalog.list_databases].
"""
def __init__(self, inner: _lancedb.DatabaseNames):
self._inner = inner
def __aiter__(self) -> "AsyncDatabaseNames":
return self
async def __anext__(self) -> str:
return await self._inner.__anext__()
def num_page_results(self) -> int:
"""Return the number of names available without another REST request."""
return self._inner.num_page_results()
def page_token(self) -> Optional[str]:
"""Return the token for the next REST request.
Before the first request, this is the supplied starting token (``None``
starts at the beginning). After the final page is fetched, it is ``None``,
even if names remain cached. Drain the cache before saving a token to
avoid skipping those names when resuming. A failed request terminates
iteration and leaves its token available for resuming a new iterator.
"""
return self._inner.page_token()
class DatabaseNames(Iterator[str]):
"""A synchronous iterator of database names with pagination state.
Returned by [Catalog.list_databases][lancedb.catalog.Catalog.list_databases].
"""
def __init__(self, inner: AsyncDatabaseNames):
self._inner = inner
def __iter__(self) -> "DatabaseNames":
return self
def __next__(self) -> str:
try:
return LOOP.run(self._inner.__anext__())
except StopAsyncIteration:
raise StopIteration from None
def num_page_results(self) -> int:
"""Return the number of names available without another REST request."""
return self._inner.num_page_results()
def page_token(self) -> Optional[str]:
"""Return the token for the next REST request.
See [page_token][lancedb.catalog.AsyncDatabaseNames.page_token]
for initial, cached, final-page, and error behavior.
"""
return self._inner.page_token()
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))
def list_databases(
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> AsyncDatabaseNames:
"""Iterate lazily over all database names using ``async for``.
``page_token`` resumes from a saved token; ``None`` starts at the beginning.
``page_limit`` sets the maximum names per REST response (not the total);
``None`` uses the server default. Request errors are raised during iteration.
Examples
--------
```python
async for name in catalog.list_databases():
print(name)
```
"""
return AsyncDatabaseNames(
self._inner.list_databases(page_token=page_token, page_limit=page_limit)
)
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)
for name in catalog.list_databases():
print(name)
```
"""
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, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> DatabaseNames:
"""Iterate lazily over database names, fetching more results as needed.
``page_token`` resumes from a saved token; ``None`` starts at the beginning.
``page_limit`` sets the maximum names per REST response (not the total);
``None`` uses the server default. Request errors are raised during iteration.
"""
return DatabaseNames(
self._inner.list_databases(page_token=page_token, page_limit=page_limit)
)
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
+1203 -133
View File
File diff suppressed because it is too large Load Diff
python/python/lancedb/embeddings/base.py
+5 -1
View File
@@ -182,7 +182,11 @@ class EmbeddingFunction(BaseModel, ABC):
def VectorField(self, **kwargs):
"""
Creates a pydantic Field that can automatically annotate
the target vector column for this embedding function
the target vector column for this embedding function.
The field can be annotated as ``list[float]`` for compatibility with
static type checkers. LanceDB will infer the fixed vector dimension from
this embedding function.
"""
return Field(json_schema_extra={"vector_column_for": self}, **kwargs)
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/open_clip.py
+2 -1
View File
@@ -7,6 +7,7 @@ import io
import os
import urllib.parse as urlparse
from typing import TYPE_CHECKING, List, Union
from urllib.request import url2pathname
import numpy as np
import pyarrow as pa
@@ -154,7 +155,7 @@ class OpenClipEmbeddings(EmbeddingFunction):
parsed = urlparse.urlparse(image)
# TODO handle drive letter on windows.
if parsed.scheme == "file":
return PIL_Image.open(parsed.path)
return PIL_Image.open(url2pathname(parsed.path))
elif parsed.scheme == "":
return PIL_Image.open(image if os.name == "nt" else parsed.path)
elif parsed.scheme.startswith("http"):
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
+1206 -179
View File
File diff suppressed because it is too large Load Diff
python/python/lancedb/index.py
+45 -7
View File
@@ -112,9 +112,42 @@ class Fm:
class FTS:
"""Describe a FTS index configuration.
`FTS` is a full-text search index that can be used on `String` columns
`FTS` supports string columns, lists of strings, and JSON columns created
from PyArrow ``pa.json_()``. JSON columns use Lance's JSON-aware tokenizer
and are stored as ``LargeBinary`` with the ``lance.json`` logical type.
Arbitrary binary columns are not supported.
For example, it works with `title`, `description`, `content`, etc.
JSON queries use ``path,str,value`` triples. Paths are relative to the JSON
column and use dots for nested object keys. For example,
``meta.tag,str,nature`` matches the string value at ``$.meta.tag``;
it does not search other paths or JSON key names. The value is processed
with the configured text tokenizer. Use ``str``, not ``string``.
Create and query a JSON index through the modern API::
import pyarrow as pa
from lancedb.index import FTS
table = db.create_table(
"documents",
pa.table({"doc": pa.array(
['{"meta":{"tag":"nature"}}'], type=pa.json_()
)}),
)
table.create_index("doc", config=FTS())
table.search(
"meta.tag,str,nature", query_type="fts", fts_columns="doc"
).to_list()
Create the index on the JSON column itself, not a JSON subpath such as
``doc.meta.tag``. JSON-path scalar indexes do not provide FTS path scoping.
Plain text such as ``nature`` is not a valid JSON FTS query. The current
Lance tokenizer can panic on malformed triples; see
https://github.com/lance-format/lance/issues/8812.
Remote tables use the same ``FTS()`` configuration and require a server
that supports JSON FTS. Wait for ``table.wait_for_index(["doc_idx"])``
before searching a newly created remote index with its default name.
Examples
--------
@@ -138,6 +171,7 @@ class FTS:
- "simple": Splits text by whitespace and punctuation.
- "whitespace": Split text by whitespace, but not punctuation.
- "raw": No tokenization. The entire text is treated as a single token.
- "code": Tokenizes source code and identifiers.
- "ngram": N-gram tokenizer for substring-style matching.
- "icu": ICU dictionary-based word segmentation.
- "icu/split": ICU segmentation with simple-style delimiter splitting.
@@ -151,12 +185,16 @@ class FTS:
ignored.
lower_case : bool, default True
Whether to convert the token to lower case. This makes queries case-insensitive.
stem : bool, default True
stem : bool, optional
Whether to stem the token. Stemming reduces words to their root form.
For example, in English "running" and "runs" would both be reduced to "run".
remove_stop_words : bool, default True
``None`` uses the base tokenizer's default: False for ``code`` and
``ngram``, True otherwise.
remove_stop_words : bool, optional
Whether to remove stop words. Stop words are common words that are often
removed from text before indexing. For example, in English "the" and "and".
``None`` uses the base tokenizer's default: False for ``code`` and
``ngram``, True otherwise.
custom_stop_words : list of str, optional
Custom words replace the built-in language stop words
and only take effect when ``remove_stop_words`` is True. ``None`` uses
@@ -197,8 +235,8 @@ class FTS:
language: str = "English"
max_token_length: Optional[int] = 40
lower_case: bool = True
stem: bool = True
remove_stop_words: bool = True
stem: Optional[bool] = None
remove_stop_words: Optional[bool] = None
ascii_folding: bool = True
ngram_min_length: int = 3
ngram_max_length: int = 3
@@ -751,7 +789,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/listing.py
+85
View File
@@ -0,0 +1,85 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Lazy, resumable resource listings."""
from __future__ import annotations
from typing import Any, AsyncIterator, Callable, Generic, Iterator, Optional, TypeVar
from .background_loop import LOOP
T = TypeVar("T")
class AsyncListing(AsyncIterator[T], Generic[T]):
"""A lazy async iterator with inspectable pagination state.
Returned by connection resource-listing methods. For example::
async for name in connection.list_views(page_limit=20):
print(name)
Requests happen during iteration; request errors terminate the iterator.
"""
def __init__(self, inner: Any, convert: Callable[[Any], T] = lambda value: value):
self._inner = inner
self._convert = convert
def __aiter__(self) -> AsyncListing[T]:
return self
async def __anext__(self) -> T:
return self._convert(await self._inner.__anext__())
def num_page_results(self) -> int:
"""Return the number of cached items available without another request."""
return self._inner.num_page_results()
def page_token(self) -> Optional[str]:
"""Return the token for the next request.
Initially this is the supplied starting token (``None`` starts at the
beginning). After fetching the final page it is ``None``, even if items
remain cached. Drain the cache before saving a token to avoid skipping
items when resuming. A failed request terminates iteration and retains
its token for retrying with a new iterator.
"""
return self._inner.page_token()
class Listing(Iterator[T], Generic[T]):
"""A lazy synchronous iterator with inspectable pagination state.
For example::
for name in connection.list_views(page_limit=20):
print(name)
Use ``list(connection.list_views())`` to collect all results.
"""
def __init__(self, inner: AsyncListing[T]):
self._inner = inner
def __iter__(self) -> Listing[T]:
return self
def __next__(self) -> T:
try:
return LOOP.run(self._inner.__anext__())
except StopAsyncIteration:
raise StopIteration from None
def num_page_results(self) -> int:
"""Return the number of cached items available without another request."""
return self._inner.num_page_results()
def page_token(self) -> Optional[str]:
"""Return the next request's token.
See [AsyncListing.page_token][lancedb.listing.AsyncListing.page_token]
for initial, cached, final-page, and error behavior.
"""
return self._inner.page_token()
python/python/lancedb/materialized_view.py
+115 -41
View File
@@ -6,11 +6,13 @@ maintained by refresh. See ``DBConnection.create_materialized_view``."""
from __future__ import annotations
import json
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
import json
from dataclasses import dataclass
from .background_loop import LOOP
from .job import AsyncJob, Job, _typed_job
if TYPE_CHECKING:
import pyarrow as pa
@@ -28,20 +30,28 @@ SelectArg = Union[
]
MaterializedViewSource = str
DEFINITION_FORMAT = 2
"""The newest stored layout this version reads: ``{"format": N, "query": "<SQL>"}``,
format 2 being a query with ``GROUP BY``. A ``kind`` key beside it is for
readers older than the format number."""
@dataclass
class MaterializedViewDefinition:
"""The query that defines a materialized view."""
"""The query that defines a materialized view, as stored::
source_table: str
"""Name of the source table, in the same database as the view."""
projections: List[Tuple[str, str]]
"""``(output column, SQL expression)`` pairs, in view schema order."""
filter: Optional[str] = None
"""SQL predicate selecting the source rows the view holds."""
limit: Optional[int] = None
"""Cap on the number of rows the view holds."""
inputs: List[str] = field(default_factory=list)
"""Source columns the projections and filter read."""
SELECT columns
FROM [ns.]table [, function(args) AS alias | , UNNEST(column) AS alias]
[WHERE predicate] [GROUP BY expr, ...] [LIMIT n]
A Function in ``FROM`` position yields one row per element it returns.
"""
query: str
"""The defining query, in the canonical spelling the server stores."""
def _definition_from_schema(
@@ -51,22 +61,64 @@ def _definition_from_schema(
raw = metadata.get(DEFINITION_META_KEY)
if raw is None:
raise ValueError(f"Table '{name}' is not a materialized view")
value = json.loads(raw)
return _definition_from_value(json.loads(raw), name)
def _definition_from_json(raw: str, name: str = "") -> MaterializedViewDefinition:
"""Parse the definition native code hands over, in its stored layout."""
return _definition_from_value(json.loads(raw), name)
def _definition_from_value(value: dict, name: str) -> MaterializedViewDefinition:
fmt = value.get("format")
if fmt is not None:
# A newer writer's layout is reported, never guessed at.
if not isinstance(fmt, int) or fmt > DEFINITION_FORMAT:
raise NotImplementedError(
f"materialized view '{name}' is stored in format {fmt}, which "
"this version of lancedb cannot refresh"
)
return MaterializedViewDefinition(query=value["query"])
# The structured layout written before the format number.
kind = value.get("kind")
if kind != "select":
if kind not in ("select", "namespaced_select"):
raise NotImplementedError(
f"materialized view '{name}' is defined by '{kind}', which this "
"version of lancedb cannot refresh"
f"materialized view '{name}' is stored in format kind '{kind}', "
"which this version of lancedb cannot refresh"
)
return MaterializedViewDefinition(
source_table=value["source_table"],
projections=[
(p["output"], p["expression"]) for p in value.get("projections", [])
],
filter=value.get("filter"),
limit=value.get("limit"),
inputs=value.get("inputs", []),
return MaterializedViewDefinition(query=_legacy_query(value))
def _legacy_ident(name: str) -> str:
if name and all(c == "_" or c.islower() or c.isdigit() for c in name):
return name
return _quote_identifier(name)
def _legacy_query(value: dict) -> str:
"""Render the pre-format structured layout as the query it described."""
projections = value.get("projections", [])
if projections:
items = []
for p in projections:
output, expression = p["output"], p["expression"]
if expression in (output, _quote_identifier(output)):
items.append(expression)
else:
items.append(f"{expression} AS {_legacy_ident(output)}")
columns = ", ".join(items)
else:
columns = "*"
table = ".".join(
_legacy_ident(part)
for part in [*value.get("source_namespace", []), value["source_table"]]
)
query = f"SELECT {columns} FROM {table}"
if value.get("filter") is not None:
query += f" WHERE {value['filter']}"
if value.get("limit") is not None:
query += f" LIMIT {value['limit']}"
return query
def _quote_identifier(name: str) -> str:
@@ -122,26 +174,38 @@ 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
self, *, source_version: Optional[int] = None
) -> "RefreshMaterializedViewResult":
"""Recompute the view from its source.
The refresh is incremental when the source's changes can be
reconciled into the view -- rows added, changed or removed since the
last one -- and otherwise rebuilds. ``full=True`` forces a rebuild;
Every refresh rebuilds the complete view. For local views,
``source_version`` refreshes to that source version instead of the
latest.
Concurrent refreshes of one view do not duplicate its rows. Two that
plan the same source rows conflict on commit, and the loser raises
rather than writing them a second time.
latest; remote SQL refreshes do not support source-version pinning.
"""
return await self._table._inner.refresh_materialized_view(
full=full, source_version=source_version
source_version=source_version
)
async def refresh_async(
self, *, 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. Remote refreshes are
submitted through the SQL service.
"""
from ._lancedb import RefreshMaterializedViewResult
return _typed_job(
await self._table._inner.refresh_materialized_view_async(
source_version=source_version
),
RefreshMaterializedViewResult.from_json,
)
@@ -167,12 +231,22 @@ 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
self, *, source_version: Optional[int] = None
) -> "RefreshMaterializedViewResult":
"""Recompute the view from its source. See
[AsyncMaterializedView.refresh][lancedb.materialized_view.AsyncMaterializedView.refresh]."""
return LOOP.run(self._async.refresh(full=full, source_version=source_version))
return LOOP.run(self._async.refresh(source_version=source_version))
def refresh_async(
self, *, 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(source_version=source_version)))
python/python/lancedb/namespace.py
+203 -16
View File
@@ -11,15 +11,15 @@ through a namespace abstraction.
from __future__ import annotations
import sys
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
from typing import Any, Dict, Iterable, List, Optional, Union
from uuid import UUID
if sys.version_info >= (3, 12):
from typing import override
else:
from overrides import override
if TYPE_CHECKING:
from lancedb.query import Query
from lancedb.query import Query
from datetime import timedelta
import pyarrow as pa
@@ -48,8 +48,12 @@ 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.listing import AsyncListing, Listing
from lancedb.sql import AsyncQuery as AsyncSqlQuery
from lancedb.sql import QueryDescription, QueryParameters
from lance_namespace import (
LanceNamespace,
connect as namespace_connect,
@@ -64,6 +68,7 @@ from lance_namespace import (
from lancedb.materialized_view import (
AsyncMaterializedView,
MaterializedView,
MaterializedViewSource,
SelectArg,
)
from lancedb.table import AsyncTable, LanceTable, Table
@@ -100,6 +105,11 @@ def _query_to_namespace_request(
if not isinstance(query, Query):
raise TypeError(f"Expected Query, got {type(query)}")
if query.search_effort is not None:
raise ValueError(
"namespace QueryTable does not support search_effort; "
"execute the query locally"
)
vector_request = None
if query.vector is not None:
@@ -185,8 +195,12 @@ def _query_to_namespace_request(
kwargs["with_row_id"] = query.with_row_id
if query.ef is not None:
kwargs["ef"] = query.ef
if query.nprobes is not None:
kwargs["nprobes"] = query.nprobes
if query.minimum_nprobes is not None:
kwargs["nprobes"] = query.minimum_nprobes
kwargs["minimum_nprobes"] = query.minimum_nprobes
if query.maximum_nprobes is not None:
kwargs["maximum_nprobes"] = query.maximum_nprobes
if query.refine_factor is not None:
kwargs["refine_factor"] = query.refine_factor
if query.lower_bound is not None:
@@ -200,6 +214,13 @@ def _query_to_namespace_request(
if query.distance_type is not None:
kwargs["distance_type"] = query.distance_type
if kwargs.get("maximum_nprobes") == 0:
# The generated 0.13 model rejects LanceDB's unbounded sentinel;
# model_copy preserves direct wire pass-through to the Lance receiver.
maximum_nprobes = kwargs.pop("maximum_nprobes")
return QueryTableRequest(**kwargs).model_copy(
update={"maximum_nprobes": maximum_nprobes}
)
return QueryTableRequest(**kwargs)
@@ -484,6 +505,14 @@ class LanceNamespaceDBConnection(DBConnection):
)
self._uri = self._inner.uri
@override
def is_open(self) -> bool:
return self._inner.is_open()
@override
def close(self) -> None:
self._inner.close()
@override
def serialize(self) -> str:
import json
@@ -509,13 +538,15 @@ class LanceNamespaceDBConnection(DBConnection):
def table_names(
self,
page_token: Optional[str] = None,
limit: int = 10,
limit: Optional[int] = None,
*,
namespace_path: Optional[List[str]] = None,
) -> Iterable[str]:
"""
List table names in the database.
If limit is None, return all table names; zero returns an empty list.
.. deprecated::
Use :meth:`list_tables` instead, which provides proper pagination support.
"""
@@ -628,11 +659,12 @@ class LanceNamespaceDBConnection(DBConnection):
def create_materialized_view(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
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 +674,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: MaterializedViewSource,
*,
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``."""
@@ -656,9 +716,44 @@ class LanceNamespaceDBConnection(DBConnection):
return view
@override
def list_materialized_views(self) -> List[str]:
"""The names of the materialized views in the root namespace."""
return LOOP.run(self._inner.list_materialized_views())
def list_materialized_views(
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> Listing[str]:
"""The names of the materialized views in the root namespace.
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
``page_token`` resumes from a saved token. Requests and errors occur during
iteration. See [pagination state][lancedb.listing.AsyncListing].
"""
return Listing(
self._inner.list_materialized_views(
page_token=page_token, page_limit=page_limit
)
)
@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):
@@ -1080,13 +1175,15 @@ class AsyncLanceNamespaceDBConnection:
async def table_names(
self,
page_token: Optional[str] = None,
limit: int = 10,
limit: Optional[int] = None,
*,
namespace_path: Optional[List[str]] = None,
) -> Iterable[str]:
"""
List table names in the namespace.
If limit is None, return all table names; zero returns an empty list.
.. deprecated::
Use :meth:`list_tables` instead, which provides proper pagination support.
"""
@@ -1185,29 +1282,84 @@ class AsyncLanceNamespaceDBConnection:
async def create_materialized_view(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
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: MaterializedViewSource,
*,
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))
await view.definition()
return view
async def list_materialized_views(self) -> List[str]:
"""The names of the materialized views in the root namespace."""
return await self._inner.list_materialized_views()
def list_materialized_views(
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> AsyncListing[str]:
"""The names of the materialized views in the root namespace.
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
``page_token`` resumes from a saved token. Requests and errors occur during
iteration. See [pagination state][lancedb.listing.AsyncListing].
"""
return AsyncListing(
self._inner.list_materialized_views(
page_token=page_token, page_limit=page_limit
)
)
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."""
@@ -1447,6 +1599,41 @@ 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,
parameters: Optional[QueryParameters] = None,
) -> AsyncRecordBatchReader:
"""Execute SQL when supported by the underlying connection."""
return await self._inner.execute_query(
query,
default_namespace_path=default_namespace_path,
parameters=parameters,
)
async def execute_query_async(
self,
query: str,
*,
default_namespace_path: Optional[List[str]] = None,
parameters: Optional[QueryParameters] = 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,
parameters=parameters,
)
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/pydantic.py
+39
View File
@@ -88,6 +88,17 @@ def Vector(
... pa.field("url", pa.utf8(), False),
... pa.field("embeddings", pa.list_(pa.float32(), 768))
... ])
Notes
-----
``Vector`` creates a type dynamically, so calls such as ``Vector(768)`` are
not valid static type annotations. For an embedding field, use the standard
``list[float]`` annotation when running mypy; ``VectorField`` supplies the
fixed dimension to LanceDB::
class MyModel(LanceModel):
text: str = embeddings.SourceField()
vector: list[float] = embeddings.VectorField()
"""
# TODO: make a public parameterized type.
@@ -313,6 +324,10 @@ def _unwrap_optional_annotation(annotation: Any) -> Any | None:
def _pydantic_to_arrow_type(field: FieldInfo) -> pa.DataType:
"""Convert a Pydantic FieldInfo to Arrow DataType"""
embedding_vector_type = _embedding_vector_to_arrow_type(field)
if embedding_vector_type is not None:
return embedding_vector_type
unwrapped = _unwrap_optional_annotation(field.annotation)
if unwrapped is not None:
return _pydantic_type_to_arrow_type(unwrapped, field)
@@ -326,8 +341,32 @@ def _pydantic_to_arrow_type(field: FieldInfo) -> pa.DataType:
return _pydantic_type_to_arrow_type(field.annotation, field)
def _embedding_vector_to_arrow_type(field: FieldInfo) -> pa.DataType | None:
"""Infer a fixed-size vector type from ``VectorField`` metadata."""
if not _is_embedding_vector_annotation(field):
return None
function = get_extras(field, "vector_column_for")
return pa.list_(pa.float32(), function.ndims())
def _is_embedding_vector_annotation(field: FieldInfo) -> bool:
if get_extras(field, "vector_column_for") is None:
return False
annotation = _unwrap_optional_annotation(field.annotation)
if annotation is None:
annotation = field.annotation
origin = getattr(annotation, "__origin__", None)
args = getattr(annotation, "__args__", ())
return origin is list and args == (float,)
def is_nullable(field: FieldInfo) -> bool:
"""Check if a Pydantic FieldInfo is nullable."""
if _is_embedding_vector_annotation(field):
return True
if _unwrap_optional_annotation(field.annotation) is not None:
return True
if isinstance(field.annotation, (_GenericAlias, GenericAlias)):
python/python/lancedb/query.py
+343 -68
View File
@@ -39,8 +39,9 @@ from .expr import Expr
from .rerankers.base import Reranker
from .rerankers.rrf import RRFReranker
from .rerankers.util import check_reranker_result
from .schema import is_blob_like_field, schema_has_blob_field
from .util import flatten_columns
from .schema import blob_column_paths, is_blob_like_field, schema_has_blob_field
from .util import _validate_query_vector, flatten_columns
from . import _wal_hybrid # WAL-PK-FUSION: delete.
from ._blob import (
BLOB_MODE_TO_HANDLING,
FetchBlobsAsync,
@@ -50,6 +51,7 @@ from ._blob import (
finalize_blob_query_table,
replace_v2_blob_columns_with_bytes,
replace_v2_blob_columns_with_bytes_sync,
strip_auto_row_ids,
validate_blob_mode,
)
from .types import BlobMode, QueryProjection
@@ -78,6 +80,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 +115,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
)
@@ -468,12 +475,12 @@ class MatchQuery(FullTextQuery):
----------
query : str
The query string to match against.
column : str
The name of the column to match against.
column : str, optional
The column to match against. None leaves column selection to query planning.
boost : float, default 1.0
The boost factor for the query.
The score of each matching document is multiplied by this value.
fuzziness : int, optional
fuzziness : int or None, default 0
The maximum edit distance for each term in the match query.
Defaults to 0 (exact match).
If None, fuzziness is applied automatically by the rules:
@@ -498,9 +505,9 @@ class MatchQuery(FullTextQuery):
"""
query: str
column: str
column: Optional[str] = None
boost: float = pydantic.Field(1.0, kw_only=True)
fuzziness: int = pydantic.Field(0, kw_only=True)
fuzziness: Optional[int] = pydantic.Field(0, kw_only=True)
max_expansions: int = pydantic.Field(50, kw_only=True)
operator: FullTextOperator = pydantic.Field(FullTextOperator.OR, kw_only=True)
prefix_length: int = pydantic.Field(0, kw_only=True)
@@ -686,8 +693,9 @@ class Query(pydantic.BaseModel):
if True then apply the filter after vector / FTS search. This is ignored for
plain SQL filtering.
nprobes : Optional[int]
The number of IVF partitions to search. If this is None then a default
number of partitions will be used.
The legacy number of IVF partitions to search. Lance sets both probe
bounds to this value. If this is None then Lance's default probe settings
will be used.
- A higher number makes search more accurate but also slower.
@@ -766,22 +774,28 @@ class Query(pydantic.BaseModel):
# distance type to use for vector search
distance_type: Optional[str] = None
# legacy number of IVF partitions to search
#
# Lance sets both probe bounds to this value. Explicit bounds can override it.
nprobes: Optional[int] = None
# which columns to return in the results (dict values may be str or Expr)
columns: QueryProjection = None
# minimum number of IVF partitions to search
#
# If None then a default value (20) will be used.
# If None then Lance's default will be used.
minimum_nprobes: Optional[int] = None
# maximum number of IVF partitions to search
#
# If None then a default value (20) will be used.
#
# If 0 then no limit will be applied and all partitions could be searched
# if needed to satisfy the limit.
# If None then Lance's default will be used. If 0 then no limit will be applied
# and all partitions could be searched if needed to satisfy the limit.
maximum_nprobes: Optional[int] = None
# Optional initial IVF effort; None preserves Auto and explicit nprobes.
search_effort: Optional[float] = None
# lower bound for distance search
lower_bound: Optional[float] = None
@@ -804,6 +818,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 +843,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
@@ -835,8 +854,10 @@ class Query(pydantic.BaseModel):
query.vector_column = req.column
query.vector = req.query_vector
query.distance_type = req.distance_type
query.nprobes = req.nprobes
query.minimum_nprobes = req.minimum_nprobes
query.maximum_nprobes = req.maximum_nprobes
query.search_effort = req.search_effort
query.lower_bound = req.lower_bound
query.upper_bound = req.upper_bound
query.ef = req.ef
@@ -853,7 +874,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)
@@ -896,6 +917,9 @@ class LanceQueryBuilder(ABC):
fast_search: bool
Skip flat search of unindexed data.
"""
if query_type != "fts":
_validate_query_vector(query)
if ordering_field_name is not None:
import warnings
@@ -1608,11 +1632,14 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
str_query: Optional[str] = None,
fast_search: bool = None,
):
_validate_query_vector(query)
super().__init__(table)
self._query = query
self._distance_type = None
self._nprobes = None
self._minimum_nprobes = None
self._maximum_nprobes = None
self._search_effort = None
self._lower_bound = None
self._upper_bound = None
self._refine_factor = None
@@ -1675,22 +1702,44 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
See discussion in [Querying an ANN Index](https://lancedb.com/docs/indexing/)
for tuning advice.
This method sets both the minimum and maximum number of probes to the same
value. See `minimum_nprobes` and `maximum_nprobes` for more fine-grained
control.
The value is retained as `nprobes` through client and server request
construction. Lance sets both probe bounds to this value. Explicit minimum
or maximum settings can override their respective bounds.
Parameters
----------
nprobes: int
The number of probes to use.
The number of probes to use. Must be greater than 0.
Returns
-------
LanceVectorQueryBuilder
The LanceQueryBuilder object.
"""
self._minimum_nprobes = nprobes
self._maximum_nprobes = nprobes
if nprobes <= 0:
raise ValueError("Invalid input, nprobes must be greater than 0")
self._nprobes = nprobes
return self
def search_effort(self, search_effort: Optional[float]) -> LanceVectorQueryBuilder:
"""Set the initial IVF search effort to a finite value in [0, 1].
None (the default) uses Auto, internally 0.5. Zero starts at the caller
minimum; one searches all partitions within the caller maximum. Later
probing may continue when too few candidates are found. Any explicit
effort, including 0.5, overrides nprobes in either call order. Explicit
minimum and maximum bounds still apply. This is not a recall guarantee
and does not change refinement or HNSW ef. MemWAL LSM queries do not
yet support this option.
Examples
--------
>>> import lancedb
>>> db = lancedb.connect("memory://")
>>> table = db.create_table("effort_vectors", [{"vector": [0.1, 0.2]}])
>>> query = table.search([0.1, 0.2]).search_effort(0.75)
"""
self._search_effort = search_effort
return self
def minimum_nprobes(self, minimum_nprobes: int) -> LanceVectorQueryBuilder:
@@ -1832,8 +1881,10 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
limit=self._limit,
distance_type=self._distance_type,
columns=self._columns,
nprobes=self._nprobes,
minimum_nprobes=self._minimum_nprobes,
maximum_nprobes=self._maximum_nprobes,
search_effort=self._search_effort,
lower_bound=self._lower_bound,
upper_bound=self._upper_bound,
refine_factor=self._refine_factor,
@@ -2190,13 +2241,21 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
self._fts_columns = fts_columns
self._norm = None
self._reranker = None
self._nprobes = None
self._minimum_nprobes = None
self._maximum_nprobes = None
self._search_effort = None
self._refine_factor = None
self._distance_type = None
self._phrase_query = None
self._lower_bound = None
self._upper_bound = None
# WAL-PK-FUSION: delete both, and the `with_row_id` override below.
# Set by `_create_query_builders` when the legs are fused on the key.
self._pk_fusion: Optional[_wal_hybrid.PkFusion] = None
# `_with_row_id` is also turned on by `rerank(return_score="all")` for
# the reranker's own use, so it cannot answer "did the caller ask?".
self._caller_requested_row_id = False
def _validate_query(self, query, vector=None, text=None):
if query is not None and (vector is not None or text is not None):
@@ -2218,6 +2277,16 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
return vector_query, text_query
def with_row_id(self, with_row_id: bool) -> Self: # WAL-PK-FUSION: delete.
"""Set whether to return row ids.
Recorded separately from `_with_row_id`, which `rerank(return_score=
"all")` also sets for its own use — on a MemWAL table only a caller who
asked is refused, and the reranker still falls back to the primary key.
"""
self._caller_requested_row_id = with_row_id
return super().with_row_id(with_row_id)
def phrase_query(self, phrase_query: bool = True) -> LanceHybridQueryBuilder:
"""Set whether to use phrase query.
@@ -2238,16 +2307,30 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
raise NotImplementedError("to_query_object not yet supported on a hybrid query")
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
# WAL-PK-FUSION: without the fallback, this body is `_run_hybrid`'s.
return _wal_hybrid.with_pk_fallback(
lambda: self._run_hybrid(timeout=timeout),
self._table,
fused=lambda: self._pk_fusion is not None,
caller_requested_row_id=self._caller_requested_row_id,
)
def _run_hybrid(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
self._create_query_builders()
fts_query, vector_query = self._fts_query, self._vector_query
# WAL-PK-FUSION: without the fallback, both legs always ask for row ids.
if self._pk_fusion is None:
fts_query = fts_query.with_row_id(True)
vector_query = vector_query.with_row_id(True)
with ThreadPoolExecutor() as executor:
fts_future = executor.submit(
self._fts_query.with_row_id(True).to_arrow, timeout=timeout
)
vector_future = executor.submit(
self._vector_query.with_row_id(True).to_arrow, timeout=timeout
)
fts_future = executor.submit(fts_query.to_arrow, timeout=timeout)
vector_future = executor.submit(vector_query.to_arrow, timeout=timeout)
fts_results = fts_future.result()
vector_results = vector_future.result()
if self._pk_fusion is not None: # WAL-PK-FUSION: delete.
vector_results, fts_results = self._pk_fusion.stamp(
vector_results, fts_results
)
results = self._combine_hybrid_results(
fts_results=fts_results,
@@ -2255,13 +2338,15 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
norm=self._norm,
fts_query=self._fts_query._query,
reranker=self._reranker,
limit=self._limit,
limit=self._limit or DEFAULT_HYBRID_LIMIT,
with_row_ids=True,
offset=self._offset,
)
return self._finish_hybrid_results(results)
def _finish_hybrid_results(self, results: pa.Table) -> pa.Table:
if self._pk_fusion is not None: # WAL-PK-FUSION: delete.
return self._pk_fusion.strip(results)
if self._user_requested_row_id():
return results
if self._blob_auto_row_id_enabled():
@@ -2435,7 +2520,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
self._norm = normalize
self._reranker = reranker
if reranker.score == "all":
self.with_row_id(True)
# WAL-PK-FUSION: without the fallback, `self.with_row_id(True)`.
# Not through `with_row_id`, which records caller intent: the
# reranker's need for row ids must not refuse a MemWAL table.
self._with_row_id = True
return self
@@ -2446,18 +2534,48 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
Higher values will yield better recall (more likely to find vectors if
they exist) at the expense of latency.
The value is retained as `nprobes` until Lance sets both probe bounds to it.
Explicit minimum or maximum settings can override their respective bounds.
Parameters
----------
nprobes: int
The number of probes to use.
The number of probes to use. Must be greater than 0.
Returns
-------
LanceHybridQueryBuilder
The LanceHybridQueryBuilder object.
"""
self._minimum_nprobes = nprobes
self._maximum_nprobes = nprobes
if nprobes <= 0:
raise ValueError("Invalid input, nprobes must be greater than 0")
self._nprobes = nprobes
return self
def search_effort(self, search_effort: Optional[float]) -> LanceHybridQueryBuilder:
"""Set the initial IVF search effort to a finite value in [0, 1].
None (the default) uses Auto, internally 0.5. Zero starts at the caller
minimum; one searches all partitions within the caller maximum. Later
probing may continue when too few candidates are found. Any explicit
effort, including 0.5, overrides nprobes in either call order. Explicit
minimum and maximum bounds still apply. This is not a recall guarantee
and does not change refinement or HNSW ef. MemWAL LSM queries do not
yet support this option.
Examples
--------
>>> import lancedb
>>> db = lancedb.connect("memory://")
>>> table = db.create_table(
... "effort_hybrid", [{"vector": [0.1, 0.2], "text": "dog"}]
... )
>>> query = (
... table.search(query_type="hybrid").vector([0.1, 0.2])
... .text("dog").search_effort(0.75)
... )
"""
self._search_effort = search_effort
return self
def minimum_nprobes(self, minimum_nprobes: int) -> LanceHybridQueryBuilder:
@@ -2702,20 +2820,31 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
)
# Apply common configurations
if self._limit:
# The final offset/limit window is sliced out of the combined,
# reranked results, so each sub-query must fetch enough rows to
# cover the skipped prefix as well as the window itself.
sub_query_limit = self._limit + (self._offset or 0)
self._vector_query.limit(sub_query_limit)
self._fts_query.limit(sub_query_limit)
if self._columns:
self._vector_query.select(self._columns)
self._fts_query.select(self._columns)
# The final offset/limit window is sliced out of the combined,
# reranked results, so each sub-query must fetch enough rows to
# cover the skipped prefix as well as the window itself.
limit = self._limit or DEFAULT_HYBRID_LIMIT
sub_query_limit = limit + (self._offset or 0)
self._vector_query.limit(sub_query_limit)
self._fts_query.limit(sub_query_limit)
# WAL-PK-FUSION: without the fallback, select `self._columns` as is.
self._pk_fusion = None
columns = self._columns
if _wal_hybrid.pk_fusion_learned(self._table, self._use_lsm):
self._pk_fusion = _wal_hybrid.PkFusion.for_query(
self._table.schema,
self._columns,
caller_requested_row_id=self._caller_requested_row_id,
)
columns = self._pk_fusion.inject(columns)
if columns:
self._vector_query.select(columns)
self._fts_query.select(columns)
if self._where:
self._vector_query.where(self._where, not self._postfilter)
self._fts_query.where(self._where, not self._postfilter)
if self._with_row_id:
# WAL-PK-FUSION: without the fallback, `if self._with_row_id:`.
if self._with_row_id and self._pk_fusion is None:
self._vector_query.with_row_id(True)
self._fts_query.with_row_id(True)
if self._use_lsm is not None:
@@ -2725,6 +2854,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
self._fts_query.phrase_query(True)
if self._distance_type:
self._vector_query.metric(self._distance_type)
if self._nprobes is not None:
self._vector_query.nprobes(self._nprobes)
if self._search_effort is not None:
self._vector_query.search_effort(self._search_effort)
if self._minimum_nprobes is not None:
self._vector_query.minimum_nprobes(self._minimum_nprobes)
if self._maximum_nprobes is not None:
@@ -3008,14 +3141,18 @@ class AsyncQueryBase(object):
If not specified, no timeout is applied. If the query does not
complete within the specified time, an error will be raised.
blob_mode: str, default "lazy"
Controls how blob columns are returned for plain scan queries.
Vector, FTS, hybrid, and other non-native query shapes keep the
existing Arrow conversion path and only support blob descriptions.
Controls how blob columns are returned. Remote queries support
"descriptions"; "bytes" and "lazy" are not yet supported.
**kwargs
Forwarded to pyarrow.Table.to_pandas after query execution and
optional flattening.
"""
validate_blob_mode(blob_mode)
if self._table is not None and not self._table._inner._is_native():
return await self._remote_to_pandas(
flatten=flatten, timeout=timeout, blob_mode=blob_mode, **kwargs
)
if hasattr(self._inner, "output_schema"):
schema = await self.output_schema()
if _blob_mode_requires_native_pandas(blob_mode, schema):
@@ -3043,6 +3180,31 @@ class AsyncQueryBase(object):
)
return tbl.to_pandas(**kwargs)
async def _remote_to_pandas(
self,
*,
flatten: Optional[Union[int, bool]],
timeout: Optional[timedelta],
blob_mode: BlobMode,
**kwargs,
) -> "pd.DataFrame":
# A live remote table can be replaced between query and blob fetch.
# Row ids and version numbers do not identify the producing table.
tbl = await self.to_arrow(timeout=timeout)
projected_blob_paths = set(blob_column_paths(tbl.schema))
if not projected_blob_paths:
return flatten_columns(tbl, flatten).to_pandas(**kwargs)
if blob_mode == "descriptions":
tbl = strip_auto_row_ids(tbl, self._blob_paths)
return flatten_columns(tbl, flatten).to_pandas(**kwargs)
raise NotImplementedError(
f"remote to_pandas(blob_mode={blob_mode!r}) cannot safely materialize "
"blob columns without a stable table snapshot; "
"use blob_mode='descriptions'"
)
async def _plain_scan_to_pandas(
self,
blob_mode: BlobMode,
@@ -3595,7 +3757,9 @@ class AsyncVectorQueryBase:
def nprobes(self, nprobes: int) -> Self:
"""
Set the number of partitions to search (probe)
Set the legacy IVF probe parameter
The number of probes must be greater than 0.
This argument is only used when the vector column has an IVF-based index.
If there is no index then this value is ignored.
@@ -3605,20 +3769,49 @@ class AsyncVectorQueryBase:
The partition whose centroids are closest to the query vector will be
exhaustiely searched to find matches. This parameter controls how many
partitions should be searched.
partitions will be searched.
Increasing this value will increase the recall of your query but will
also increase the latency of your query. The default value is 20. This
default is good for many cases but the best value to use will depend on
your data and the recall that you need to achieve.
also increase the latency of your query. If this method is not called,
Lance's adaptive probe defaults are used.
For best results we recommend tuning this parameter with a benchmark against
your actual data to find the smallest possible value that will still give
you the desired recall.
LanceDB retains this as `nprobes` through local and remote request
construction. Lance sets both probe bounds to this value. Explicit minimum
or maximum settings can override their respective bounds.
"""
self._inner.nprobes(nprobes)
return self
def search_effort(self, search_effort: Optional[float]) -> Self:
"""Set the initial IVF search effort to a finite value in [0, 1].
None preserves Auto (internally 0.5) and fixed nprobes. Any explicit
effort, including 0.5, overrides nprobes in either call order; explicit
minimum and maximum bounds remain effective. Zero starts at the caller
minimum and one searches all partitions within the maximum. Later
probing may continue to obtain enough candidates. This is not a recall
guarantee and does not change refinement or HNSW ef. MemWAL LSM queries
do not yet support this option.
Examples
--------
>>> import asyncio
>>> import lancedb
>>> async def build_query():
... db = await lancedb.connect_async("memory://")
... table = await db.create_table(
... "effort_async", [{"vector": [0.1, 0.2]}]
... )
... return table.vector_search([0.1, 0.2]).search_effort(0.75)
>>> query = asyncio.run(build_query())
"""
self._inner.search_effort(search_effort)
return self
def minimum_nprobes(self, minimum_nprobes: int) -> Self:
"""Set the minimum number of probes to use.
@@ -3856,6 +4049,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
self._inner = inner
self._norm = "score"
self._reranker = RRFReranker()
# WAL-PK-FUSION: delete.
# Set by `_create_child_queries` when the legs are fused on the key.
self._pk_fusion: Optional[_wal_hybrid.PkFusion] = None
def rerank(
self, reranker: Reranker = RRFReranker(), normalize: str = "score"
@@ -3887,14 +4083,90 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
return self
async 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. The key the legs are fused on, if any, is left
on `self._pk_fusion`.
"""
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
# WAL-PK-FUSION: without the fallback, keep only the two `with_row_id`
# calls.
self._pk_fusion = None
if _wal_hybrid.pk_fusion_learned(self._table, fts_req.use_lsm):
self._pk_fusion = _wal_hybrid.PkFusion.for_query(
await self._table.schema(),
_query_request_projection(fts_req),
caller_requested_row_id=self._user_requested_row_id(),
)
if self._pk_fusion is None:
fts_query.with_row_id()
vec_query.with_row_id()
else:
# Ask for the key columns instead, and never for `_rowid`.
# `select` carries the whole projection; `select_source_columns`
# keeps only plain column references, so rebuilding from it would
# drop computed ones.
columns = self._pk_fusion.inject(fts_req.select)
if self._pk_fusion.injected:
fts_query.select(columns)
vec_query.select(columns)
# 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)
# WAL-PK-FUSION: without the fallback, this body is `_run_hybrid`'s.
return await _wal_hybrid.with_pk_fallback_async(
lambda: self._run_hybrid(
max_batch_length=max_batch_length, timeout=timeout
),
self._table,
fused=lambda: self._pk_fusion is not None,
caller_requested_row_id=self._user_requested_row_id(),
)
async def _run_hybrid(
self,
*,
max_batch_length: Optional[int] = None,
timeout: Optional[timedelta] = None,
) -> AsyncRecordBatchReader:
fts_query, vec_query, limit, offset = await self._create_child_queries()
req = fts_query._inner.to_query_request()
blob_auto_row_id = False
@@ -3914,13 +4186,14 @@ 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),
)
if self._pk_fusion is not None: # WAL-PK-FUSION: delete.
vector_results, fts_results = self._pk_fusion.stamp(
vector_results, fts_results
)
result = LanceHybridQueryBuilder._combine_hybrid_results(
fts_results=fts_results,
@@ -3928,10 +4201,13 @@ 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 (
if self._pk_fusion is not None: # WAL-PK-FUSION: delete this branch.
result = self._pk_fusion.strip(result)
elif (
not self._user_requested_row_id()
and not blob_auto_row_id
and "_rowid" in result.column_names
@@ -3958,14 +4234,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 +4256,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, _, _ = await 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 +4285,12 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
-------
plan : str
"""
fts_query, vec_query, _, _ = await 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
+395 -50
View File
@@ -2,13 +2,27 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from __future__ import annotations
from dataclasses import replace
from datetime import timedelta
import json
import logging
from concurrent.futures import ThreadPoolExecutor
import sys
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
from typing import (
TYPE_CHECKING,
Any,
Dict,
Iterable,
List,
Optional,
Sequence,
Tuple,
Union,
)
from urllib.parse import urlparse
from uuid import UUID
import warnings
if sys.version_info >= (3, 12):
@@ -22,13 +36,18 @@ from lancedb.remote import ClientConfig, RetryConfig, TimeoutConfig, TlsConfig
import pyarrow as pa
from ..common import DATA
from lancedb.listing import Listing
from ..db import DBConnection, LOOP
from ..functions import FunctionVersion, UdfDefinition
from ..job import AsyncJob, Job
from ..materialized_view import MaterializedView, SelectArg
from ..sql import Query as SqlQuery
from ..sql import QueryDescription, QueryParameters
from ..materialized_view import MaterializedView, MaterializedViewSource, SelectArg
from ..secrets import EnvVarSecret, SecretInfo
from ..view import ViewDescription
if TYPE_CHECKING:
from .._lancedb import JobDescription, JobInfo
from .._lancedb import JobInfo
from ..embeddings import EmbeddingFunctionConfig
from lance_namespace import (
LanceNamespace,
@@ -116,6 +135,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 +181,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 +196,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 +255,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,
}
@@ -316,7 +379,7 @@ class RemoteDBConnection(DBConnection):
page_token: Optional[str] = None,
limit: Optional[int] = None,
) -> ListTablesResponse:
"""List all tables in this database with pagination support.
"""List a page of tables in this database.
Parameters
----------
@@ -327,7 +390,8 @@ class RemoteDBConnection(DBConnection):
Token for pagination. Use the token from a previous response
to get the next page of results.
limit: int, optional
The maximum number of results to return.
The maximum number of results to return, default 100.
Zero returns an empty page without a continuation token.
Returns
-------
@@ -346,7 +410,7 @@ class RemoteDBConnection(DBConnection):
def table_names(
self,
page_token: Optional[str] = None,
limit: int = 10,
limit: Optional[int] = None,
*,
namespace_path: Optional[List[str]] = None,
) -> Iterable[str]:
@@ -362,8 +426,9 @@ class RemoteDBConnection(DBConnection):
Empty list represents root namespace.
page_token: str
The last token to start the new page.
limit: int, default 10
The maximum number of tables to return for each page.
limit: int, optional
The maximum number of tables to return. None returns all tables;
zero returns an empty list.
Returns
-------
@@ -423,12 +488,12 @@ class RemoteDBConnection(DBConnection):
namespace_path = []
if storage_options is not None:
logging.info(
"storage_options is ignored in LanceDb Cloud"
"storage_options is ignored for remote tables"
" (storage is managed; set storage_options on connect() instead)"
)
if index_cache_size is not None:
logging.info(
"index_cache_size is ignored in LanceDb Cloud"
"index_cache_size is ignored for remote tables"
" (there is no local cache to configure)"
)
@@ -548,6 +613,9 @@ class RemoteDBConnection(DBConnection):
One of "error", "drop", "fill".
fill_value: float
The value to use when filling vectors. Only used if on_bad_vectors="fill".
embedding_functions: list of EmbeddingFunctionConfig, optional
The embedding functions to store in the table schema and use to
generate vectors when creating the table or adding data.
Returns
-------
@@ -622,12 +690,6 @@ class RemoteDBConnection(DBConnection):
if namespace_path is None:
namespace_path = []
validate_table_name(name)
if embedding_functions is not None:
logging.warning(
"embedding_functions is not yet supported on LanceDB Cloud."
"Please vote https://github.com/lancedb/lancedb/issues/626 "
"for this feature."
)
from .table import RemoteTable
@@ -640,6 +702,7 @@ class RemoteDBConnection(DBConnection):
schema=schema,
on_bad_vectors=on_bad_vectors,
fill_value=fill_value,
embedding_functions=embedding_functions,
)
)
return RemoteTable(
@@ -653,27 +716,97 @@ class RemoteDBConnection(DBConnection):
def create_materialized_view(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
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: MaterializedViewSource,
*,
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, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> Listing[str]:
"""List materialized views.
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
``page_token`` resumes from a saved token. Requests and errors occur during
iteration. See [pagination state][lancedb.listing.AsyncListing].
"""
return Listing(
self._conn.list_materialized_views(
page_token=page_token, page_limit=page_limit
)
)
@override
def list_materialized_views(self) -> List[str]:
raise NotImplementedError(
"materialized views are supported only on local databases"
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,35 +865,200 @@ 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,
namespace_path: Optional[List[str]] = None,
) -> Job[FunctionVersion]:
job = LOOP.run(
self._conn.create_function_async(
definition, secrets=secrets, namespace_path=namespace_path
)
)
return Job(job)
@override
def get_function(self, name: str, *, version: str) -> FunctionVersion:
return LOOP.run(self._conn.get_function(name, version=version))
def get_function(
self,
name: str,
*,
version: str,
namespace_path: Optional[List[str]] = None,
) -> FunctionVersion:
return LOOP.run(
self._conn.get_function(
name, version=version, 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())
def list_functions(
self,
*,
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
page_limit: Optional[int] = None,
) -> Listing[FunctionVersion]:
"""List functions.
@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.
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
``page_token`` resumes from a saved token. Requests and errors occur during
iteration. See [pagination state][lancedb.listing.AsyncListing].
"""
return LOOP.run(self._conn.get_job(job_id))
return Listing(
self._conn.list_functions(
namespace_path=namespace_path,
page_token=page_token,
page_limit=page_limit,
)
)
@override
def drop_function(
self,
name: str,
*,
version: str,
namespace_path: Optional[List[str]] = None,
) -> bool:
return LOOP.run(
self._conn.drop_function(
name, version=version, namespace_path=namespace_path
)
)
@override
def drop_function_async(
self,
name: str,
*,
version: str,
namespace_path: Optional[List[str]] = None,
) -> Tuple[bool, Job]:
dropped, job = LOOP.run(
self._conn.drop_function_async(
name, version=version, namespace_path=namespace_path
)
)
return dropped, Job(job)
@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,
page_token: Optional[str] = None,
page_limit: Optional[int] = None,
) -> Listing[str]:
"""List secrets.
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
``page_token`` resumes from a saved token. Requests and errors occur during
iteration. See [pagination state][lancedb.listing.AsyncListing].
"""
return Listing(
self._conn.list_secrets(
namespace_path=namespace_path,
page_token=page_token,
page_limit=page_limit,
)
)
@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 create_view(
self, name: str, query: str, *, namespace_path: Optional[List[str]] = None
) -> ViewDescription:
return LOOP.run(
self._conn.create_view(name, query, namespace_path=namespace_path)
)
@override
def describe_view(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> ViewDescription:
return LOOP.run(self._conn.describe_view(name, namespace_path=namespace_path))
@override
def drop_view(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> None:
LOOP.run(self._conn.drop_view(name, namespace_path=namespace_path))
@override
def drop_view_async(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> Job[None]:
job = LOOP.run(self._conn.drop_view_async(name, namespace_path=namespace_path))
return Job(job)
@override
def list_views(
self,
*,
namespace_path: Optional[List[str]] = None,
page_token: Optional[str] = None,
page_limit: Optional[int] = None,
) -> Listing[str]:
"""List views.
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
``page_token`` resumes from a saved token. Requests and errors occur during
iteration. See [pagination state][lancedb.listing.AsyncListing].
"""
return Listing(
self._conn.list_views(
namespace_path=namespace_path,
page_token=page_token,
page_limit=page_limit,
)
)
@override
def list_jobs(
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> Listing[JobInfo]:
"""List server-side jobs across the database's tables.
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
``page_token`` resumes from a saved token. Requests and errors occur during
iteration. See [pagination state][lancedb.listing.AsyncListing].
"""
return Listing(
self._conn.list_jobs(page_token=page_token, page_limit=page_limit)
)
@override
def cancel_job(self, job_id: str) -> bool:
@@ -773,12 +1071,55 @@ 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 pause_job(self, job_id: str) -> str:
"""Pause a server-side job by id.
Lists history across all jobs when `job_id` is None.
Returns "pausing", "already_paused", or "committing".
"""
return LOOP.run(self._conn.job_history(job_id))
return LOOP.run(self._conn.pause_job(job_id))
@override
def resume_job(self, job_id: str) -> str:
"""Resume a paused server-side job by id.
Returns "resumed", "still_pausing", or "not_paused".
"""
return LOOP.run(self._conn.resume_job(job_id))
@override
def execute_query_async(
self,
query: str,
*,
default_namespace_path: Optional[List[str]] = None,
parameters: Optional[QueryParameters] = None,
) -> SqlQuery:
"""Start executing SQL through this remote connection.
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.
``parameters`` binds values to ``$1`` / ``$name`` placeholders; see
[DBConnection.execute_query_async][lancedb.db.DBConnection.execute_query_async].
"""
return SqlQuery(
LOOP.run(
self._conn.execute_query_async(
query,
default_namespace_path=default_namespace_path,
parameters=parameters,
)
)
)
@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:
@@ -793,6 +1134,10 @@ class RemoteDBConnection(DBConnection):
"""
return LOOP.run(self._conn.namespace_client())
async def close(self):
"""Close the connection to the database."""
@override
def is_open(self) -> bool:
return self._conn.is_open()
@override
def close(self) -> None:
self._conn.close()
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
+88 -54
View File
@@ -42,6 +42,7 @@ from lancedb.index import (
BTree,
Bitmap,
HnswFlat,
HnswPq,
HnswSq,
IvfFlat,
IvfPq,
@@ -63,12 +64,22 @@ from lancedb.table import _normalize_progress
from ..query import (
AnalyzePlanDistributedMetrics,
DocumentGranularity,
FullTextQuery,
LanceQueryBuilder,
LanceTakeQueryBuilder,
LanceVectorQueryBuilder,
)
from ..table import AsyncTable, BlobMode, Branches, IndexStatistics, Query, Table, Tags
from ..table import (
AsyncTable,
BlobMode,
Branches,
IndexStatistics,
Query,
Table,
Tags,
)
from ..types import BaseTokenizerType
from ..util import infer_vector_column_name
class RemoteTable(Table):
@@ -229,12 +240,12 @@ class RemoteTable(Table):
return LOOP.run(self._table.list_versions())
def to_arrow(self) -> pa.Table:
"""to_arrow() is not yet supported on LanceDB cloud."""
raise NotImplementedError("to_arrow() is not yet supported on LanceDB cloud.")
"""to_arrow() is not supported for remote tables."""
raise NotImplementedError("to_arrow() is not supported for remote tables.")
def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs):
"""to_pandas() is not yet supported on LanceDB cloud."""
raise NotImplementedError("to_pandas() is not yet supported on LanceDB cloud.")
"""to_pandas() is not supported for remote tables."""
raise NotImplementedError("to_pandas() is not supported for remote tables.")
def checkout(self, version: Union[int, str]):
result = LOOP.run(self._table.checkout(version))
@@ -343,8 +354,8 @@ class RemoteTable(Table):
language: str = "English",
max_token_length: Optional[int] = 40,
lower_case: bool = True,
stem: bool = True,
remove_stop_words: bool = True,
stem: Optional[bool] = None,
remove_stop_words: Optional[bool] = None,
custom_stop_words: Optional[List[str]] = None,
ascii_folding: bool = True,
ngram_min_length: int = 3,
@@ -491,14 +502,8 @@ class RemoteTable(Table):
if accelerator is not None:
logging.warning(
"GPU accelerator is not yet supported on LanceDB cloud."
"If you have 100M+ vectors to index,"
"please contact us at contact@lancedb.com"
)
if replace is not None:
logging.warning(
"replace is not supported on LanceDB cloud."
"Existing indexes will always be replaced."
"GPU accelerator is not supported for remote tables "
"and will be ignored."
)
idx_type = index_type.upper()
@@ -518,9 +523,11 @@ class RemoteTable(Table):
elif idx_type == "IVF_SQ":
config = IvfSq(distance_type=metric, num_partitions=num_partitions)
elif idx_type == "IVF_HNSW_PQ":
raise ValueError(
"IVF_HNSW_PQ is not supported on LanceDB cloud."
"Please use IVF_HNSW_SQ instead."
config = HnswPq(
distance_type=metric,
num_partitions=num_partitions,
num_sub_vectors=num_sub_vectors,
num_bits=num_bits,
)
elif idx_type == "IVF_HNSW_SQ":
config = HnswSq(distance_type=metric, num_partitions=num_partitions)
@@ -540,6 +547,7 @@ class RemoteTable(Table):
LOOP.run(
self._table.create_index(
column,
replace=replace,
config=config,
wait_timeout=wait_timeout,
name=name,
@@ -648,7 +656,7 @@ class RemoteTable(Table):
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Not supported on LanceDB Cloud. Setting this raises.
Not supported for remote tables. Setting this raises.
Returns
-------
@@ -678,7 +686,7 @@ class RemoteTable(Table):
vector_column_name: Optional[str] = None,
query_type="auto",
fts_columns: Optional[Union[str, List[str]]] = None,
fast_search: bool = False,
fast_search: Optional[bool] = None,
) -> LanceVectorQueryBuilder:
"""Create a search query to find the nearest neighbors
of the given query vector. We currently support
@@ -711,7 +719,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
@@ -725,11 +733,17 @@ class RemoteTable(Table):
- If the table has multiple vector columns then the *vector_column_name*
needs to be specified. Otherwise, an error is raised.
query_type: str, default "auto"
The type of search: "auto", "vector", "fts", or "hybrid".
With "auto", a string query uses the embedding function for the
inferred or specified vector column if available, otherwise it uses
full-text search.
fast_search: bool, optional
Skip a flat search of unindexed data. This may improve
search performance but search results will not include unindexed data.
- *default False*.
- *default None*, which leaves fast search disabled.
Returns
-------
@@ -744,12 +758,27 @@ class RemoteTable(Table):
- and also the "_distance" column which is the distance between the query
vector and the returned vector.
"""
if isinstance(query, FullTextQuery):
query_type = "fts"
elif (
query_type == "auto"
and isinstance(query, str)
and not self.embedding_functions
):
query_type = "fts"
vector_column_name = infer_vector_column_name(
schema=self.schema,
query_type=query_type,
query=query,
vector_column_name=vector_column_name,
)
return LanceQueryBuilder.create(
self,
query,
query_type,
vector_column_name=vector_column_name,
fts_columns=fts_columns,
fts_columns=fts_columns or [],
fast_search=fast_search,
)
@@ -924,26 +953,20 @@ class RemoteTable(Table):
def cleanup_old_versions(self, *_):
"""
cleanup_old_versions() is a no-op on LanceDB Cloud.
Tables are automatically cleaned up and optimized.
cleanup_old_versions() is a no-op for remote tables.
"""
warnings.warn(
"cleanup_old_versions() is a no-op on LanceDB Cloud. "
"Tables are automatically cleaned up and optimized.",
"cleanup_old_versions() is a no-op for remote tables.",
stacklevel=2,
)
pass
def compact_files(self, *_):
"""
compact_files() is a no-op on LanceDB Cloud.
Tables are automatically compacted and optimized.
compact_files() is a no-op for remote tables.
"""
warnings.warn(
"compact_files() is a no-op on LanceDB Cloud. "
"Tables are automatically compacted and optimized.",
"compact_files() is a no-op for remote tables.",
stacklevel=2,
)
pass
@@ -954,17 +977,13 @@ class RemoteTable(Table):
cleanup_older_than: Optional[timedelta] = None,
delete_unverified: bool = False,
):
"""
optimize() is a no-op on LanceDB Cloud.
Indices are optimized automatically.
"""
warnings.warn(
"optimize() is a no-op on LanceDB Cloud. "
"Indices are optimized automatically.",
stacklevel=2,
"""optimize() is not supported for remote tables."""
return LOOP.run(
self._table.optimize(
cleanup_older_than=cleanup_older_than,
delete_unverified=delete_unverified,
)
)
pass
def count_rows(self, filter: Optional[str] = None) -> int:
return LOOP.run(self._table.count_rows(filter))
@@ -973,6 +992,9 @@ class RemoteTable(Table):
self,
transforms: Dict[str, str | FunctionApplication]
| FunctionApplication
| pa.Field
| List[pa.Field]
| pa.Schema
| None = None,
*,
computed: Dict[str, str] | None = None,
@@ -985,6 +1007,16 @@ class RemoteTable(Table):
def refresh_column_async(self, column: str) -> Job[RefreshColumnResult]:
return Job(LOOP.run(self._table.refresh_column_async(column)))
def function_errors(
self,
job_id: Optional[str] = None,
column: Optional[str] = None,
limit: Optional[int] = None,
):
return LOOP.run(
self._table.function_errors(job_id=job_id, column=column, limit=limit)
)
def alter_columns(
self, *alterations: Iterable[Dict[str, str]]
) -> AlterColumnsResult:
@@ -999,7 +1031,7 @@ class RemoteTable(Table):
return LOOP.run(self._table.drop_columns(columns))
def set_unenforced_primary_key(self, columns: Union[str, Iterable[str]]) -> None:
"""Not supported on LanceDB Cloud."""
"""Not supported for remote tables."""
return LOOP.run(self._table.set_unenforced_primary_key(columns))
def set_lsm_write_spec(self, spec: "LsmWriteSpec") -> None:
@@ -1014,6 +1046,14 @@ class RemoteTable(Table):
"""Read the installed LsmWriteSpec, or ``None``."""
return LOOP.run(self._table.get_lsm_write_spec())
# WAL-PK-FUSION: delete both hooks.
def _hybrid_pk_fusion_learned(self) -> bool:
# `self._table` is the async table, which owns the Rust handle.
return self._table._hybrid_pk_fusion_learned()
def _note_hybrid_pk_fusion(self) -> None:
self._table._note_hybrid_pk_fusion()
def checkpoint_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""
@@ -1037,7 +1077,7 @@ class RemoteTable(Table):
)
def close_lsm_writers(self) -> None:
"""No-op on LanceDB Cloud (no local shard writers)."""
"""No-op for remote tables (no local shard writers)."""
return LOOP.run(self._table.close_lsm_writers())
def drop_index(self, index_name: str):
@@ -1105,14 +1145,10 @@ class RemoteTable(Table):
return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids))
def uses_v2_manifest_paths(self) -> bool:
raise NotImplementedError(
"uses_v2_manifest_paths() is not supported on the LanceDB Cloud"
)
return LOOP.run(self._table.uses_v2_manifest_paths())
def migrate_v2_manifest_paths(self):
raise NotImplementedError(
"migrate_v2_manifest_paths() is not supported on the LanceDB Cloud"
)
return LOOP.run(self._table.migrate_manifest_paths_v2())
def blob_columns(self) -> list[str]:
return LOOP.run(self._table.blob_columns())
@@ -1123,9 +1159,7 @@ class RemoteTable(Table):
return LOOP.run(self._table.fetch_blobs(column, row_ids))
def fetch_blob_ranges(self, column: str, requests) -> pa.LargeBinaryArray:
raise NotImplementedError(
"fetch_blob_ranges() is not supported on LanceDB Cloud"
)
return LOOP.run(self._table.fetch_blob_ranges(column, list(requests)))
def fetch_blob_files(
self, column: str, row_ids: Union[list[int], pa.Table]
python/python/lancedb/rerankers/__init__.py
+2
View File
@@ -11,6 +11,7 @@ from .jinaai import JinaReranker
from .rrf import RRFReranker
from .mrr import MRRReranker
from .answerdotai import AnswerdotaiRerankers
from .typesafe import TypeSafeReranker
from .voyageai import VoyageAIReranker
from .watsonx import WatsonxReranker
@@ -27,6 +28,7 @@ __all__ = [
"JinaReranker",
"RRFReranker",
"AnswerdotaiRerankers",
"TypeSafeReranker",
"VoyageAIReranker",
"MRRReranker",
"WatsonxReranker",
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/rerankers/cross_encoder.py
+12 -1
View File
@@ -31,6 +31,12 @@ class CrossEncoderReranker(Reranker):
trust_remote_code : bool, default True
If True, will trust the remote code to be safe. If False, will not trust
the remote code and will not run it
batch_size : int, default 32
The number of query-passage pairs scored per forward pass.
**kwargs
Additional keyword arguments passed to
``sentence_transformers.CrossEncoder``, for example ``max_length`` or
``model_kwargs={"torch_dtype": "bfloat16"}``.
"""
def __init__(
@@ -40,6 +46,8 @@ class CrossEncoderReranker(Reranker):
device: Union[str, None] = None,
return_score="relevance",
trust_remote_code: bool = True,
batch_size: int = 32,
**kwargs,
):
super().__init__(return_score)
torch = attempt_import_or_raise("torch")
@@ -47,6 +55,8 @@ class CrossEncoderReranker(Reranker):
self.column = column
self.device = device
self.trust_remote_code = trust_remote_code
self.batch_size = batch_size
self.kwargs = kwargs
if self.device is None:
self.device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -61,6 +71,7 @@ class CrossEncoderReranker(Reranker):
self.model_name,
device=self.device,
trust_remote_code=self.trust_remote_code,
**self.kwargs,
)
return cross_encoder
@@ -71,7 +82,7 @@ class CrossEncoderReranker(Reranker):
return result_set
passages = result_set[self.column].to_pylist()
cross_inp = [[query, passage] for passage in passages]
cross_scores = self.model.predict(cross_inp)
cross_scores = self.model.predict(cross_inp, batch_size=self.batch_size)
result_set = result_set.append_column(
"_relevance_score", pa.array(cross_scores, type=pa.float32())
)
python/python/lancedb/rerankers/rrf.py
+20 -4
View File
@@ -107,8 +107,24 @@ class RRFReranker(Reranker):
`search().with_row_id(True)`"
)
combined = pa.concat_tables(vector_results, **self._concat_tables_args)
empty_table = pa.Table.from_arrays([], names=[])
reranked = self.rerank_hybrid(query, combined, empty_table)
# RRF ranks each result list on its own, so a document's rank is its
# position within the list it came from, not within the concatenation.
rrf_score_map = defaultdict(float)
for result in vector_results:
for i, result_id in enumerate(result["_rowid"].to_pylist(), 1):
rrf_score_map[result_id] += 1 / (i + self.K)
return reranked
combined = pa.concat_tables(vector_results, **self._concat_tables_args)
combined = self._deduplicate(combined)
relevance_scores = [
rrf_score_map[row_id] for row_id in combined["_rowid"].to_pylist()
]
combined = combined.append_column(
"_relevance_score", pa.array(relevance_scores, type=pa.float32())
)
combined = combined.sort_by([("_relevance_score", "descending")])
if self.score == "relevance":
combined = self._keep_relevance_score(combined)
return combined
python/python/lancedb/rerankers/typesafe.py
+230
View File
@@ -0,0 +1,230 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from concurrent.futures import ThreadPoolExecutor
from functools import cached_property
from itertools import chain
from typing import Any, Dict, List, Mapping, Optional
import pyarrow as pa
from ..util import attempt_import_or_raise
from .base import Reranker
DEFAULT_INSTRUCTIONS = (
"Does the document contain information that answers or directly addresses "
"the query?"
)
DEFAULT_CRITERIA = {
"true": "The document answers the query or states facts the query asks about.",
"false": "The document is off-topic, or only shares keywords or a general "
"subject with the query without addressing it.",
}
# The key TypeSafe returns the answer under. It is not sent to the model.
_QUESTION_ID = "relevance"
class TypeSafeReranker(Reranker):
"""
Reranks the results using the TypeSafe System One API.
https://docs.typesafe.ai/cookbooks/rerank_typesafe
Each result is scored independently: TypeSafe reads the query together with
the result's ``column`` value and answers a yes/no ("noul") question about
whether the document is relevant. The probability of "yes", between 0 and 1,
becomes the ``_relevance_score``. Because every score is an absolute
probability rather than a position in the list, scores are comparable
across queries and can be thresholded. They are model estimates, though,
and can vary slightly between identical calls, so results with close scores
may swap places when the same search is repeated.
By default, one request is sent per non-null result. Set ``batch_size=40``
to score up to 40 documents per request, with at most ``max_concurrency``
requests in flight. Each question sees only its own document and the shared
query. Null documents keep a score of zero without an API request; empty
strings are scored normally. API and response-validation errors propagate.
Parameters
----------
model_name : str, default "jev-latest"
The TypeSafe model to use.
column : str, default "text"
The name of the column holding the document text to score.
instructions : str, optional
The yes/no question asked about each query and document pair. The state
TypeSafe reads with ``batch_size=1`` is
``{"query": <query>, "document": <column value>}``.
Defaults to a generic relevance question. See ``batch_size`` for the
different location of the document in batched requests.
criteria : Mapping[str, str], optional
What a "yes" and a "no" mean, as a mapping with the keys ``"true"`` and
``"false"``. Domain-specific criteria usually rank better than the
generic default. Pass an empty mapping to send no criteria.
return_score : str, default "relevance"
Options are "relevance" or "all". If "all", keeps the vector and FTS
scores alongside the relevance score.
api_key : str, optional
The API key to use. If None, the TypeSafe SDK reads the
``TYPESAFE_API_KEY`` environment variable.
max_concurrency : int, default 8
The maximum number of TypeSafe requests in flight for one rerank call.
batch_size : int, default 1
Positive integer limiting non-null documents per request. The default
preserves the original request format. With a value greater than one,
state is ``{"query": <query>}`` and each question's instructions are
``{"question": <instructions>, "document": <column value>}``.
Instructions and criteria are preserved verbatim. Custom prompts that
explicitly reference the document in request state (for example,
``state.document``) must be adapted before opting into batching.
The TypeSafe SDK handles retries in both modes.
Examples
--------
>>> reranker = TypeSafeReranker(batch_size=40, max_concurrency=8)
"""
def __init__(
self,
model_name: str = "jev-latest",
column: str = "text",
instructions: Optional[str] = None,
criteria: Optional[Mapping[str, str]] = None,
return_score: str = "relevance",
api_key: Optional[str] = None,
max_concurrency: int = 8,
batch_size: int = 1,
):
super().__init__(return_score)
if max_concurrency < 1:
raise ValueError("max_concurrency must be at least 1")
if (
isinstance(batch_size, bool)
or not isinstance(batch_size, int)
or batch_size < 1
):
raise ValueError("batch_size must be a positive integer")
criteria = DEFAULT_CRITERIA if criteria is None else dict(criteria)
unknown = set(criteria) - {"true", "false"}
if unknown:
raise ValueError(
f"criteria keys must be 'true' or 'false', got {sorted(unknown)}"
)
self.model_name = model_name
self.column = column
self.instructions = instructions or DEFAULT_INSTRUCTIONS
self.criteria = criteria
self.api_key = api_key
self.max_concurrency = max_concurrency
self.batch_size = batch_size
def __str__(self):
return f"TypeSafeReranker(model_name={self.model_name})"
@cached_property
def _client(self):
typesafe_sdk = attempt_import_or_raise("typesafe_sdk", "typesafe-sdk")
return typesafe_sdk.TypeSafeClient(api_key=self.api_key)
@cached_property
def _question(self) -> Dict[str, Any]:
question: Dict[str, Any] = {"type": "noul", "instructions": self.instructions}
if self.criteria:
question["criteria"] = self.criteria
return question
def _score_batch(self, query: str, documents: List[str]) -> List[float]:
if self.batch_size == 1:
state = {"query": query, "document": documents[0]}
questions = {_QUESTION_ID: self._question}
else:
state = {"query": query}
# IDs only need to be unique within this request, even for duplicates.
questions = {
str(index): {
**self._question,
"instructions": {
"question": self.instructions,
"document": document,
},
}
for index, document in enumerate(documents)
}
response = self._client.system_one(
state=state,
questions=questions,
model=self.model_name,
)
expected = set(questions)
actual = set(response.answers)
if actual != expected:
raise ValueError(
"TypeSafe returned mismatched answer IDs: "
f"missing {sorted(expected - actual)}, "
f"unexpected {sorted(actual - expected)}"
)
scores = []
for question_id in questions:
score = getattr(response.answers[question_id], "noul", None)
if (
isinstance(score, bool)
or not isinstance(score, (int, float))
or not 0 <= score <= 1
):
raise ValueError(
"TypeSafe returned an invalid relevance probability "
f"for answer {question_id!r}: {score!r}"
)
scores.append(float(score))
return scores
def _rerank(self, result_set: pa.Table, query: str) -> pa.Table:
result_set = self._handle_empty_results(result_set)
if len(result_set) == 0:
return result_set
docs = result_set[self.column].to_pylist()
documents = [doc for doc in docs if doc is not None]
batches = [
documents[start : start + self.batch_size]
for start in range(0, len(documents), self.batch_size)
]
# Rerankers are also called synchronously from inside the async query
# APIs, so the requests run on threads rather than on an event loop.
with ThreadPoolExecutor(max_workers=self.max_concurrency) as pool:
batch_scores = pool.map(
lambda batch: self._score_batch(query, batch), batches
)
# Both map and _score_batch preserve input order.
scored_documents = chain.from_iterable(batch_scores)
scores = [0.0 if doc is None else next(scored_documents) for doc in docs]
result_set = result_set.append_column(
"_relevance_score", pa.array(scores, type=pa.float32())
)
return result_set.sort_by([("_relevance_score", "descending")])
def rerank_hybrid(
self,
query: str,
vector_results: pa.Table,
fts_results: pa.Table,
):
if self.score == "all":
combined_results = self._merge_and_keep_scores(vector_results, fts_results)
else:
combined_results = self.merge_results(vector_results, fts_results)
combined_results = self._rerank(combined_results, query)
if self.score == "relevance":
combined_results = self._keep_relevance_score(combined_results)
return combined_results
def rerank_vector(self, query: str, vector_results: pa.Table):
vector_results = self._rerank(vector_results, query)
if self.score == "relevance":
vector_results = vector_results.drop_columns(["_distance"])
return vector_results
def rerank_fts(self, query: str, fts_results: pa.Table):
fts_results = self._rerank(fts_results, query)
if self.score == "relevance":
fts_results = fts_results.drop_columns(["_score"])
return fts_results
python/python/lancedb/schema.py
+101 -34
View File
@@ -4,30 +4,34 @@
"""Schema helpers for Lance blob columns."""
import importlib
from typing import TYPE_CHECKING
import pyarrow as pa
import pyarrow.ipc
if TYPE_CHECKING:
from lance.blob import BlobType as BlobType
_BLOB_EXTENSION_NAME = "lance.blob.v2"
_BLOB_V1_KEY = "lance-encoding:blob"
_ARROW_EXT_NAME_KEY = "ARROW:extension:name"
_BLOB_V2_STORAGE_TYPE = pa.struct(
[
pa.field("data", pa.large_binary(), nullable=True),
pa.field("uri", pa.utf8(), nullable=True),
pa.field("position", pa.uint64(), nullable=True),
pa.field("size", pa.uint64(), nullable=True),
]
)
_resolved_blob_type = None
class BlobType(pa.ExtensionType):
"""PyArrow extension type for a Lance blob v2 column.
Queries return descriptors; call :meth:`~lancedb.table.Table.fetch_blob_files`
for lazy reads or :meth:`~lancedb.table.Table.fetch_blobs` for eager bytes.
"""
class _FallbackBlobType(pa.ExtensionType):
"""lance.blob.v2 extension type used when pylance is not installed."""
def __init__(self) -> None:
storage_type = pa.struct(
[
pa.field("data", pa.large_binary(), nullable=True),
pa.field("uri", pa.utf8(), nullable=True),
pa.field("position", pa.uint64(), nullable=True),
pa.field("size", pa.uint64(), nullable=True),
]
)
super().__init__(storage_type, _BLOB_EXTENSION_NAME)
pa.ExtensionType.__init__(self, _BLOB_V2_STORAGE_TYPE, _BLOB_EXTENSION_NAME)
def __arrow_ext_serialize__(self) -> bytes:
return b""
@@ -35,23 +39,16 @@ class BlobType(pa.ExtensionType):
@classmethod
def __arrow_ext_deserialize__(
cls, storage_type: pa.DataType, serialized: bytes
) -> "BlobType":
) -> "_FallbackBlobType":
return cls()
def __reduce__(self):
# Ensure pickle round-trips on older pyarrow (apache/arrow#35599).
return type(self).__arrow_ext_deserialize__, (
self.storage_type,
self.__arrow_ext_serialize__(),
)
try:
pa.register_extension_type(BlobType()) # type: ignore[arg-type]
except pa.ArrowKeyError:
pass
def _metadata_value(metadata: dict, key: str):
return metadata.get(key.encode()) or metadata.get(key)
@@ -92,43 +89,105 @@ def is_blob_like_field(field: pa.Field) -> bool:
return is_blob_v2_field(field) or _metadata_marks_legacy_blob(field.metadata or {})
def _collect_blob_paths(schema: pa.Schema, is_blob) -> list[str]:
paths: list[str] = []
def _collect_blob_paths(schema: pa.Schema, is_blob) -> list[tuple[str, bool]]:
"""Walk the schema and return (path, has_list_ancestor) for each blob field."""
paths: list[tuple[str, bool]] = []
def walk(fields, prefix: str) -> None:
def walk(fields, prefix: str, has_list_ancestor: bool) -> None:
for field in fields:
path = f"{prefix}.{field.name}" if prefix else field.name
if is_blob(field):
paths.append(path)
paths.append((path, has_list_ancestor))
elif pa.types.is_struct(field.type):
walk(field.type, path)
walk(field.type, path, has_list_ancestor)
elif (
pa.types.is_list(field.type)
or pa.types.is_large_list(field.type)
or pa.types.is_fixed_size_list(field.type)
):
walk([field.type.value_field], path)
walk([field.type.value_field], path, True)
walk(schema, "")
walk(schema, "", False)
return paths
def blob_column_paths(schema: pa.Schema) -> list[str]:
"""Dotted paths of blob-like columns (v2 extension or legacy metadata)."""
return _collect_blob_paths(schema, is_blob_like_field)
return [path for path, _ in _collect_blob_paths(schema, is_blob_like_field)]
def blob_v2_column_paths(schema: pa.Schema) -> list[str]:
return _collect_blob_paths(schema, is_blob_v2_field)
return [path for path, _ in _collect_blob_paths(schema, is_blob_v2_field)]
def row_addressable_blob_v2_paths(schema: pa.Schema) -> list[str]:
"""Blob v2 paths with one blob addressable by table row id.
``fetch_blobs`` and the descriptor row-id ride-along address one blob per
row, so a blob inside a list container has no row-id slot and no fetch
path. Those columns still store and query as raw descriptors.
"""
return [
path
for path, has_list_ancestor in _collect_blob_paths(schema, is_blob_v2_field)
if not has_list_ancestor
]
def schema_has_blob_field(schema: pa.Schema) -> bool:
return bool(blob_column_paths(schema))
def _deserialize_registered_type(extension_type: pa.ExtensionType) -> pa.DataType:
"""Return the type Arrow reconstructs for this extension name."""
schema = pa.schema([pa.field("value", extension_type)])
restored = pa.ipc.read_schema(schema.serialize())
return restored.field("value").type
def _resolve_blob_type():
"""Return the BlobType class this process should use.
pylance's class when it owns the lance.blob.v2 registry entry,
otherwise LanceDB's fallback. A different registered class is an error.
"""
global _resolved_blob_type
if _resolved_blob_type is not None:
return _resolved_blob_type
try:
blob_module = importlib.import_module("lance.blob")
except ModuleNotFoundError as err:
if err.name not in ("lance", "lance.blob"):
raise
else:
blob_type = getattr(blob_module, "BlobType", None)
if blob_type is not None:
registered_type = _deserialize_registered_type(blob_type())
if type(registered_type) is not blob_type:
registered_cls = type(registered_type)
raise ValueError(
"lance.blob.v2 is already registered by "
f"{registered_cls.__module__}.{registered_cls.__qualname__}"
)
_resolved_blob_type = blob_type
return blob_type
try:
pa.register_extension_type(_FallbackBlobType()) # type: ignore[arg-type]
except pa.ArrowKeyError as err:
raise ValueError(
"lance.blob.v2 is already registered by another extension class"
) from err
_resolved_blob_type = _FallbackBlobType
return _resolved_blob_type
def blob(name: str, nullable: bool = True) -> pa.Field:
"""Create a Lance blob v2 column field."""
return pa.field(name, BlobType(), nullable=nullable)
"""Create a Lance blob v2 column field.
When pylance is installed this is ``lance.blob.BlobType``.
"""
blob_type = _resolve_blob_type()
return pa.field(name, blob_type(), nullable=nullable)
def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType:
@@ -155,3 +214,11 @@ def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataTyp
... ])
"""
return pa.list_(value_type, dimension)
def __getattr__(name: str):
if name == "BlobType":
blob_type = _resolve_blob_type()
globals()["BlobType"] = blob_type
return blob_type
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
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
+166
View File
@@ -0,0 +1,166 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Handles to SQL queries running on a remote database."""
from typing import Any, Mapping, Optional, Sequence, Union
from uuid import UUID
import numpy as np
import pyarrow as pa
from lancedb.background_loop import LOOP
from . import _lancedb
from .arrow import AsyncRecordBatchReader
QueryDescription = _lancedb.QueryDescription
QueryParameters = Union[pa.RecordBatch, pa.Table, Sequence[Any], Mapping[str, Any]]
"""Values for a statement's ``$1`` / ``$name`` placeholders.
A sequence binds by position (``$1``, ``$2``, ...), a mapping binds by name
(``$name``), and a one-row ``pyarrow.RecordBatch`` or ``pyarrow.Table`` binds
both ways: column ``i`` is ``$<i + 1>`` and, when its name is not a number,
also ``$<name>``.
"""
def _parameter_array(value: Any) -> pa.Array:
"""One parameter as a one-element array of its own type."""
if isinstance(value, pa.Scalar):
return pa.array([value.as_py()], type=value.type)
# A 1-D array is a vector: keep its element type and make it the
# fixed-size list a vector column stores, rather than a list of Python
# floats that would come back as float64.
if isinstance(value, pa.ChunkedArray):
value = value.combine_chunks()
if isinstance(value, np.ndarray):
if value.ndim == 0:
return pa.array(value.reshape(1))
if value.ndim != 1:
raise ValueError(
"a numpy query parameter must be a scalar or a 1-D vector, "
f"got {value.ndim} dimensions"
)
value = pa.array(value)
if isinstance(value, pa.Array):
return pa.FixedSizeListArray.from_arrays(value, len(value))
return pa.array([value])
def to_parameter_batch(
parameters: Optional[QueryParameters],
) -> Optional[pa.RecordBatch]:
"""Turn ``parameters`` into the one-row batch a parameterized query sends.
Scalars keep the type pyarrow infers for them (numpy scalars keep their
dtype), a ``pyarrow.Scalar`` keeps its own type, and a 1-D numpy or
pyarrow array becomes a ``FixedSizeList`` of its element type, which is
what a vector column stores.
"""
if parameters is None:
return None
batch: pa.RecordBatch
if isinstance(parameters, pa.RecordBatch):
batch = parameters
elif isinstance(parameters, pa.Table):
batches = parameters.combine_chunks().to_batches()
batch = (
batches[0]
if len(batches) == 1
else pa.RecordBatch.from_pylist([], schema=parameters.schema)
)
elif isinstance(parameters, Mapping):
batch = pa.RecordBatch.from_arrays(
[_parameter_array(value) for value in parameters.values()],
names=[str(name) for name in parameters.keys()],
)
elif isinstance(parameters, Sequence) and not isinstance(parameters, (str, bytes)):
batch = pa.RecordBatch.from_arrays(
[_parameter_array(value) for value in parameters],
names=[str(position + 1) for position in range(len(parameters))],
)
else:
raise TypeError(
"query parameters must be a sequence, a mapping, or a one-row "
f"pyarrow RecordBatch or Table, not {type(parameters).__name__}"
)
if batch.num_rows != 1:
raise ValueError(
f"query parameters must be exactly one row, got {batch.num_rows}"
)
return batch
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", "QueryParameters"]
python/python/lancedb/table.py
+823 -144
View File
File diff suppressed because it is too large Load Diff
python/python/lancedb/util.py
+28
View File
@@ -5,10 +5,12 @@
import binascii
import functools
import importlib
import math
import os
import pathlib
import warnings
from datetime import date, datetime
from decimal import Decimal
from functools import singledispatch
from typing import Tuple, Union, Optional, Any, List
from urllib.parse import urlparse
@@ -279,6 +281,14 @@ def infer_vector_column_dim(data_type: pa.DataType) -> Optional[int]:
return None
def _validate_query_vector(query: Any) -> None:
"""Reject empty vector inputs before vector-column inference or execution."""
if (isinstance(query, list) and not query) or (
isinstance(query, np.ndarray) and query.size == 0
):
raise ValueError("Query vector must not be empty")
def _query_vector_dim(query: Optional[Any]) -> Optional[int]:
if query is None:
return None
@@ -300,6 +310,9 @@ def infer_vector_column_name(
query: Optional[Any], # inferred later in query builder
vector_column_name: Optional[str],
):
if query_type != "fts":
_validate_query_vector(query)
if vector_column_name is not None:
return vector_column_name
@@ -351,9 +364,24 @@ def _(value: int):
@value_to_sql.register(float)
def _(value: float):
# str() gives "nan" / "inf", which SQL would read as column names.
if math.isnan(value):
return "CAST('NaN' AS DOUBLE)"
if math.isinf(value):
return (
"CAST('Infinity' AS DOUBLE)" if value > 0 else "CAST('-Infinity' AS DOUBLE)"
)
return str(value)
@value_to_sql.register(Decimal)
def _(value: Decimal):
if not value.is_finite():
raise ValueError("Non-finite Decimal values cannot be converted to SQL")
# Quote exact decimal text so the target-column cast avoids float rounding.
return value_to_sql(format(value, "f"))
@value_to_sql.register(bool)
def _(value: bool):
return str(value).upper()
python/python/lancedb/view.py
+40
View File
@@ -0,0 +1,40 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Views: a named query a database stores and plans on every read.
A view holds no rows. What it stores is the statement that defines it and the
schema that statement resolved to, so a reader sees the sources as they are
now. See ``DBConnection.create_view``.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, List
if TYPE_CHECKING:
import pyarrow as pa
@dataclass(frozen=True)
class ViewDescription:
"""What a database records about one view."""
name: str
"""The view's name within its namespace."""
query: str
"""The defining query, as the database stores it."""
default_database: str
"""The database that unqualified table names in ``query`` resolve against."""
schema: "pa.Schema"
"""The schema the defining query resolved to when the view was created."""
namespace_path: List[str] = field(default_factory=list)
"""The namespace holding the view; empty is the root namespace."""
default_namespace_path: List[str] = field(default_factory=list)
"""The namespace path those unqualified names resolve against.
Recorded with the view because it outlives the session that declared it: a
reader resolving the query against its own default namespace could read a
different table than the view was defined over.
"""
python/python/tests/test_blob.py
+805 -2
View File
@@ -2,10 +2,16 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import io
import subprocess
import sys
import textwrap
import lance
import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
import pytest
from lance.blob import Blob, BlobType as LanceBlobType
import lancedb
from lancedb._blob import (
@@ -18,6 +24,20 @@ from lancedb.index import FTS
from lancedb.schema import blob_column_paths, blob_v2_column_paths
_HIDE_LANCE_BLOB = """\
import importlib.abc
import sys
class _MissingLanceBlob(importlib.abc.MetaPathFinder):
def find_spec(self, fullname, path, target=None):
if fullname == "lance.blob" or fullname.startswith("lance.blob."):
raise ModuleNotFoundError(fullname, name="lance.blob")
sys.modules.pop("lance.blob", None)
sys.meta_path.insert(0, _MissingLanceBlob())
"""
def _blob_table(name, rows):
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
@@ -47,10 +67,204 @@ 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)
assert field.type.extension_name == "lance.blob.v2"
assert lancedb.BlobType is LanceBlobType
assert type(field.type) is LanceBlobType
def test_blob_type_works_without_pylance():
script = _HIDE_LANCE_BLOB + textwrap.dedent(
"""\
import lancedb
import pyarrow as pa
field = lancedb.blob("image")
if not isinstance(field.type, pa.ExtensionType):
raise SystemExit("expected an extension type")
if field.type.extension_name != "lance.blob.v2":
raise SystemExit(field.type.extension_name)
if lancedb.BlobType is not type(field.type):
raise SystemExit("BlobType is not the field type class")
if lancedb.BlobType.__module__ != "lancedb.schema":
raise SystemExit(lancedb.BlobType.__module__)
db = lancedb.connect("memory:///")
table = db.create_table(
"images",
schema=pa.schema([pa.field("id", pa.int64()), field]),
)
table.add([{"id": 1, "image": b"hello"}])
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute([{"id": 1, "image": b"updated"}, {"id": 2, "image": b"inserted"}])
)
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
raise SystemExit(
f"merge_insert rows updated={result.num_updated_rows} "
f"inserted={result.num_inserted_rows}"
)
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_blob_resolves_pylance_type_without_eager_import():
script = textwrap.dedent(
"""\
import sys
import lancedb
if "lance.blob" in sys.modules:
raise SystemExit("import lancedb imported lance.blob")
field = lancedb.blob("image")
from lance.blob import BlobType
if type(field.type) is not BlobType:
raise SystemExit(f"{type(field.type)} is not {BlobType}")
import lance
image = lance.blob_array([b"x"])
if type(image.type) is not BlobType:
raise SystemExit("blob_array used a different class")
if type(image.type) is not type(field.type):
raise SystemExit("field and array classes differ")
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_blob_fallback_fails_if_name_already_registered():
script = _HIDE_LANCE_BLOB + textwrap.dedent(
"""\
import pyarrow as pa
class OtherBlobType(pa.ExtensionType):
def __init__(self):
super().__init__(
pa.struct([pa.field("data", pa.large_binary())]),
"lance.blob.v2",
)
def __arrow_ext_serialize__(self):
return b""
@classmethod
def __arrow_ext_deserialize__(cls, storage_type, serialized):
return cls()
pa.register_extension_type(OtherBlobType())
import lancedb
try:
lancedb.blob("image")
except ValueError as err:
if "already registered" not in str(err):
raise SystemExit(err)
else:
raise SystemExit("expected ValueError")
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_blob_type_rejects_competing_registration_with_pylance():
script = textwrap.dedent(
"""\
import pyarrow as pa
import pyarrow.ipc
class OtherBlobType(pa.ExtensionType):
def __init__(self):
super().__init__(
pa.struct(
[
pa.field("data", pa.large_binary()),
pa.field("uri", pa.utf8()),
pa.field("position", pa.uint64()),
pa.field("size", pa.uint64()),
]
),
"lance.blob.v2",
)
def __arrow_ext_serialize__(self):
return b""
@classmethod
def __arrow_ext_deserialize__(cls, storage_type, serialized):
return cls()
pa.register_extension_type(OtherBlobType())
from lance.blob import BlobType
if BlobType is OtherBlobType:
raise SystemExit("pylance BlobType was replaced")
schema = pa.schema([pa.field("value", BlobType())])
restored = pa.ipc.read_schema(schema.serialize())
if type(restored.field("value").type) is not OtherBlobType:
raise SystemExit(type(restored.field("value").type))
import lancedb
try:
lancedb.blob("image")
except ValueError as err:
if "__main__.OtherBlobType" not in str(err):
raise SystemExit(err)
else:
raise SystemExit("expected ValueError")
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_blob_v2_column_paths_include_list_children():
@@ -84,7 +298,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()),
@@ -203,6 +420,292 @@ def test_fetch_blobs_round_trip():
assert [blobs[0].as_py(), blobs[1].as_py()] == [b"alpha", b"beta"]
def test_merge_insert_writes_python_bytes():
table = _blob_table("merge_bytes", [{"id": 1, "image": b"before"}])
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute([{"id": 1, "image": b"updated"}, {"id": 2, "image": b"inserted"}])
)
assert result.num_updated_rows == 1
assert result.num_inserted_rows == 1
by_id = _row_ids_by_id(table)
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
assert blobs.to_pylist() == [b"updated", b"inserted"]
def test_merge_insert_bytes_after_reopen_without_touching_blob_type(tmp_path):
db = lancedb.connect(tmp_path)
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("images", schema=schema)
table.add([{"id": 1, "image": b"hello"}])
script = textwrap.dedent(
f"""\
import lancedb
db = lancedb.connect({str(tmp_path)!r})
table = db.open_table("images")
image_type = table.schema.field("image").type
if type(image_type).__name__ != "StructType":
raise SystemExit(f"expected StructType, got {{type(image_type)}}")
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(
[{{"id": 1, "image": b"updated"}}, {{"id": 2, "image": b"inserted"}}]
)
)
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
raise SystemExit(
f"rows updated={{result.num_updated_rows}} "
f"inserted={{result.num_inserted_rows}}"
)
hits = table.search().with_row_id(True).limit(10).to_arrow()
by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
if blobs.to_pylist() != [b"updated", b"inserted"]:
raise SystemExit(blobs.to_pylist())
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_merge_insert_bytes_after_reopen_without_pylance(tmp_path):
db = lancedb.connect(tmp_path)
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("images", schema=schema)
table.add([{"id": 1, "image": b"hello"}])
script = _HIDE_LANCE_BLOB + textwrap.dedent(
f"""\
import lancedb
db = lancedb.connect({str(tmp_path)!r})
table = db.open_table("images")
image_type = table.schema.field("image").type
if type(image_type).__name__ != "StructType":
raise SystemExit(f"expected StructType, got {{type(image_type)}}")
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(
[{{"id": 1, "image": b"updated"}}, {{"id": 2, "image": b"inserted"}}]
)
)
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
raise SystemExit(
f"rows updated={{result.num_updated_rows}} "
f"inserted={{result.num_inserted_rows}}"
)
hits = table.search().with_row_id(True).limit(10).to_arrow()
by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
if blobs.to_pylist() != [b"updated", b"inserted"]:
raise SystemExit(blobs.to_pylist())
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_merge_insert_blob_array_into_reopened_unregistered_table(tmp_path):
db = lancedb.connect(tmp_path)
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("images", schema=schema)
table.add([{"id": 1, "image": b"before"}])
script = textwrap.dedent(
f"""\
import pyarrow as pa
import lancedb
db = lancedb.connect({str(tmp_path)!r})
table = db.open_table("images")
image_type = table.schema.field("image").type
if type(image_type).__name__ != "StructType":
raise SystemExit(
f"expected StructType before lance import, got {{type(image_type)}}"
)
import lance
updates = pa.Table.from_arrays(
[
pa.array([1, 2], type=pa.int64()),
lance.blob_array([b"updated", b"inserted"]),
],
names=["id", "image"],
)
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(updates)
)
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
raise SystemExit(
f"rows updated={{result.num_updated_rows}} "
f"inserted={{result.num_inserted_rows}}"
)
hits = table.search().with_row_id(True).limit(10).to_arrow()
by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
if blobs.to_pylist() != [b"updated", b"inserted"]:
raise SystemExit(blobs.to_pylist())
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_add_all_null_blob_column():
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("all_null", schema=schema)
table.add([{"id": 1, "image": None}, {"id": 2, "image": None}])
by_id = _row_ids_by_id(table)
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
assert blobs.to_pylist() == [None, None]
def test_create_table_nested_blob_schema_without_rows():
db = lancedb.connect("memory:///")
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("info", pa.struct([lancedb.blob("blob")])),
pa.field("images", pa.list_(lancedb.blob("image"))),
]
)
table = db.create_table("nested_empty", schema=schema)
assert table.count_rows() == 0
def test_merge_insert_nested_blob_dicts():
db = lancedb.connect("memory:///")
info = pa.StructArray.from_arrays(
[
pa.array(["first"], type=pa.string()),
_blob_array("blob", [b"before"]),
],
names=["name", "blob"],
)
data = pa.Table.from_arrays(
[pa.array([1], type=pa.int64()), info],
names=["id", "info"],
)
table = db.create_table("nested_merge", data=data)
result = (
table.merge_insert("id")
.when_matched_update_all()
.execute([{"id": 1, "info": {"name": "first", "blob": b"after"}}])
)
assert result.num_updated_rows == 1
by_id = _row_ids_by_id(table)
blobs = table.fetch_blobs("info.blob", [by_id[1]])
assert blobs.to_pylist() == [b"after"]
def _list_blob_table(name):
db = lancedb.connect("memory:///")
blob_field = lancedb.blob("image")
images = pa.ListArray.from_arrays(
pa.array([0, 1], type=pa.int32()), _blob_array("image", [b"before"])
)
data = pa.Table.from_arrays(
[pa.array([1], type=pa.int64()), images],
schema=pa.schema(
[pa.field("id", pa.int64()), pa.field("images", pa.list_(blob_field))]
),
)
return db.create_table(name, data=data)
def test_merge_insert_list_blob_dicts():
table = _list_blob_table("list_merge")
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute([{"id": 1, "images": [b"one", b"two"]}, {"id": 2, "images": None}])
)
assert result.num_updated_rows == 1
assert result.num_inserted_rows == 1
hits = table.search().limit(10).to_arrow()
sizes = {
row["id"]: None if row["images"] is None else [d["size"] for d in row["images"]]
for row in hits.to_pylist()
}
assert sizes == {1: [3, 3], 2: None}
def test_list_blob_column_queries_as_raw_descriptors():
table = _list_blob_table("list_query")
hits = table.search().limit(10).to_arrow()
element = hits.schema.field("images").type.value_type
assert pa.types.is_struct(element)
assert "_lance_row_id" not in element.names
with pytest.raises(ValueError, match="expected struct before segment"):
table.fetch_blobs("images.image", [0])
def test_row_addressable_paths_exclude_list_children():
from lancedb.schema import row_addressable_blob_v2_paths
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("info", pa.struct([lancedb.blob("blob")])),
pa.field("images", pa.list_(lancedb.blob("image"))),
]
)
assert blob_v2_column_paths(schema) == ["info.blob", "images.image"]
assert row_addressable_blob_v2_paths(schema) == ["info.blob"]
def test_merge_insert_writes_pylance_blob_array():
table = _blob_table("merge_pylance", [{"id": 1, "image": b"before"}])
image = lance.blob_array([b"updated", b"inserted"])
assert type(image.type) is LanceBlobType
assert type(image.type) is type(lancedb.BlobType())
updates = pa.Table.from_arrays(
[pa.array([1, 2], type=pa.int64()), image], names=["id", "image"]
)
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(updates)
)
assert result.num_updated_rows == 1
assert result.num_inserted_rows == 1
by_id = _row_ids_by_id(table)
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
assert blobs.to_pylist() == [b"updated", b"inserted"]
def test_fetch_blobs_accepts_query_result():
table = _blob_table("from_result", [{"id": 1, "image": b"gamma"}])
hits = table.search().limit(10).to_arrow()
@@ -211,6 +714,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",
@@ -230,6 +752,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",
@@ -259,8 +855,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():
@@ -309,6 +922,70 @@ def test_fetch_blob_files_lazy_read():
handles = table.fetch_blob_files("image", [by_id[1]])
assert len(handles) == 1
assert handles[0].read() == payload
handles[0].close()
assert handles[0].closed
with pytest.raises(RuntimeError, match="already closed"):
handles[0].read()
def test_blob_file_finalizer_and_close_on_runtime_worker():
# A tokio panic in IOBase.__del__ is printed but does not fail the add call.
# Run in a child process so the assertion catches that panic on stderr.
script = textwrap.dedent(
"""\
import gc
import threading
import lancedb
import pyarrow as pa
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("finalize", schema=schema)
table.add([{"id": 1, "image": b"first"}])
hits = table.search().with_row_id(True).limit(1).to_arrow()
row_id = hits["_rowid"][0].as_py()
gc_handle, explicit_handle = table.fetch_blob_files("image", [row_id, row_id])
gc.disable()
cycle = [gc_handle]
cycle.append(cycle)
del gc_handle, cycle
observed = []
def on_progress(_):
if not observed:
observed.append((
threading.current_thread() is not threading.main_thread(),
explicit_handle.closed,
gc.collect(),
))
explicit_handle.close()
table.add([{"id": 2, "image": b"second"}], progress=on_progress)
gc.enable()
assert len(observed) == 1, observed
assert observed[0][0] is True, observed
assert observed[0][1] is False, observed
assert observed[0][2] > 0, observed
assert explicit_handle.closed
explicit_handle.close() # Close remains idempotent after the worker call.
try:
explicit_handle.read()
except RuntimeError as error:
assert "already closed" in str(error)
else:
raise AssertionError("a closed blob file was readable")
print("worker close completed")
"""
)
result = subprocess.run(
[sys.executable, "-c", script], capture_output=True, text=True, check=False
)
assert result.returncode == 0, result.stderr
assert "Cannot start a runtime from within a runtime" not in result.stderr
assert "worker close completed" in result.stdout
def test_fetch_blob_files_null_alignment():
@@ -756,3 +1433,129 @@ def test_add_external_uri_string_round_trips_with_flag(tmp_path):
hits = table.search().to_arrow()
blobs = table.fetch_blobs("image", hits)
assert blobs[0].as_py() == payload
@pytest.mark.parametrize("as_pandas", [False, True])
def test_add_bytes_and_uri_in_one_batch_preserves_external_uri(tmp_path, as_pandas):
payload = b"external-payload"
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(payload)
rows = [
{"id": 1, "image": b"inline"},
{"id": 2, "image": blob_path.as_uri()},
]
db = lancedb.connect(tmp_path / "db")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("mixed", schema=schema)
table.add(
pd.DataFrame(rows) if as_pandas else rows,
allow_external_blob_outside_bases=True,
)
by_id = _row_ids_by_id(table)
assert table.fetch_blobs("image", [by_id[1], by_id[2]]).to_pylist() == [
b"inline",
payload,
]
@pytest.mark.parametrize("as_pandas", [False, True])
def test_add_mixed_python_blob_values_preserves_external_uris(tmp_path, as_pandas):
payload = b"external-payload"
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(payload)
uri = blob_path.as_uri()
rows = [
{"id": 1, "image": b"inline"},
{"id": 2, "image": uri},
{"id": 3, "image": {"data": b"from-dict"}},
{"id": 4, "image": Blob.from_bytes(b"from-blob")},
{"id": 5, "image": Blob.from_uri(uri)},
{"id": 6, "image": {"uri": uri}},
{"id": 7, "image": None},
]
db = lancedb.connect(tmp_path / "db")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("mixed", schema=schema)
table.add(
pd.DataFrame(rows) if as_pandas else rows,
allow_external_blob_outside_bases=True,
)
by_id = _row_ids_by_id(table)
assert table.fetch_blobs("image", [by_id[i] for i in range(1, 8)]).to_pylist() == [
b"inline",
payload,
b"from-dict",
b"from-blob",
payload,
payload,
None,
]
@pytest.mark.parametrize("as_pandas", [False, True])
def test_merge_insert_mixed_python_blobs_validates_external_uri(tmp_path, as_pandas):
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(b"external-payload")
table = _blob_table("merge_mixed_uri", [{"id": 1, "image": b"before"}])
rows = [
{"id": 1, "image": b"updated"},
{"id": 2, "image": blob_path.as_uri()},
]
with pytest.raises(ValueError, match="outside registered external bases"):
(
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(pd.DataFrame(rows) if as_pandas else rows)
)
assert table.count_rows() == 1
@pytest.mark.parametrize("as_pandas", [False, True])
def test_merge_insert_python_blob_dicts_and_objects(as_pandas):
table = _blob_table("merge_python_blobs", [{"id": 1, "image": b"before"}])
rows = [
{"id": 1, "image": b"updated"},
{"id": 2, "image": {"data": b"from-dict"}},
{"id": 3, "image": Blob.from_bytes(b"from-blob")},
]
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(pd.DataFrame(rows) if as_pandas else rows)
)
assert result.num_updated_rows == 1
assert result.num_inserted_rows == 2
by_id = _row_ids_by_id(table)
assert table.fetch_blobs("image", [by_id[i] for i in range(1, 4)]).to_pylist() == [
b"updated",
b"from-dict",
b"from-blob",
]
def test_create_table_with_python_blob_dict_and_object():
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table(
"create_python_blobs",
data=[
{"id": 1, "image": {"data": b"from-dict"}},
{"id": 2, "image": Blob.from_bytes(b"from-blob")},
],
schema=schema,
)
by_id = _row_ids_by_id(table)
assert table.fetch_blobs("image", [by_id[1], by_id[2]]).to_pylist() == [
b"from-dict",
b"from-blob",
]
python/python/tests/test_catalog.py
+265
View File
@@ -0,0 +1,265 @@
# 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"]}),
(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() == []
assert list(catalog.list_databases()) == ["team/search"]
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"
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)
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
@pytest.mark.parametrize("page_token,page_limit", [(None, None), ("start/token", 2)])
async def test_database_names_iteration(
catalog_server, asynchronous, page_token, page_limit
):
endpoint, requests, responses = catalog_server
responses.extend(
[
(200, {"namespaces": ["first", "second"], "page_token": "a/b"}),
(200, {"namespaces": [], "page_token": "next"}),
(200, {"namespaces": ["last", "final"], "page_token": ""}),
]
)
catalog = (
await lancedb.connect_catalog_async(endpoint)
if asynchronous
else lancedb.connect_catalog(endpoint)
)
names = catalog.list_databases(page_token=page_token, page_limit=page_limit)
assert requests == []
assert names.num_page_results() == 0
assert names.page_token() == page_token
for cached, expected in [(1, "first"), (0, "second")]:
assert (await names.__anext__() if asynchronous else next(names)) == expected
assert names.num_page_results() == cached
assert names.page_token() == "a/b"
assert len(requests) == 1
assert (await names.__anext__() if asynchronous else next(names)) == "last"
assert names.num_page_results() == 1
assert names.page_token() is None
assert (await names.__anext__() if asynchronous else next(names)) == "final"
assert names.num_page_results() == 0
assert names.page_token() is None
for _ in range(2):
with pytest.raises(StopAsyncIteration if asynchronous else StopIteration):
if asynchronous:
await names.__anext__()
else:
next(names)
from urllib.parse import parse_qs, urlsplit
for request, token in zip(requests, [page_token, "a/b", "next"]):
url = urlsplit(request[0])
assert url.path == "/v1/namespace/$/list"
expected = {}
if token is not None:
expected["page_token"] = [token]
if page_limit is not None:
expected["limit"] = [str(page_limit)]
assert parse_qs(url.query) == expected
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("fail", [False, True])
@pytest.mark.asyncio
async def test_database_names_empty_or_error(catalog_server, asynchronous, fail):
endpoint, requests, responses = catalog_server
responses.append(
(401, {"error": "unauthorized"}) if fail else (200, {"namespaces": []})
)
catalog = (
await lancedb.connect_catalog_async(endpoint)
if asynchronous
else lancedb.connect_catalog(endpoint)
)
names = catalog.list_databases()
if fail:
with pytest.raises(HttpError):
if asynchronous:
await names.__anext__()
else:
next(names)
assert ([name async for name in names] if asynchronous else list(names)) == []
assert len(requests) == 1
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_database_names_resume(catalog_server, asynchronous):
endpoint, requests, responses = catalog_server
responses.extend(
[
(200, {"namespaces": ["first"], "page_token": "resume/token"}),
(200, {"namespaces": ["second"]}),
]
)
catalog = (
await lancedb.connect_catalog_async(endpoint)
if asynchronous
else lancedb.connect_catalog(endpoint)
)
names = catalog.list_databases(page_limit=1)
assert (await names.__anext__() if asynchronous else next(names)) == "first"
assert names.num_page_results() == 0
resumed = catalog.list_databases(page_token=names.page_token(), page_limit=1)
assert resumed.page_token() == "resume/token"
assert len(requests) == 1
assert ([name async for name in resumed] if asynchronous else list(resumed)) == [
"second"
]
assert resumed.page_token() is None
assert resumed.num_page_results() == 0
assert requests[1][0] == "/v1/namespace/$/list?limit=1&page_token=resume%2Ftoken"
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("page_limit", [0, 2**32 - 1])
@pytest.mark.asyncio
async def test_database_names_invalid_limit(catalog_server, asynchronous, page_limit):
endpoint, requests, _ = catalog_server
catalog = (
await lancedb.connect_catalog_async(endpoint)
if asynchronous
else lancedb.connect_catalog(endpoint)
)
names = catalog.list_databases(page_limit=page_limit)
with pytest.raises(ValueError, match="limit"):
if asynchronous:
await names.__anext__()
else:
next(names)
assert names.num_page_results() == 0
assert requests == []
python/python/tests/test_db.py
+133
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import contextlib
import inspect
import re
import sys
@@ -244,6 +245,40 @@ def test_db_contains_and_len_include_all_table_name_pages(tmp_db: lancedb.DBConn
assert "does_not_exist" not in tmp_db
@pytest.mark.asyncio
async def test_table_listing_defaults_and_zero_limit(tmp_path):
names = [f"t{i:03}" for i in range(130)]
for name in names:
(tmp_path / f"{name}.lance").mkdir()
db = lancedb.connect(tmp_path)
assert db.table_names() == names
assert db.table_names(limit=None) == names
assert db.table_names(limit=0) == []
first = db.list_tables()
assert first.tables == names[:100]
assert first.page_token
second = db.list_tables(page_token=first.page_token)
assert second.tables == names[100:]
assert second.page_token is None
assert db.list_tables(limit=200).tables == names
empty = db.list_tables(limit=0)
assert empty.tables == []
assert empty.page_token is None
adb = await lancedb.connect_async(tmp_path)
assert await adb.table_names() == names
assert await adb.table_names(limit=0) == []
first = await adb.list_tables()
assert first.tables == names[:100]
second = await adb.list_tables(page_token=first.page_token)
assert second.tables == names[100:]
assert second.page_token is None
empty = await adb.list_tables(limit=0)
assert empty.tables == []
assert empty.page_token is None
def test_db_contains_stops_after_matching_table_page(
tmp_db: lancedb.DBConnection, monkeypatch
):
@@ -350,6 +385,40 @@ async def test_create_table_from_iterator_async(mem_db_async: lancedb.AsyncConne
assert await table.count_rows() == 10
def test_create_exist_ok_mode(tmp_db: lancedb.DBConnection):
db = tmp_db
data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}]
table = db.create_table("t1", data, mode="exist_ok")
assert table.count_rows() == 2
original = table.to_arrow()
reopened = db.create_table("t1", [{"vector": [9.0, 8.0]}], mode="exist_ok")
assert reopened.to_arrow().equals(original)
def test_create_empty_exist_ok_mode(tmp_db: lancedb.DBConnection):
schema = pa.schema([pa.field("id", pa.int64())])
table = tmp_db.create_table("empty", schema=schema, mode="exist_ok")
assert table.schema == schema
assert table.count_rows() == 0
table.add([{"id": 7}])
reopened = tmp_db.create_table("empty", schema=schema, mode="exist_ok")
assert reopened.to_arrow().to_pylist() == [{"id": 7}]
bad_schema = pa.schema([pa.field("other", pa.int64())])
with pytest.raises(ValueError):
tmp_db.create_table("empty", schema=bad_schema, mode="exist_ok")
def test_create_table_invalid_mode(tmp_db: lancedb.DBConnection):
db = tmp_db
data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}]
with pytest.raises(ValueError, match="Invalid mode bogus"):
db.create_table("tb", data, mode="bogus")
assert "tb" not in db
def test_create_exist_ok(tmp_db: lancedb.DBConnection):
from conftest import pandas_string_type
@@ -405,6 +474,49 @@ async def test_connect(tmp_path):
assert str(db) == f"ListingDatabase(uri={tmp_path}, read_consistency_interval=5s)"
@pytest.fixture(params=["memory", "local", "namespace", "from_inner"])
def sync_db(request, tmp_path):
if request.param == "memory":
return lancedb.connect("memory://")
if request.param == "namespace":
return lancedb.connect_namespace("dir", {"root": str(tmp_path)})
db = lancedb.connect(tmp_path)
if request.param == "from_inner":
return lancedb.db.LanceDBConnection.from_inner(db._inner, None)
return db
def test_sync_close(sync_db):
assert sync_db.is_open()
assert sync_db.close() is None
assert not sync_db.is_open()
assert sync_db.close() is None
with pytest.warns(DeprecationWarning, match="table_names"):
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.table_names()
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.list_tables()
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.open_table("test")
@pytest.mark.parametrize("raise_error", [False, True])
def test_sync_context_manager(sync_db, raise_error):
with contextlib.ExitStack() as stack:
if raise_error:
stack.enter_context(pytest.raises(ValueError, match="test error"))
with sync_db as db:
assert db is sync_db
assert db.is_open()
if raise_error:
raise ValueError("test error")
assert not sync_db.is_open()
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.list_tables()
@pytest.mark.asyncio
async def test_close(mem_db_async: lancedb.AsyncConnection):
assert mem_db_async.is_open()
@@ -495,6 +607,27 @@ async def test_create_exist_ok_async(tmp_db_async: lancedb.AsyncConnection):
# await db.create_table("test", schema=bad_schema, exist_ok=True)
@pytest.mark.parametrize("enable_v2_manifest_paths", [False, True])
@pytest.mark.parametrize("empty", [False, True])
def test_create_table_deprecated_v2_manifest_paths(
tmp_path, enable_v2_manifest_paths, empty
):
db = lancedb.connect(tmp_path)
data = None if empty else [{"id": 1}]
with pytest.warns(DeprecationWarning, match="enable_v2_manifest_paths"):
table = db.create_table(
"test",
data=data,
schema=pa.schema([("id", pa.int64())]),
enable_v2_manifest_paths=enable_v2_manifest_paths,
)
assert table.uses_v2_manifest_paths() == enable_v2_manifest_paths
assert table.to_arrow().to_pylist() == ([] if empty else data)
assert db.open_table("test").uses_v2_manifest_paths() == enable_v2_manifest_paths
@pytest.mark.asyncio
async def test_create_table_v2_manifest_paths_async(tmp_path):
db_with_v2_paths = await lancedb.connect_async(
python/python/tests/test_e2e_remote_db.py
+65
View File
@@ -2,13 +2,78 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import os
from datetime import timedelta
from uuid import uuid4
import lancedb
import numpy as np
import pyarrow as pa
import pytest
from lancedb import LanceDBConnection
from lancedb.index import FTS
# TODO: setup integ test mark and script
@pytest.mark.skipif(
not os.environ.get("LANCEDB_URI"), reason="Set LANCEDB_URI to a JSON FTS server"
)
def test_json_fts_against_server():
"""Requires a writable database and capacity to run the distributed indexer.
Configure LANCEDB_URI, LANCEDB_API_KEY, and optionally LANCEDB_REGION and
LANCEDB_HOST_OVERRIDE before running this test.
"""
uri = os.environ["LANCEDB_URI"]
assert uri.startswith("db://"), "LANCEDB_URI must identify a remote database"
conn = lancedb.connect(
uri,
api_key=os.environ.get("LANCEDB_API_KEY"),
region=os.environ.get("LANCEDB_REGION", "us-east-1"),
host_override=os.environ.get("LANCEDB_HOST_OVERRIDE"),
)
name = f"json_fts_{uuid4().hex}"
table = conn.create_table(
name,
pa.table(
{
"id": [0, 1, 2],
"doc": pa.array(
[
'{"meta":{"tag":"nature"},"author":"alice"}',
'{"meta":{"tag":"city"},"author":"nature"}',
'{"meta":{"tag":"nature"},"author":"bob"}',
],
type=pa.json_(),
),
"raw": pa.array([b"nature"] * 3, type=pa.large_binary()),
}
),
)
try:
table.create_index("doc", config=FTS(), name="doc_fts")
table.wait_for_index(["doc_fts"], timeout=timedelta(minutes=5))
stats = table.index_stats("doc_fts")
assert stats.num_indexed_rows == 3
assert stats.num_unindexed_rows == 0
assert stats.index_type == "FTS"
for query, expected in [
("meta.tag,str,nature", [0, 2]),
("meta.tag,str,alice", []),
("meta.tag,str,tag", []),
]:
search = table.search(query, query_type="fts", fts_columns="doc")
assert sorted(row["id"] for row in search.to_list()) == expected
plan = search.explain_plan(True)
assert "MatchQuery" in plan, plan
assert "FlatMatchQuery" not in plan, plan
with pytest.raises(ValueError, match="FTS.*lance.json.*raw binary"):
table.create_index("raw", config=FTS())
finally:
conn.drop_table(name)
@pytest.mark.skip(reason="Need to set up a local server")
def test_against_local_server():
conn = LanceDBConnection("lancedb+http://localhost:10024")
python/python/tests/test_embeddings.py
+17 -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()
@@ -633,6 +633,18 @@ def test_url_retrieve_downloads_image():
assert img.size[0] > 0 and img.size[1] > 0
def test_open_clip_opens_percent_encoded_file_uri(tmp_path):
"""OpenCLIP should decode local file URIs before opening them."""
Image = pytest.importorskip("PIL.Image")
from lancedb.embeddings.open_clip import OpenClipEmbeddings
image_path = tmp_path / "test image.png"
Image.new("RGB", (4, 4), color="red").save(image_path, format="PNG")
with OpenClipEmbeddings._to_pil(None, image_path.as_uri()) as image:
assert image.size == (4, 4)
def test_jina_generate_image_input_dict_local_path(tmp_path):
"""
JinaEmbeddings._generate_image_input_dict must accept a local image path
python/python/tests/test_first_class_function_slice1.py
+140 -13
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()
@@ -65,6 +83,18 @@ def test_public_function_values_are_in_api_reference():
FunctionBinding,
False,
),
(
"remote_initialized_function_application.json",
"remote_initialized_function_application.canonical.json",
FunctionApplication,
False,
),
(
"remote_initialized_function_binding.json",
"remote_initialized_function_binding.canonical.json",
FunctionBinding,
False,
),
(
"remote_refresh_job.json",
"remote_refresh_result.canonical.json",
@@ -93,16 +123,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():
@@ -121,6 +162,29 @@ def test_function_version_binds_named_columns_as_one_immutable_application():
] == [("text", "column", "documents.body")]
def test_function_version_names_its_namespace_as_a_separate_field():
"""A Function is named by its parts, as a Secret is; a root one states no
namespace."""
value = job_result("remote_function_job.json")
root = FunctionVersion.from_json(json.dumps(value))
assert root.namespace_path == ()
assert "namespace_path" not in json.loads(root.to_canonical_json())
nested = FunctionVersion(**{**value, "namespace_path": ["analytics", "features"]})
assert nested.name == "embed"
assert nested.namespace_path == ("analytics", "features")
assert json.loads(nested.to_canonical_json())["namespace_path"] == [
"analytics",
"features",
]
application = nested(text=col("documents.body"))
assert application.function.name == "embed"
assert application.function.namespace_path == ("analytics", "features")
assert json.loads(application.to_canonical_json())["function"][
"namespace_path"
] == ["analytics", "features"]
def test_function_version_binding_validates_names_and_direct_columns():
version = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
@@ -128,7 +192,7 @@ def test_function_version_binding_validates_names_and_direct_columns():
with pytest.raises(TypeError, match=r"missing inputs: \['text'\]"):
version()
with pytest.raises(TypeError, match=r"unknown inputs: \['body'\]"):
with pytest.raises(TypeError, match=r"unknown arguments: \['body'\]"):
version(text=col("text"), body=col("body"))
with pytest.raises(TypeError, match="direct col"):
version(text=col("text").lower())
@@ -137,7 +201,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 +246,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 +286,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
@@ -240,6 +304,35 @@ def test_binding_and_refresh_result_keep_stable_remote_fields():
assert RefreshColumnResult.from_json(result.to_canonical_json()) == result
def test_initialization_travels_on_the_application_and_binding():
application = FunctionApplication.from_json(
fixture("remote_initialized_function_application.json")
)
assert application.initialization["temperature"] == 0.25
assert application.initialization["retry"]["backoff"] == (1, 2, 4)
application._ensure_declarable()
binding = FunctionBinding.from_json(
fixture("remote_initialized_function_binding.json")
)
assert binding.initialization_row() == {
"model": "text-embedding-3-small",
"dimensions": 512,
}
assert (
FunctionBinding.from_json(
fixture("remote_function_binding.json")
).initialization_row()
is None
)
signature = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
).signature
assert signature.initialization == ()
assert "initialization" not in json.loads(signature.to_canonical_json())
def test_function_literal_numeric_domain_matches_rust():
with pytest.raises(ValueError, match="floating-point Function literals"):
FunctionApplication.from_json(fixture("remote_function_application_float.json"))
@@ -276,6 +369,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 +457,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
+1352 -40
View File
File diff suppressed because it is too large Load Diff
python/python/tests/test_fts.py
+188 -5
View File
@@ -13,6 +13,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os
import random
import shutil
@@ -44,6 +45,75 @@ from utils import exception_output
TEST_LANGUAGE_MODEL_HOME = Path(__file__).parent / "models"
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
@pytest.mark.parametrize("block_size", [128, 256])
def test_json_fts(tmp_path, block_size):
db = ldb.connect(tmp_path)
data = pa.table(
{
"id": [0, 1, 2, 3],
"doc": pa.array(
[
'{"title":"brown fox","meta":{"tag":"nature"},"author":"alice"}',
'{"title":"blue car","meta":{"tag":"city"}}',
'{"title":"brown bear","meta":{"tag":"nature"}}',
'{"title":"red fox","author":"alice"}',
],
type=pa.json_(),
),
}
)
table = db.create_table("json_fts", data)
table.create_index("doc", config=FTS(block_size=block_size), name="doc_fts")
table = db.open_table("json_fts")
index = table.list_indices()[0]
assert index.name == "doc_fts"
assert index.index_type == "FTS"
assert index.columns == ["doc"]
stats = table.index_stats("doc_fts")
assert stats.num_indexed_rows == 4
assert stats.num_unindexed_rows == 0
for query, expected in [
("title,str,brown", [0, 2]),
("meta.tag,str,nature", [0, 2]),
("author,str,alice", [0, 3]),
("title,str,alice", []),
("title,str,title", []),
]:
search = table.search(query, query_type="fts", fts_columns="doc")
assert sorted(row["id"] for row in search.to_list()) == expected
plan = search.explain_plan(True)
assert "MatchQuery" in plan, plan
assert "FlatMatchQuery" not in plan, plan
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
@pytest.mark.asyncio
async def test_json_fts_async(tmp_path):
db = await ldb.connect_async(tmp_path)
table = await db.create_table(
"json_fts",
pa.table({"doc": pa.array(['{"meta":{"tag":"nature"}}'], type=pa.json_())}),
)
await table.create_index("doc", config=FTS())
search = table.query().nearest_to_text(MatchQuery("meta.tag,str,nature", "doc"))
assert len(await search.to_list()) == 1
plan = await search.explain_plan(True)
assert "MatchQuery" in plan, plan
assert "FlatMatchQuery" not in plan, plan
@pytest.mark.parametrize("dtype", [pa.binary(), pa.large_binary()])
def test_fts_rejects_raw_binary(tmp_path, dtype):
table = ldb.connect(tmp_path).create_table(
"binary_fts", pa.table({"doc": pa.array([b'{"title":"brown"}'], type=dtype)})
)
with pytest.raises(ValueError, match="FTS.*lance.json.*raw binary"):
table.create_index("doc", config=FTS())
assert table.list_indices() == []
@pytest.fixture
def table(tmp_path) -> ldb.table.LanceTable:
# Use local random state to avoid affecting other tests
@@ -894,6 +964,78 @@ def test_tokenize_uses_simple_index_tokenizer(mem_db: DBConnection):
]
@pytest.mark.parametrize("legacy", [False, True])
@pytest.mark.parametrize(
("base_tokenizer", "options", "expected_filters", "expected_tokens"),
[
(
"code",
{},
(False, False),
[("running", 0), ("in", 1), ("cafe", 2)],
),
(
"code",
{"stem": None, "remove_stop_words": None},
(False, False),
[("running", 0), ("in", 1), ("cafe", 2)],
),
(
"code",
{"stem": True},
(True, False),
[("run", 0), ("in", 1), ("cafe", 2)],
),
(
"code",
{"remove_stop_words": True},
(False, True),
[("running", 0), ("cafe", 2)],
),
(
"code",
{"stem": True, "remove_stop_words": True},
(True, True),
[("run", 0), ("cafe", 2)],
),
(
"simple",
{},
(True, True),
[("run", 0), ("cafe", 2)],
),
(
"simple",
{"stem": False, "remove_stop_words": False},
(False, False),
[("running", 0), ("in", 1), ("cafe", 2)],
),
],
)
def test_fts_tokenizer_filter_defaults(
mem_db, legacy, base_tokenizer, options, expected_filters, expected_tokens
):
text = "Running in café"
table = mem_db.create_table("test_filter_defaults", data=[{"text": text}])
if legacy:
with pytest.warns(DeprecationWarning, match="create_fts_index"):
table.create_fts_index("text", base_tokenizer=base_tokenizer, **options)
else:
table.create_index("text", config=FTS(base_tokenizer=base_tokenizer, **options))
details = table.list_indices()[0].index_details
assert (details["stem"], details["remove_stop_words"]) == expected_filters
assert [
(token.text, token.position) for token in table.tokenize(text, column="text")
] == expected_tokens
assert [
(token.text, token.position)
for token in ldb.tokenize(text, base_tokenizer=base_tokenizer, **options)
] == expected_tokens
if base_tokenizer == "code" and not expected_filters[1]:
assert table.search("in", query_type="fts").to_list()[0]["text"] == text
def test_tokenize_uses_icu_index_tokenizer_by_name(mem_db: DBConnection):
data = pa.table({"text": ["Hello, こんにちは世界!"]})
table = mem_db.create_table("test_tokenize_icu", data=data)
@@ -1011,8 +1153,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 +1181,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()
@@ -1127,6 +1279,34 @@ def test_fts_lindera_tokenizer(
assert [row["text"] for row in results] == ["成田国際空港"]
@pytest.mark.parametrize("fuzziness", [None, 0, 1, 2])
def test_match_query_fuzziness_to_json(fuzziness):
query = MatchQuery("puppy", "text", fuzziness=fuzziness)
assert query.fuzziness == fuzziness
assert json.loads(query.to_json())["match"]["fuzziness"] == fuzziness
@pytest.mark.parametrize(
"options, expected",
[
pytest.param({}, ["ab"], id="default-exact"),
pytest.param({"fuzziness": 0}, ["ab"], id="explicit-exact"),
pytest.param({"fuzziness": None}, ["ab"], id="automatic"),
pytest.param({"fuzziness": 1}, ["ab", "ac"], id="one-edit"),
pytest.param({"fuzziness": 2}, ["ab", "ac", "zz"], id="two-edits"),
],
)
def test_match_query_fuzziness_search(tmp_db, options, expected):
# Automatic fuzziness is zero for two-character terms.
table = tmp_db.create_table("fuzzy", pa.table({"text": ["ab", "ac", "zz"]}))
table.create_index("text", config=FTS(stem=False, remove_stop_words=False))
results = table.search(MatchQuery("ab", "text", **options)).to_list()
assert sorted(row["text"] for row in results) == sorted(expected)
def test_fts_query_to_json():
"""Test that FTS query to_json() produces valid JSON strings with exact format."""
@@ -1257,7 +1437,10 @@ def test_fts_query_to_json():
json_str = multi_match.to_json()
expected = (
'{"multi_match":{"query":"python","columns":["tags","title"],'
'"boost":[1.0,1.0]}}'
'"boost":[1.0,1.0],"match_queries":[{"column":"tags","terms":"python",'
'"boost":1.0,"fuzziness":0,"max_expansions":50,"operator":"Or",'
'"prefix_length":0},{"column":"title","terms":"python","boost":1.0,'
'"fuzziness":0,"max_expansions":50,"operator":"Or","prefix_length":0}]}}'
)
assert json_str == expected
python/python/tests/test_function_packaging.py
+527
View File
@@ -0,0 +1,527 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Class Functions, initialization, and the source that travels with a Function."""
from __future__ import annotations
import base64
import importlib
import json
import subprocess
import sys
import textwrap
import warnings
from pathlib import Path
from typing import Annotated, Optional
import pyarrow as pa
import pytest
from lancedb import col
from lancedb.functions import FunctionVersion, udf
FIXTURES = (
Path(__file__).parents[3]
/ "rust"
/ "lancedb"
/ "tests"
/ "fixtures"
/ "first_class_functions"
/ "v1"
)
def artifact(definition) -> bytes:
return base64.b64decode(definition.registration_request.artifact.content.data)
def bundle_files(definition) -> dict[str, str]:
assert definition.registration_request.artifact.kind == "python_bundle"
return json.loads(artifact(definition))["files"]
@pytest.fixture
def helper_package(tmp_path, monkeypatch):
"""A local package on sys.path, the way a project's own helpers are."""
def write(name: str, files: dict[str, str]):
root = tmp_path / name
for path, source in files.items():
target = root / path
target.parent.mkdir(parents=True, exist_ok=True)
target.write_text(textwrap.dedent(source))
monkeypatch.syspath_prepend(str(tmp_path))
for module in [
module for module in sys.modules if module.split(".")[0] == name
]:
del sys.modules[module]
importlib.invalidate_caches()
return importlib.import_module(name)
return write
@udf
class scale:
def __init__(
self,
factor: float,
label: str,
offset: float = 0.0,
tags: Optional[list[str]] = None,
retry: Annotated[dict, pa.struct([("attempts", pa.int32())])] = None,
):
self.factor, self.offset = factor, offset
def __call__(self, value: float) -> float:
return value * self.factor + self.offset
class NoCall:
pass
class BadClose:
def __call__(self, value: int) -> int:
return value
def close(self, force):
pass
def test_class_function_declares_inputs_output_and_initialization():
signature = scale.registration_request.signature
assert [(p.name, p.arrow_type) for p in signature.inputs] == [("value", "float64")]
assert [
(field.name, field.arrow_type, field.nullable)
for field in signature.initialization
] == [
("factor", "float64", False),
("label", "utf8", False),
("offset", "float64", True),
("tags", "list<utf8>", True),
(
"retry",
'{"fields":[{"name":"attempts","nullable":true,"type":{"type":"int32"}}],'
'"type":"struct"}',
True,
),
]
# Calling the definition constructs the class, so local tests are ordinary.
assert scale(factor=2.0, label="x")(1.5) == 3.0
assert scale.registration_request.artifact.entrypoint == "scale"
source = artifact(scale).decode()
assert source.startswith("from __future__ import annotations\n")
assert "class scale:" in source and "@udf" not in source
def test_function_without_initialization_keeps_its_wire_form():
@udf
def double(value: int) -> int:
return value * 2
request = json.loads(double.registration_request.to_canonical_json())
assert "initialization" not in request["signature"]
assert request["artifact"]["kind"] == "python_callable"
@pytest.mark.parametrize(
("init", "message"),
[
("def __init__(self, blob: bytes): pass", "initialization takes booleans"),
("def __init__(self, *args: int): pass", "named and non-variadic"),
("def __init__(self, factor): pass", "missing Function initialization"),
("def __init__(self, value: int): pass", "share names \\['value'\\]"),
],
)
def test_class_function_rejects_unsupported_initialization(init, message):
namespace = {}
exec(
textwrap.dedent(
f"""
class bad:
{init}
def __call__(self, value: int) -> int:
return value
"""
),
namespace,
)
with pytest.raises((TypeError, ValueError), match=message):
udf(namespace["bad"])
def test_class_function_requires_call_and_a_plain_close():
with pytest.raises(TypeError, match="synchronous __call__"):
udf(NoCall)
with pytest.raises(TypeError, match="close must be a method taking no"):
udf(BadClose)
def test_nested_class_and_closure_are_rejected_with_the_reason():
def build():
class nested:
def __call__(self, value: int) -> int:
return value
return nested
with pytest.raises(ValueError, match="defined at module level"):
udf(build())
factor = 3
with pytest.raises(ValueError, match="capture closure values"):
@udf
def closes_over(value: int) -> int:
return value * factor
def test_binding_passes_initialization_as_constants():
version = FunctionVersion.from_json(
json.dumps(
{
**json.loads(
(FIXTURES / "remote_function_version.canonical.json").read_text()
),
"signature": {
"inputs": [
{"name": "text", "arrow_type": "utf8", "nullable": True}
],
"output": {
"kind": "scalar",
"arrow_type": "list<float32>",
"nullable": False,
},
"initialization": [
{"name": "model", "arrow_type": "utf8", "nullable": False},
{"name": "dimensions", "arrow_type": "int32", "nullable": True},
],
},
}
)
)
application = version(text=col("body"), model="small", dimensions=(512))
assert dict(application.initialization) == {"model": "small", "dimensions": 512}
assert json.loads(application.to_canonical_json())["initialization"] == {
"dimensions": 512,
"model": "small",
}
assert version(text=col("body"), model="small").initialization == {"model": "small"}
with pytest.raises(TypeError, match=r"missing initialization: \['model'\]"):
version(text=col("body"))
with pytest.raises(TypeError, match="takes a constant"):
version(text=col("body"), model=col("model"))
with pytest.raises(ValueError, match="must be finite"):
version(text=col("body"), model="small", dimensions=float("nan"))
with pytest.raises(TypeError, match="takes booleans"):
version(text=col("body"), model=object())
def test_code_packages_ship_as_a_bundle_and_version_with_their_source(
tmp_path, helper_package
):
files = {
"__init__.py": "",
"rate.py": """
class TokenBucket:
def __init__(self, tokens_per_minute):
self.tokens_per_minute = tokens_per_minute
""",
"__pycache__/rate.cpython-312.pyc": "",
"fixtures/data.json": "{}",
"not-a-module.py": "",
}
helper_package("fx_helpers", files)
(tmp_path / "fx_helpers_functions.py").write_text(
textwrap.dedent(
"""
import fx_helpers
from fx_helpers.rate import TokenBucket
from lancedb import udf
@udf(code=[fx_helpers])
class limited:
def __init__(self, tpm: int):
self.limiter = TokenBucket(tpm)
def __call__(self, value: int) -> int:
return value
"""
)
)
def define():
sys.modules.pop("fx_helpers_functions", None)
return importlib.import_module("fx_helpers_functions").limited
first = define()
shipped = bundle_files(first)
assert sorted(shipped) == [
"fx_helpers/__init__.py",
"fx_helpers/rate.py",
"user_function.py",
]
assert (
"from fx_helpers.rate import TokenBucket as TokenBucket"
in (shipped["user_function.py"])
)
# Only the helper changes: a new artifact digest, the same environment.
helper_package(
"fx_helpers", {**files, "rate.py": textwrap.dedent(files["rate.py"]) + "# v2\n"}
)
second = define()
assert (
first.registration_request.artifact.digest
!= second.registration_request.artifact.digest
)
assert first.registration_request.runtime == second.registration_request.runtime
def test_code_rejects_submodules_duplicates_and_non_modules(helper_package):
helpers = helper_package("fx_rejects", {"__init__.py": "", "sub.py": ""})
import fx_rejects.sub
for code, message in [
([fx_rejects.sub], "top-level modules"),
([helpers, helpers], "more than once"),
(["fx_rejects"], "imported modules"),
([json], "cannot ship module 'json'"),
]:
with pytest.raises((TypeError, ValueError), match=message):
@udf(code=code)
def f(value: int) -> int:
return value
def test_local_module_outside_code_warns_at_registration(tmp_path, helper_package):
"""The worker most likely cannot import a local module that is neither
shipped nor declared, but registration still succeeds: a declared
distribution may install it under another name, and callers may rewrite
the packaged source before registering it."""
helper_package(
"fx_local", {"__init__.py": "def normalize(value):\n return value\n"}
)
def define(name, options):
(tmp_path / f"{name}.py").write_text(
textwrap.dedent(
f"""
from fx_local import normalize
from lancedb import udf
@udf({options})
def {name}(value: int) -> int:
return normalize(value)
"""
)
)
return getattr(importlib.import_module(name), name)
with pytest.warns(UserWarning, match=r"local module.*code=\[fx_local\]"):
local = define("fx_uses_local", "")
assert "from fx_local import normalize" in artifact(local).decode()
# A declared package that provides the module is trusted to install it.
with warnings.catch_warnings():
warnings.simplefilter("error")
installed = define("fx_uses_installed", "pip=['fx-local==1.0']")
assert "from fx_local import normalize" in artifact(installed).decode()
MMLB_HELPERS = {
"__init__.py": "",
"rate.py": """
class TokenBucket:
constructed = 0
def __init__(self, tokens_per_minute):
TokenBucket.constructed += 1
self.tokens_per_minute = tokens_per_minute
self.acquired = 0
def acquire(self, tokens):
self.acquired += tokens
""",
"embedding.py": """
from fx_mmlb.rate import TokenBucket
class OpenAIEmbedder:
def __init__(self, model: str, tpm_limit: int, max_attempts: int = 6):
self.model = model
self.limiter = TokenBucket(tpm_limit)
self.max_attempts = max_attempts
def embed(self, texts):
self.limiter.acquire(sum(len(text) for text in texts))
return [[float(len(text))] for text in texts]
""",
}
MMLB_FUNCTIONS = """
import pyarrow as pa
import fx_mmlb
from fx_mmlb.embedding import OpenAIEmbedder
from lancedb import udf
TEXT = pa.schema([pa.field("text", pa.string())])
EMBEDDING = pa.list_(pa.field("item", pa.float64(), nullable=False))
@udf(code=[fx_mmlb], input_schema=TEXT, output_schema=EMBEDDING)
class openai_embedding_batch(OpenAIEmbedder):
def __call__(self, text: pa.Array) -> pa.Array:
item = pa.field("item", pa.float64(), nullable=False)
return pa.array(self.embed(text.to_pylist()), pa.list_(item))
@udf(code=[fx_mmlb])
class openai_embedding(OpenAIEmbedder):
def __call__(self, text: str) -> list[float]:
return self.embed([text])[0]
"""
def test_shared_helpers_construct_state_on_every_path(tmp_path, helper_package):
"""The MMLB rate limiter, shared by a row and a batch Function: both are
initialized through the same helper, so neither path can skip it."""
helper_package("fx_mmlb", MMLB_HELPERS)
(tmp_path / "fx_mmlb_functions.py").write_text(MMLB_FUNCTIONS)
functions = importlib.import_module("fx_mmlb_functions")
signatures = [
functions.openai_embedding_batch.registration_request.signature,
functions.openai_embedding.registration_request.signature,
]
for signature in signatures:
assert [field.name for field in signature.initialization] == [
"model",
"tpm_limit",
"max_attempts",
]
# Each bundle imports on its own, in a process that has never seen the
# client's modules, the way the worker's image code layer does.
for definition, call in [
(
functions.openai_embedding_batch,
"instance(pa.array(['ab', 'abc'])).to_pylist()",
),
(functions.openai_embedding, "[instance('ab'), instance('abc')]"),
]:
worker = tmp_path / definition.__name__
for path, source in bundle_files(definition).items():
(worker / path).parent.mkdir(parents=True, exist_ok=True)
(worker / path).write_text(source)
probe = textwrap.dedent(
f"""
import sys
sys.path.insert(0, {str(worker)!r})
import pyarrow as pa
import user_function
from fx_mmlb.rate import TokenBucket
cls = getattr(user_function, {definition.__name__!r})
instance = cls(model="small", tpm_limit=1000)
print(repr(({call}, TokenBucket.constructed, instance.limiter.acquired)))
"""
)
result = subprocess.run(
[sys.executable, "-I", "-c", probe],
capture_output=True,
text=True,
cwd=worker,
)
assert result.returncode == 0, result.stderr
assert result.stdout.strip() == "([[2.0], [3.0]], 1, 5)"
def test_main_definitions_travel_by_source(tmp_path):
"""Notebook cells and scripts define helpers in __main__, which a worker
cannot import; they are packaged by source, dependencies first."""
script = tmp_path / "notebook.py"
script.write_text(
textwrap.dedent(
"""
import json
import math
from dataclasses import dataclass
from lancedb import udf
SCALE = 2
@dataclass
class Settings:
factor: int
def helper(value):
return math.floor(value) * SCALE
class Base:
def run(self, value):
return helper(value) * self.settings.factor
@udf
class scaled(Base):
def __init__(self, factor: int):
self.settings = Settings(factor)
def __call__(self, value: float) -> int:
return self.run(value)
print(json.dumps(scaled.registration_request.to_canonical_json()))
"""
)
)
result = subprocess.run(
[sys.executable, str(script)], capture_output=True, text=True
)
assert result.returncode == 0, result.stderr
request = json.loads(json.loads(result.stdout))
assert request["artifact"]["kind"] == "python_callable"
source = base64.b64decode(request["artifact"]["content"]["data"]).decode()
assert source.index("def helper") < source.index("class Base")
assert source.index("class Base") < source.index("class scaled")
assert source.index("class Settings") < source.index("class scaled")
assert "@dataclass" in source
assert "SCALE = 2" in source
namespace = {}
exec(compile(source, "user_function.py", "exec"), namespace)
assert namespace["scaled"](factor=3)(2.7) == 12
def test_main_lambdas_are_rejected(tmp_path):
script = tmp_path / "lambda.py"
script.write_text(
textwrap.dedent(
"""
from lancedb import udf
double = lambda value: value * 2
@udf
def uses_lambda(value: int) -> int:
return double(value)
"""
)
)
result = subprocess.run(
[sys.executable, str(script)], capture_output=True, text=True
)
assert result.returncode != 0
assert "not lambdas, closures, or nested definitions" in result.stderr
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
+128
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.
@@ -228,6 +315,33 @@ def test_hybrid_query_offset(sync_table: Table):
assert offset_result["_rowid"].to_pylist() == full["_rowid"].to_pylist()[2:]
def test_hybrid_query_default_limit(sync_table: Table):
# The vector and FTS legs match disjoint rows, so fusing their results
# yields more rows than the default limit.
new_rows = []
for i in range(20):
new_rows.append({"text": "close_vec", "vector": [0.1, 0.1]})
new_rows.append({"text": "dog", "vector": [50.0 + i, 50.0 + i]})
sync_table.add(new_rows)
def query():
return (
sync_table.search(query_type="hybrid")
.vector([0.1, 0.1])
.text("dog")
.with_row_id(True)
)
# Like the async hybrid query, the default limit is 10.
result = query().to_arrow()
assert len(result) == 10
# The offset window is taken after the default limit is applied.
full = query().limit(15).to_arrow()
offset_result = query().offset(5).to_arrow()
assert offset_result["_rowid"].to_pylist() == full["_rowid"].to_pylist()[5:]
def test_hybrid_query_minimum_nprobes_zero_raises(sync_table: Table):
# minimum_nprobes(0) must raise the same validation error a plain vector
# query raises, not silently no-op because 0 is falsy.
@@ -241,6 +355,20 @@ def test_hybrid_query_minimum_nprobes_zero_raises(sync_table: Table):
)
def test_hybrid_query_nprobes_remains_independent(sync_table: Table):
query = sync_table.search(query_type="hybrid").nprobes(20)
assert query._nprobes == 20
assert query._minimum_nprobes is None
assert query._maximum_nprobes is None
query = query.minimum_nprobes(5).nprobes(20)
assert query._nprobes == 20
assert query._minimum_nprobes == 5
assert query._maximum_nprobes is None
def test_hybrid_query_distance_range(sync_table: Table):
reranker = RRFReranker(return_score="all")
result = (
python/python/tests/test_listing.py
+194
View File
@@ -0,0 +1,194 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import asyncio
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import json
from pathlib import Path
from threading import Thread
from urllib.parse import parse_qs, urlsplit
import lancedb
from lancedb.remote.errors import HttpError
import pytest
@pytest.fixture
def listing_server():
requests, responses = [], []
class Handler(BaseHTTPRequestHandler):
def handle_request(self):
body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
requests.append((self.command, self.path, json.loads(body) if body else {}))
status, response = responses.pop(0)
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.end_headers()
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 response_page(kind, names, token):
if kind == "secrets":
items = [{"name": name} for name in names]
elif kind == "jobs":
items = [
{
"job_id": name,
"table": "t",
"job_type": "create_index",
"state": "in_progress",
"created_at_millis": 1,
}
for name in names
]
elif kind == "functions":
fixture = (
Path(__file__).resolve().parents[3]
/ "rust/lancedb/tests/fixtures/first_class_functions/v1"
/ "remote_function_version.canonical.json"
)
definition = json.loads(fixture.read_text())
items = [
{"name": "embed", "version": "1", "definition": definition} for _ in names
]
else:
items = names
key = "views" if kind == "materialized_views" else kind
return {key: items, "page_token": token}
async def connection(endpoint, asynchronous):
kwargs = dict(
api_key="fake",
host_override=endpoint,
client_config={"retry_config": {"retries": 0}},
)
if asynchronous:
return await lancedb.connect_async("db://dev", **kwargs)
return lancedb.connect("db://dev", **kwargs)
async def advance(items, asynchronous):
return await items.__anext__() if asynchronous else next(items)
KINDS = ["secrets", "views", "jobs", "functions", "materialized_views"]
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", KINDS)
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_listing_lazy_pages_and_resume(listing_server, kind, asynchronous):
endpoint, requests, responses = listing_server
responses.extend(
(200, response_page(kind, names, token))
for names, token in [(["a", "b"], "empty"), ([], "last"), (["c", "d"], "")]
)
db = await connection(endpoint, asynchronous)
method = getattr(db, "list_" + kind)
kwargs = (
{} if kind in ("jobs", "materialized_views") else {"namespace_path": ["team"]}
)
items = method(page_token="start/token", page_limit=2, **kwargs)
assert requests == []
assert items.page_token() == "start/token"
assert items.num_page_results() == 0
first = await advance(items, asynchronous)
if kind == "functions":
assert first.name == "embed"
assert first.version == "1"
elif kind == "jobs":
assert first.job_id == "a"
assert first.state == "running"
else:
assert first == "a"
assert items.num_page_results() == 1
assert items.page_token() == "empty"
await advance(items, asynchronous)
assert items.num_page_results() == 0
assert len(requests) == 1
resumed = method(page_token=items.page_token(), page_limit=2, **kwargs)
await advance(resumed, asynchronous)
assert len(requests) == 3
assert resumed.page_token() is None
assert resumed.num_page_results() == 1
remaining = [item async for item in resumed] if asynchronous else list(resumed)
assert len(remaining) == 1
assert resumed.num_page_results() == 0
assert ([item async for item in resumed] if asynchronous else list(resumed)) == []
for (verb, url, body), token in zip(requests, ["start/token", "empty", "last"]):
path = urlsplit(url)
if kind == "jobs":
assert verb == "POST"
assert path.path == "/v1/jobs/list"
assert body == {"page_token": token, "limit": 2}
else:
namespace = "$" if kind == "materialized_views" else "team"
resource = kind.removesuffix("s")
assert path.path == f"/v1/namespace/{namespace}/{resource}/list"
expected = {"page_token": [token], "limit": ["2"]}
if kind == "functions":
expected["include_definition"] = ["true"]
assert parse_qs(path.query) == expected
assert responses == []
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", KINDS)
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_listing_error_retains_token(listing_server, kind, asynchronous):
endpoint, requests, responses = listing_server
responses.extend([(401, {}), (200, response_page(kind, ["a"], None))])
db = await connection(endpoint, asynchronous)
method = getattr(db, "list_" + kind)
items = method(page_token="retry")
with pytest.raises(HttpError):
await advance(items, asynchronous)
assert items.page_token() == "retry"
assert items.num_page_results() == 0
assert ([item async for item in items] if asynchronous else list(items)) == []
assert len(requests) == 1
resumed = method(page_token=items.page_token())
assert len([item async for item in resumed] if asynchronous else list(resumed)) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", KINDS)
async def test_listing_serializes_concurrent_advances(listing_server, kind):
endpoint, requests, responses = listing_server
responses.append((200, response_page(kind, ["a", "b"], None)))
db = await connection(endpoint, True)
items = getattr(db, "list_" + kind)()
results = await asyncio.gather(
items.__anext__(), items.__anext__(), items.__anext__(), return_exceptions=True
)
assert (
len([result for result in results if isinstance(result, StopAsyncIteration)])
== 1
)
assert len(requests) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", KINDS)
@pytest.mark.parametrize("page_limit", [0, 2**32 - 1])
async def test_listing_invalid_limit_does_not_fetch(listing_server, kind, page_limit):
endpoint, requests, _ = listing_server
db = await connection(endpoint, True)
items = getattr(db, "list_" + kind)(page_limit=page_limit)
with pytest.raises(ValueError, match="limit"):
await items.__anext__()
assert requests == []
python/python/tests/test_lsm_write_spec.py
+7 -7
View File
@@ -52,9 +52,9 @@ def test_set_lsm_write_spec_validates(tmp_path):
with pytest.raises(Exception, match="num_buckets"):
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 1025))
# Happy path then mutation rejected.
# Happy path, then a set over the installed spec is refused.
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 4))
with pytest.raises(Exception, match="mutation"):
with pytest.raises(Exception, match="already set"):
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 8))
@@ -176,18 +176,18 @@ def test_get_lsm_write_spec(tmp_path):
table.unset_lsm_write_spec()
assert table.get_lsm_write_spec() is None
# Identity round-trips (column recovered from the schema). Leaving the
# maintained set to be inferred picks up the index on the table, so the
# spec reads back naming it rather than as "infer".
# Identity round-trips (column recovered from the schema). An unnamed
# maintained set reads back as None: what is stored is the intent to
# maintain every index, not the set it resolves to now.
table.set_lsm_write_spec(LsmWriteSpec.identity("id"))
spec = table.get_lsm_write_spec()
assert spec.spec_type == "identity"
assert spec.column == "id"
assert spec.maintained_indexes == [idx_name]
assert spec.maintained_indexes is None
table.unset_lsm_write_spec()
# Unsharded round-trips (no routing column). Opting out is distinct from
# the inferred default.
# the unnamed default.
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
spec = table.get_lsm_write_spec()
assert spec.spec_type == "unsharded"
python/python/tests/test_materialized_views.py
+334 -50
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,134 @@ 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 list(db.list_materialized_views()) == ["daily_sales"]
assert requests == ["/v1/namespace/$/materialized_view/list"]
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,52 +165,67 @@ 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 list(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")
view = db.open_materialized_view("adults")
assert view.definition == MaterializedViewDefinition(
source_table="people",
projections=[("name", "`name`"), ("age", "`age`")],
filter="age >= 18",
inputs=["age", "name"],
query="SELECT name, age FROM people WHERE age >= 18"
)
def test_incremental_refresh_after_append(tmp_path):
def test_full_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}])
result = view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
assert result.mode == "rebuild"
assert result.rows_written == 4
assert view.table.count_rows() == 4
assert view.refresh().mode == "no_op"
assert view.refresh().mode == "rebuild"
def test_incremental_refresh_after_update(tmp_path):
def test_full_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})
result = view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
assert result.mode == "rebuild"
assert result.rows_written == 3
rows = view.table.search().to_list()
assert sorted(row["age"] for row in rows) == [8, 36, 85]
@@ -87,7 +236,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})
@@ -97,13 +246,43 @@ def test_legacy_storage_source_update_rebuilds(tmp_path):
assert sorted(row["age"] for row in rows) == [8, 36]
def test_direct_writes_to_a_view_are_refused(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view("copy", "people")
table = view.table
row = [{"name": "eve", "age": 1}]
refused = pytest.raises(NotImplementedError, match="is a materialized view")
with refused:
table.add(row)
with refused:
table.update(where="name = 'kid'", values={"age": 8})
with refused:
table.delete("age < 18")
with refused:
table.merge_insert("name").when_not_matched_insert_all().execute(row)
with refused:
table.add_columns({"older": "age + 1"})
with refused:
table.drop_columns(["age"])
assert table.count_rows() == 3
db.open_table("people").add(row)
assert view.refresh().rows_written == 4
def test_list_and_not_a_view(tmp_path):
db = make_db(tmp_path)
db.create_materialized_view("adults", "people", where="age >= 18")
assert db.list_materialized_views() == ["adults"]
assert list(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 list(db.list_materialized_views()) == []
def test_invalid_expression_fails_at_create(tmp_path):
@@ -119,7 +298,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"
@@ -127,29 +309,53 @@ async def test_async_create_refresh_and_open(tmp_path):
reopened = await db.open_materialized_view("shouts")
definition = await reopened.definition()
assert definition.projections == [("shout", "upper(name)")]
assert await db.list_materialized_views() == ["shouts"]
assert definition.query == "SELECT upper(name) AS shout FROM people"
assert [item async for item in db.list_materialized_views()] == ["shouts"]
@pytest.mark.asyncio
async def test_async_incremental(tmp_path):
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}])
view = await db.create_materialized_view("copy", "people")
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 [item async for item in db.list_materialized_views()]
@pytest.mark.asyncio
async def test_async_refresh_rebuilds(tmp_path):
db = await lancedb.connect_async(tmp_path)
await db.create_table("people", [{"name": "ada", "age": 36}])
view = await db.create_materialized_view("copy", "people", with_no_data=True)
await view.refresh()
table = await db.open_table("people")
await table.add([{"name": "alan", "age": 41}])
result = await view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
assert result.mode == "rebuild"
assert result.rows_written == 2
def test_source_requires_stable_row_ids(tmp_path):
def test_source_does_not_require_stable_row_ids(tmp_path):
db = lancedb.connect(tmp_path)
db.create_table("plain", [{"x": 1}])
with pytest.raises(Exception, match="stable row ids"):
db.create_materialized_view("v", "plain")
view = db.create_materialized_view("v", "plain")
assert view.table.count_rows() == 1
def test_bare_select_names_are_quoted(tmp_path):
@@ -157,7 +363,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,25 +375,12 @@ 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")
view.refresh()
rows = view.table.search().to_list()
assert set(rows[0]) - {"__source_row_id"} == {"name"}
assert set(rows[0]) == {"name"}
assert sorted(row["name"] for row in rows) == ["ada", "grace", "kid"]
@@ -195,7 +391,7 @@ async def test_async_scalar_select_is_one_column(tmp_path):
view = await db.create_materialized_view("just_name", "people", select="name")
await view.refresh()
rows = await view.table.query().to_list()
assert set(rows[0]) - {"__source_row_id"} == {"name"}
assert set(rows[0]) == {"name"}
def test_limit_above_i64_max_is_refused(tmp_path):
@@ -228,13 +424,24 @@ def test_namespace_connection_materialized_views(tmp_path):
view = db.create_materialized_view("adults", "people", where="age >= 18")
view.refresh()
assert view.table.count_rows() == 1
assert db.list_materialized_views() == ["adults"]
assert list(db.list_materialized_views()) == ["adults"]
reopened = db.open_materialized_view("adults")
assert reopened.definition.source_table == "people"
assert reopened.definition.query.startswith("SELECT name, age FROM ")
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 list(db.list_materialized_views()) == []
@pytest.mark.asyncio
async def test_async_namespace_connection_materialized_views(tmp_path):
@@ -252,10 +459,10 @@ async def test_async_namespace_connection_materialized_views(tmp_path):
view = await db.create_materialized_view("adults", "people", where="age >= 18")
await view.refresh()
assert await view.table.count_rows() == 1
assert await db.list_materialized_views() == ["adults"]
assert [item async for item in db.list_materialized_views()] == ["adults"]
reopened = await db.open_materialized_view("adults")
assert (await reopened.definition()).source_table == "people"
assert (await reopened.definition()).query.startswith("SELECT name, age FROM ")
# The view's table came through the namespace, not straight from the
# inner connection: a bare inner table carries no namespace context, so
@@ -266,3 +473,80 @@ 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 [item async for item in db.list_materialized_views()] == []
def test_stored_queries_and_legacy_layouts_are_read():
import json
import pyarrow as pa
from lancedb.materialized_view import _definition_from_schema
def read(definition: dict) -> MaterializedViewDefinition:
schema = pa.schema([pa.field("id", pa.int32())]).with_metadata(
{b"mv.definition": json.dumps(definition).encode()}
)
return _definition_from_schema(schema, "v")
query = "SELECT id, c.chunk FROM ns.docs, UNNEST(chunks) AS c WHERE id > 1"
assert read({"format": 1, "query": query}).query == query
# The structured layout written before the format number reads as the
# query it described, under either of its kind tags.
assert (
read(
{
"kind": "namespaced_select",
"source_table": "people",
"source_namespace": ["ns"],
"projections": [
{"output": "name", "expression": "`name`"},
{"output": "Shout", "expression": "upper(name)"},
],
"filter": "age >= 18",
"limit": 10,
}
).query
== "SELECT `name`, upper(name) AS `Shout` FROM ns.people "
"WHERE age >= 18 LIMIT 10"
)
assert read({"kind": "select", "source_table": "people"}).query == (
"SELECT * FROM people"
)
# A newer writer's layout is reported, never guessed at.
for newer in (
{"format": 3, "query": query},
{"kind": "select_v3", "source_table": "people"},
):
with pytest.raises(NotImplementedError, match="cannot refresh"):
read(newer)
def test_local_materialized_view_listing_pages_and_resume(tmp_path):
db = make_db(tmp_path)
db.create_materialized_view("a_view", "people", with_no_data=True)
db.create_materialized_view("z_view", "people", with_no_data=True)
db.create_table("ordinary", [{"id": 1}])
names = db.list_materialized_views(page_limit=1)
assert names.num_page_results() == 0
first = next(names)
assert names.num_page_results() == 0
assert names.page_token() is not None
resumed = db.list_materialized_views(page_token=names.page_token(), page_limit=1)
assert sorted([first, *resumed]) == ["a_view", "z_view"]
assert resumed.page_token() is None
assert list(resumed) == []
python/python/tests/test_namespace.py
+199 -2
View File
@@ -6,11 +6,17 @@
import tempfile
import shutil
import importlib
import http.server
import json
import threading
from urllib.parse import parse_qs, unquote, urlparse
import pytest
import pyarrow as pa
import lancedb
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
from lancedb.namespace import _MAX_QUERY_K
from lancedb.listing import AsyncListing, Listing
from lancedb.namespace import _MAX_QUERY_K, _query_to_namespace_request
from lancedb.query import Query
from lancedb.table import AsyncTable, LanceTable
@@ -72,6 +78,132 @@ def _namespace_lance_table(namespace_client: _NamespaceClient) -> LanceTable:
return table
@pytest.mark.asyncio
@pytest.mark.parametrize("namespace_path", [[], ["test_ns"]])
async def test_namespace_table_names_defaults_and_limits(tmp_path, namespace_path):
db = lancedb.connect_namespace("dir", {"root": str(tmp_path)})
if namespace_path:
db.create_namespace(namespace_path)
names = [f"t{idx:03}" for idx in range(12)]
schema = pa.schema([pa.field("id", pa.int64())])
for name in names:
db.create_table(name, schema=schema, namespace_path=namespace_path)
assert list(db.table_names(namespace_path=namespace_path)) == names
assert list(db.table_names(limit=None, namespace_path=namespace_path)) == names
assert list(db.table_names(limit=0, namespace_path=namespace_path)) == []
assert list(db.table_names(limit=2, namespace_path=namespace_path)) == names[:2]
assert (
list(
db.table_names(page_token=names[1], limit=2, namespace_path=namespace_path)
)
== names[2:4]
)
adb = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)})
assert list(await adb.table_names(namespace_path=namespace_path)) == names
assert (
list(await adb.table_names(limit=None, namespace_path=namespace_path)) == names
)
assert list(await adb.table_names(limit=0, namespace_path=namespace_path)) == []
assert (
list(await adb.table_names(limit=2, namespace_path=namespace_path)) == names[:2]
)
assert (
list(
await adb.table_names(
page_token=names[1], limit=2, namespace_path=namespace_path
)
)
== names[2:4]
)
@pytest.fixture
def paginated_rest_namespace():
names = [f"t{idx:03}" for idx in reversed(range(12))]
requests = []
class Handler(http.server.BaseHTTPRequestHandler):
def do_GET(self):
url = urlparse(self.path)
query = parse_qs(url.query)
requests.append((unquote(url.path), query))
token = query.get("page_token", [None])[0]
if token not in {None, "opaque-next"}:
self.send_error(400, "Unknown page token")
return
start = 0 if token is None else 10
body = {"tables": names[start : start + 10]}
if start == 0:
body["page_token"] = "opaque-next"
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.end_headers()
self.wfile.write(json.dumps(body).encode())
def log_message(self, *args):
pass
with http.server.HTTPServer(("127.0.0.1", 0), Handler) as server:
thread = threading.Thread(target=server.serve_forever)
thread.start()
try:
yield {"uri": f"http://127.0.0.1:{server.server_port}"}, requests
finally:
server.shutdown()
thread.join()
@pytest.mark.asyncio
@pytest.mark.parametrize("namespace_path", [[], ["n"]])
async def test_rest_namespace_table_names_collects_pages(
paginated_rest_namespace, namespace_path
):
properties, requests = paginated_rest_namespace
names = [f"t{idx:03}" for idx in range(12)]
db = lancedb.connect_namespace("rest", properties)
first = db.list_tables(namespace_path=namespace_path)
second = db.list_tables(namespace_path=namespace_path, page_token=first.page_token)
assert len(first.tables) == 10
assert first.page_token == "opaque-next"
assert sorted(first.tables + second.tables) == names
assert second.page_token is None
assert list(db.table_names(namespace_path=namespace_path)) == names
assert list(db.table_names(limit=None, namespace_path=namespace_path)) == names
assert list(db.table_names(limit=11, namespace_path=namespace_path)) == names[:11]
assert (
list(db.table_names(page_token="t005", limit=2, namespace_path=namespace_path))
== names[6:8]
)
adb = lancedb.connect_namespace_async("rest", properties)
assert list(await adb.table_names(namespace_path=namespace_path)) == names
assert (
list(await adb.table_names(limit=None, namespace_path=namespace_path)) == names
)
assert (
list(await adb.table_names(limit=11, namespace_path=namespace_path))
== names[:11]
)
assert (
list(
await adb.table_names(
page_token="t005", limit=2, namespace_path=namespace_path
)
)
== names[6:8]
)
request_count = len(requests)
assert list(db.table_names(limit=0, namespace_path=namespace_path)) == []
assert list(await adb.table_names(limit=0, namespace_path=namespace_path)) == []
assert len(requests) == request_count
identifier = "$" if not namespace_path else "n"
assert all(path == f"/v1/namespace/{identifier}/table/list" for path, _ in requests)
class TestNamespaceConnection:
"""Test namespace-based LanceDB connection using DirectoryNamespace."""
@@ -104,6 +236,18 @@ class TestNamespaceConnection:
assert isinstance(db, lancedb.LanceNamespaceDBConnection)
assert len(list(db.table_names())) == 0
def test_list_materialized_views(self):
db = lancedb.connect_namespace("dir", {"root": self.temp_dir})
db.create_table("source", [{"id": 1}])
db.create_materialized_view("view", "source", with_no_data=True)
names = db.list_materialized_views(page_limit=1)
assert isinstance(names, Listing)
assert names.num_page_results() == 0
assert list(names) == ["view"]
assert names.num_page_results() == 0
assert names.page_token() is None
def test_sync_builtin_namespace_uses_rust_without_python_client(self, monkeypatch):
"""Built-in sync namespace connections should not construct or call the
Python namespace client for normal namespace/table management."""
@@ -193,7 +337,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]
@@ -599,6 +749,18 @@ class TestAsyncNamespaceConnection:
table_names = await db.table_names()
assert len(list(table_names)) == 0
async def test_list_materialized_views(self):
db = lancedb.connect_namespace_async("dir", {"root": self.temp_dir})
await db.create_table("source", [{"id": 1}])
await db.create_materialized_view("view", "source", with_no_data=True)
names = db.list_materialized_views(page_limit=1)
assert isinstance(names, AsyncListing)
assert names.num_page_results() == 0
assert [name async for name in names] == ["view"]
assert names.num_page_results() == 0
assert names.page_token() is None
async def test_async_builtin_namespace_uses_rust_without_python_client(
self, monkeypatch
):
@@ -950,6 +1112,29 @@ class TestPushdownOperations:
]
assert all(r.k <= 2**31 - 1 for r in namespace_client.requests)
def test_probe_fields_pass_through_independently(self):
request = _query_to_namespace_request(
["geneva", "hist"],
Query(
vector=[1.0, 2.0],
nprobes=20,
minimum_nprobes=3,
maximum_nprobes=10,
),
)
assert request.nprobes == 20
assert request.minimum_nprobes == 3
assert request.maximum_nprobes == 10
request = _query_to_namespace_request(
["geneva", "hist"],
Query(vector=[1.0, 2.0], nprobes=5, maximum_nprobes=0),
)
assert request.nprobes == 5
assert request.maximum_nprobes == 0
assert request.to_dict()["maximum_nprobes"] == 0
@pytest.mark.asyncio
class TestAsyncPushdownOperations:
@@ -1024,3 +1209,15 @@ def test_local_table_to_arrow_and_to_pandas_are_unchanged(tmp_path):
assert table.to_arrow().column("id").to_pylist() == [1, 2]
assert table.to_pandas()["id"].tolist() == [1, 2]
@pytest.mark.parametrize("effort", [0.0, 0.5, 1.0])
def test_namespace_request_rejects_search_effort(effort):
from lancedb.query import Query
with pytest.raises(
ValueError, match="namespace QueryTable does not support search_effort"
):
_query_to_namespace_request(
["test"], Query(vector=[1.0, 0.0], search_effort=effort)
)
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_pydantic.py
+20
View File
@@ -9,6 +9,7 @@ from typing import List, Optional, Tuple
import pyarrow as pa
import pydantic
import pytest
from lancedb.conftest import MockTextEmbeddingFunction
from lancedb.pydantic import (
LanceModel,
MultiVector,
@@ -425,6 +426,25 @@ def test_bare_vector_raises_clear_error():
exec("class TestModel(LanceModel):\n vector: Vector", namespace)
def test_embedding_vector_list_annotation():
embedding = MockTextEmbeddingFunction.create()
class StaticTypingModel(LanceModel):
text: str = embedding.SourceField()
vector: list[float] = embedding.VectorField()
schema = pydantic_to_schema(StaticTypingModel)
assert schema == pa.schema(
[
pa.field("text", pa.utf8(), False),
pa.field("vector", pa.list_(pa.float32(), embedding.ndims()), True),
]
)
model = StaticTypingModel(text="hello", vector=[0.0] * embedding.ndims())
assert model.vector == [0.0] * embedding.ndims()
def test_fixed_size_list_field():
class TestModel(pydantic.BaseModel):
vec: Vector(16)
python/python/tests/test_query.py
+318 -38
View File
@@ -24,6 +24,7 @@ from lancedb.pydantic import LanceModel, Vector
from lancedb.query import (
AsyncFTSQuery,
AsyncHybridQuery,
AsyncQuery,
AsyncQueryBase,
AsyncVectorQuery,
ColumnOrdering,
@@ -40,6 +41,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 +124,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 +284,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 +333,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 +360,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 +380,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 +419,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 +473,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 +497,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 +512,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"):
@@ -946,6 +975,19 @@ def test_query_builder_with_filter(table):
assert all(np.array(rs[0]["vector"]) == [3, 4])
@pytest.mark.parametrize("query_type", ["vector", "hybrid"])
@pytest.mark.parametrize("nprobes", [0, -1])
def test_nprobes_nonpositive_sync(table, query_type, nprobes):
if query_type == "hybrid":
query = table.search(query_type="hybrid").vector([0, 0]).text("a")
else:
query = table.search([0, 0])
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(nprobes)
def test_invalid_nprobes_sync(table):
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
LanceVectorQueryBuilder(table, [0, 0], "vector").minimum_nprobes(0).to_list()
@@ -953,12 +995,22 @@ def test_invalid_nprobes_sync(table):
ValueError,
match="maximum_nprobes must be greater than or equal to minimum_nprobes",
):
LanceVectorQueryBuilder(table, [0, 0], "vector").maximum_nprobes(5).to_list()
(
LanceVectorQueryBuilder(table, [0, 0], "vector")
.minimum_nprobes(10)
.maximum_nprobes(5)
.to_list()
)
with pytest.raises(
ValueError,
match="minimum_nprobes must be less than or equal to maximum_nprobes",
match="maximum_nprobes must be greater than or equal to minimum_nprobes",
):
LanceVectorQueryBuilder(table, [0, 0], "vector").minimum_nprobes(100).to_list()
(
LanceVectorQueryBuilder(table, [0, 0], "vector")
.maximum_nprobes(5)
.minimum_nprobes(100)
.to_list()
)
def test_nprobes_works_sync(table):
@@ -977,6 +1029,18 @@ def test_multiple_nprobes_calls_works_sync(table):
).minimum_nprobes(20).to_list()
@pytest.mark.asyncio
@pytest.mark.parametrize("hybrid", [False, True])
async def test_nprobes_zero_async(table_async: AsyncTable, hybrid):
query = table_async.query().nearest_to([0, 0])
if hybrid:
query = query.nearest_to_text("dog")
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(0)
@pytest.mark.asyncio
async def test_invalid_nprobes_async(table_async: AsyncTable):
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
@@ -985,12 +1049,22 @@ async def test_invalid_nprobes_async(table_async: AsyncTable):
ValueError,
match="maximum_nprobes must be greater than or equal to minimum_nprobes",
):
await table_async.vector_search([0, 0]).maximum_nprobes(5).to_list()
await (
table_async.vector_search([0, 0])
.minimum_nprobes(10)
.maximum_nprobes(5)
.to_list()
)
with pytest.raises(
ValueError,
match="minimum_nprobes must be less than or equal to maximum_nprobes",
):
await table_async.vector_search([0, 0]).minimum_nprobes(100).to_list()
await (
table_async.vector_search([0, 0])
.maximum_nprobes(5)
.minimum_nprobes(100)
.to_list()
)
def test_query_builder_with_prefilter(table):
@@ -1251,6 +1325,27 @@ async def test_query_to_polars_async(table_async: AsyncTable):
assert df.shape == (0, num_columns)
@pytest.mark.asyncio
@pytest.mark.parametrize("search_kwargs", [{}, {"query": None}])
@pytest.mark.parametrize("with_vector", [False, True])
async def test_async_search_without_query(mem_db_async, search_kwargs, with_vector):
data = [{"id": i} for i in range(10)]
if with_vector:
for row in data:
row["vector"] = [float(row["id"]), 1.0]
table = await mem_db_async.create_table("test", data)
query = await table.search(**search_kwargs)
assert isinstance(query, AsyncQuery)
assert await query.limit(3).to_arrow() == await table.query().limit(3).to_arrow()
query = await table.search(**search_kwargs)
assert await query.where("id >= 8").select(["id"]).limit(3).to_list() == [
{"id": 8},
{"id": 9},
]
@pytest.mark.asyncio
async def test_none_query(table_async: AsyncTable):
with pytest.raises(ValueError):
@@ -1560,8 +1655,9 @@ def test_query_serialization_sync(table: lancedb.table.Table):
q,
vector_column="vector",
vector=[5.0, 6.0],
minimum_nprobes=10,
maximum_nprobes=10,
nprobes=10,
minimum_nprobes=None,
maximum_nprobes=None,
refine_factor=5,
)
@@ -1579,8 +1675,19 @@ def test_query_serialization_sync(table: lancedb.table.Table):
q,
vector_column="vector",
vector=[5.0, 6.0],
minimum_nprobes=50,
maximum_nprobes=50,
nprobes=50,
minimum_nprobes=None,
maximum_nprobes=None,
)
q = table.search([5.0, 6.0]).minimum_nprobes(5).nprobes(50).to_query_object()
check_set_props(
q,
vector_column="vector",
vector=[5.0, 6.0],
nprobes=50,
minimum_nprobes=5,
maximum_nprobes=None,
)
q = table.search([5.0, 6.0]).maximum_nprobes(10).to_query_object()
@@ -1592,6 +1699,15 @@ def test_query_serialization_sync(table: lancedb.table.Table):
minimum_nprobes=None,
)
q = table.search([5.0, 6.0]).nprobes(5).maximum_nprobes(0).to_query_object()
check_set_props(
q,
vector_column="vector",
vector=[5.0, 6.0],
nprobes=5,
maximum_nprobes=0,
)
q = table.search([5.0, 6.0]).distance_range(0.0, 1.0).to_query_object()
check_set_props(
q, vector_column="vector", vector=[5.0, 6.0], lower_bound=0.0, upper_bound=1.0
@@ -1641,8 +1757,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
limit=10,
vector=sample_vector,
postfilter=False,
minimum_nprobes=20,
maximum_nprobes=20,
minimum_nprobes=None,
maximum_nprobes=None,
with_row_id=False,
bypass_vector_index=False,
)
@@ -1652,8 +1768,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
q,
vector=sample_vector,
postfilter=False,
minimum_nprobes=20,
maximum_nprobes=20,
minimum_nprobes=None,
maximum_nprobes=None,
with_row_id=False,
bypass_vector_index=False,
limit=10,
@@ -1664,8 +1780,9 @@ async def test_query_serialization_async(table_async: AsyncTable):
q,
vector=sample_vector,
postfilter=False,
minimum_nprobes=50,
maximum_nprobes=50,
nprobes=50,
minimum_nprobes=None,
maximum_nprobes=None,
with_row_id=False,
bypass_vector_index=False,
limit=10,
@@ -1684,8 +1801,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
filter="id = 1",
postfilter=True,
vector=sample_vector,
minimum_nprobes=20,
maximum_nprobes=20,
minimum_nprobes=None,
maximum_nprobes=None,
with_row_id=False,
bypass_vector_index=False,
)
@@ -1699,8 +1816,9 @@ async def test_query_serialization_async(table_async: AsyncTable):
check_set_props(
q,
vector=sample_vector,
minimum_nprobes=10,
maximum_nprobes=10,
nprobes=10,
minimum_nprobes=None,
maximum_nprobes=None,
refine_factor=5,
postfilter=False,
with_row_id=False,
@@ -1713,7 +1831,24 @@ async def test_query_serialization_async(table_async: AsyncTable):
q,
vector=sample_vector,
minimum_nprobes=5,
maximum_nprobes=20,
maximum_nprobes=None,
postfilter=False,
with_row_id=False,
bypass_vector_index=False,
limit=10,
)
q = (
(await table_async.search([5.0, 6.0]))
.nprobes(5)
.maximum_nprobes(0)
.to_query_object()
)
check_set_props(
q,
vector=sample_vector,
nprobes=5,
maximum_nprobes=0,
postfilter=False,
with_row_id=False,
bypass_vector_index=False,
@@ -1731,8 +1866,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
lower_bound=0.0,
upper_bound=1.0,
postfilter=False,
minimum_nprobes=20,
maximum_nprobes=20,
minimum_nprobes=None,
maximum_nprobes=None,
with_row_id=False,
bypass_vector_index=False,
limit=10,
@@ -1744,8 +1879,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
distance_type="cosine",
vector=sample_vector,
postfilter=False,
minimum_nprobes=20,
maximum_nprobes=20,
minimum_nprobes=None,
maximum_nprobes=None,
with_row_id=False,
bypass_vector_index=False,
limit=10,
@@ -1757,8 +1892,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
ef=7,
vector=sample_vector,
postfilter=False,
minimum_nprobes=20,
maximum_nprobes=20,
minimum_nprobes=None,
maximum_nprobes=None,
with_row_id=False,
bypass_vector_index=False,
limit=10,
@@ -1770,8 +1905,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
bypass_vector_index=True,
vector=sample_vector,
postfilter=False,
minimum_nprobes=20,
maximum_nprobes=20,
minimum_nprobes=None,
maximum_nprobes=None,
with_row_id=False,
limit=10,
)
@@ -1923,6 +2058,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())
@@ -2048,6 +2198,12 @@ def test_ensure_vector_query_nested_empty_list():
ensure_vector_query([[]])
@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)])
def test_vector_query_builder_empty_vector(table, query):
with pytest.raises(ValueError, match="^Query vector must not be empty$"):
LanceVectorQueryBuilder(table, query, "vector")
def test_fast_search(tmp_path):
db = lancedb.connect(tmp_path)
@@ -2109,3 +2265,127 @@ def test_blob_v2_to_batches_row_id(tmp_db):
assert "_rowid" in hits.column_names
blobs = table.fetch_blobs("blob", hits)
assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"]
@pytest.mark.parametrize("metric", ["l2", "cosine", "dot"])
@pytest.mark.parametrize(
"effort,probes", [(None, 4), (0, 1), (0.25, 2), (0.5, 4), (0.75, 8), (1, 16)]
)
def test_search_effort_probe_budget(tmp_path, monkeypatch, metric, effort, probes):
import re
from lancedb.index import IvfFlat
monkeypatch.setenv("LANCE_AUTO_PROBE_MARGIN", "0")
monkeypatch.setenv("LANCE_AUTO_MIN_INITIAL_NPROBES", "4")
monkeypatch.setenv("LANCE_AUTO_MAX_INITIAL_NPROBES", "4")
rng = np.random.default_rng(42)
vectors = np.repeat(np.eye(16, dtype=np.float32), 32, axis=0)
vectors += rng.normal(0, 0.001, vectors.shape).astype(np.float32)
data = pa.table(
{
"id": np.arange(512),
"text": ["dog"] * 512,
"vector": pa.FixedSizeListArray.from_arrays(pa.array(vectors.ravel()), 16),
}
)
table = lancedb.connect(tmp_path).create_table("effort", data.slice(0, 256))
table.add(data.slice(256))
table.create_index(
"vector", config=IvfFlat(distance_type=metric, num_partitions=16)
)
vector = np.eye(16, dtype=np.float32)[0]
def search(value=effort):
return table.search(vector).distance_type(metric).search_effort(value).limit(8)
def partitions(query):
vector_plan = query.analyze_plan().split("FTS Search Plan:", 1)[0]
values = re.findall(r"partitions_searched=(\d+)", vector_plan)
assert values
return sum(map(int, values))
assert partitions(search()) == probes
exact = set(
table.search(vector)
.distance_type(metric)
.bypass_vector_index()
.limit(8)
.to_arrow()["id"]
.to_pylist()
)
actual = set(search().to_arrow()["id"].to_pylist())
assert len(actual & exact) / 8 >= 0.5
expected = 1 if effort is None else probes
assert partitions(search().nprobes(1)) == expected
assert (
partitions(
table.search(vector)
.distance_type(metric)
.nprobes(1)
.search_effort(effort)
.limit(8)
)
== expected
)
assert partitions(search().nprobes(1).search_effort(None)) == 1
assert partitions(search().maximum_nprobes(2)) <= 2
assert partitions(search().minimum_nprobes(3).maximum_nprobes(3)) == 3
assert partitions(search(0).nprobes(16).maximum_nprobes(0)) == 1
if effort is None:
assert search().to_arrow() == search(0.5).to_arrow()
filtered = search(0).maximum_nprobes(2).where("id % 8 = 0", prefilter=True)
assert partitions(filtered) <= 2
assert all(row["id"] % 8 == 0 for row in filtered.to_list())
if metric == "l2" and effort == 0.5:
table.create_index("text", config=FTS())
hybrid = (
table.search(query_type="hybrid")
.vector(vector)
.text("dog")
.nprobes(1)
.search_effort(0.5)
.limit(8)
)
assert partitions(hybrid) == 4
assert partitions(hybrid.search_effort(None)) == 1
@pytest.mark.parametrize(
"effort", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]
)
def test_search_effort_invalid_sync(table, effort):
with pytest.raises(ValueError, match="search_effort must be finite and in"):
table.search([0, 0]).search_effort(effort).to_list()
@pytest.mark.asyncio
@pytest.mark.parametrize("hybrid", [False, True])
@pytest.mark.parametrize("effort", [None, 0, 0.25, 0.5, 0.75, 1])
async def test_search_effort_async_roundtrip(table_async, hybrid, effort):
query = table_async.vector_search([0, 0])
if hybrid:
query = query.nearest_to_text("dog")
query = query.nprobes(5).search_effort(effort).minimum_nprobes(2).maximum_nprobes(3)
request = query.to_query_object()
assert request.search_effort == effort
assert request.nprobes == 5
assert request.minimum_nprobes == 2
assert request.maximum_nprobes == 3
assert query.search_effort(None).to_query_object().search_effort is None
query.search_effort(effort)
if not hybrid:
assert len(await query.limit(1).to_list()) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("hybrid", [False, True])
@pytest.mark.parametrize(
"effort", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]
)
async def test_search_effort_invalid_async(table_async, hybrid, effort):
query = table_async.vector_search([0, 0])
if hybrid:
query = query.nearest_to_text("dog")
with pytest.raises(ValueError, match="search_effort must be finite and in"):
query.search_effort(effort)
python/python/tests/test_remote_db.py
+1382 -73
View File
File diff suppressed because it is too large Load Diff
python/python/tests/test_rerankers.py
+160 -4
View File
@@ -21,6 +21,7 @@ from lancedb.rerankers import (
OpenaiReranker,
JinaReranker,
AnswerdotaiRerankers,
TypeSafeReranker,
VoyageAIReranker,
MRRReranker,
WatsonxReranker,
@@ -81,7 +82,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",
@@ -101,7 +102,15 @@ def get_test_table(tmp_path):
return table, MyTable
def _run_test_reranker(reranker, table, query, query_vector, schema):
def _run_test_reranker(
reranker, table, query, query_vector, schema, deterministic=True
):
"""Exercise a reranker across search types.
Set ``deterministic=False`` for rerankers whose scores can vary between
identical calls, such as remote model APIs; repeated searches are then not
expected to return the same order.
"""
# Hybrid search setting
result1 = (
table.search(query, query_type="hybrid", vector_column_name="vector")
@@ -113,7 +122,8 @@ def _run_test_reranker(reranker, table, query, query_vector, schema):
.rerank(reranker=reranker)
.to_pydantic(schema)
)
assert result1 == result2
if deterministic:
assert result1 == result2
query_vector = table.to_pandas()["vector"][0]
result = (
@@ -203,7 +213,9 @@ def _run_test_reranker(reranker, table, query, query_vector, schema):
)
assert len(result_deduped) <= 20
result_arrow = reranker.rerank_multivector([rs1.to_arrow(), rs2.to_arrow()], query)
assert len(result) == 20 and result == result_arrow
assert len(result) == 20 and len(result_arrow) == 20
if deterministic:
assert result == result_arrow
def _run_test_hybrid_reranker(reranker, tmp_path):
@@ -383,6 +395,38 @@ def test_mrr_multivector_rewards_consensus():
assert result["_rowid"].to_pylist()[0] == 2
def test_rrf_multivector_ranks_each_result_list():
# RRF scores a document by its rank *within each* result list. The top hit
# of the second list must score the same as the top hit of the first list,
# not as if it were ranked after every row of the first list.
reranker = RRFReranker(K=60)
def ranking(row_ids):
return pa.table({"_rowid": pa.array(row_ids, type=pa.int64())})
rs1 = ranking([1, 2, 3])
rs2 = ranking([4, 5, 6])
result = reranker.rerank_multivector([rs1, rs2])
scores = dict(
zip(result["_rowid"].to_pylist(), result["_relevance_score"].to_pylist())
)
assert scores[1] == pytest.approx(1 / 61)
assert scores[4] == pytest.approx(1 / 61)
assert scores[5] == pytest.approx(1 / 62)
# the second-ranked hit of the second list beats the third hit of the first
assert scores[5] > scores[3]
# A document found by both lists sums its reciprocal ranks from each list.
result = reranker.rerank_multivector([ranking([1, 2]), ranking([2, 3])])
scores = dict(
zip(result["_rowid"].to_pylist(), result["_relevance_score"].to_pylist())
)
assert scores[2] == pytest.approx(1 / 62 + 1 / 61)
assert result["_rowid"].to_pylist() == [2, 1, 3]
def test_rrf_reranker_distance():
data = pa.table(
{
@@ -469,6 +513,14 @@ def test_cross_encoder_reranker(tmp_path):
_run_test_reranker(reranker, table, "single player experience", None, schema)
def test_cross_encoder_reranker_kwargs(tmp_path):
pytest.importorskip("sentence_transformers")
reranker = CrossEncoderReranker(batch_size=4, max_length=64)
assert reranker.model.max_length == 64
table, schema = get_test_table(tmp_path)
_run_test_reranker(reranker, table, "single player experience", None, schema)
def test_colbert_reranker(tmp_path):
pytest.importorskip("rerankers")
reranker = ColbertReranker()
@@ -515,6 +567,110 @@ def test_voyageai_reranker(tmp_path):
_run_test_reranker(reranker, table, "single player experience", None, schema)
class _FakeTypeSafeClient:
"""Stands in for ``typesafe_sdk.TypeSafeClient``; scores by word overlap."""
def __init__(self):
self.requests = []
def system_one(self, state, questions, model):
self.requests.append((state, questions, model))
query_words = set(state["query"].lower().split())
answers = {}
for key, question in questions.items():
document = (
state["document"]
if "document" in state
else question["instructions"]["document"]
)
doc_words = set(document.lower().split())
noul = len(query_words & doc_words) / len(query_words)
answers[key] = type("NoulAnswer", (), {"noul": noul})()
return type("SystemOneResponse", (), {"answers": answers})()
@pytest.mark.parametrize("batch_size", [1, 40])
@pytest.mark.parametrize("return_score", ["relevance", "all"])
def test_typesafe_reranker_with_fake_client(tmp_path, batch_size, return_score):
reranker = TypeSafeReranker(
max_concurrency=4, batch_size=batch_size, return_score=return_score
)
reranker._client = _FakeTypeSafeClient()
table, schema = get_test_table(tmp_path)
_run_test_reranker(reranker, table, "single player experience", None, schema)
state, questions, model = reranker._client.requests[0]
assert model == "jev-latest"
if batch_size == 1:
assert set(state) == {"query", "document"}
assert questions == {
"relevance": {
"type": "noul",
"instructions": reranker.instructions,
"criteria": reranker.criteria,
}
}
else:
assert set(state) == {"query"}
assert 1 <= len(questions) <= batch_size
def test_typesafe_reranker_scores_each_row():
reranker = TypeSafeReranker(
column="body",
instructions="Is this about cats?",
criteria={},
return_score="all",
)
reranker._client = _FakeTypeSafeClient()
results = pa.table(
{
"body": ["dogs bark", "cats purr and cats nap", None, "cats"],
"_distance": [0.1, 0.2, 0.3, 0.4],
}
)
reranked = reranker.rerank_vector("cats nap", results)
assert reranked["body"].to_pylist() == [
"cats purr and cats nap",
"cats",
"dogs bark",
None,
]
assert reranked["_relevance_score"].to_pylist() == [1.0, 0.5, 0.0, 0.0]
assert reranked["_distance"].to_pylist() == [0.2, 0.4, 0.1, 0.3]
# Null documents are scored 0 without a request.
assert len(reranker._client.requests) == 3
assert reranker._client.requests[0][1] == {
"relevance": {"type": "noul", "instructions": "Is this about cats?"}
}
def test_typesafe_reranker_rejects_invalid_arguments():
with pytest.raises(ValueError, match="criteria keys"):
TypeSafeReranker(criteria={"yes": "relevant"})
with pytest.raises(ValueError, match="max_concurrency"):
TypeSafeReranker(max_concurrency=0)
@pytest.mark.skipif(
os.environ.get("TYPESAFE_API_KEY") is None, reason="TYPESAFE_API_KEY not set"
)
def test_typesafe_reranker(tmp_path):
pytest.importorskip("typesafe_sdk")
reranker = TypeSafeReranker()
table, schema = get_test_table(tmp_path)
_run_test_reranker(
reranker,
table,
"single player experience",
None,
schema,
deterministic=False,
)
def test_empty_result_reranker():
pytest.importorskip("sentence_transformers")
db = lancedb.connect("memory://")
python/python/tests/test_sql.py
+269
View File
@@ -0,0 +1,269 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import threading
from uuid import UUID
import numpy as np
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, to_parameter_batch
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:
def __init__(self):
self.parameters = []
async def execute_query_async(
self, query, *, default_namespace_path=None, parameters=None
):
self.parameters.append(parameters)
return FakeNativeQuery()
class FakeAsyncConnection:
async def execute_query_async(
self, query, *, default_namespace_path=None, parameters=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
)
def test_positional_parameters_keep_their_types():
vector = np.array([0.1, 0.2, 0.3], dtype=np.float32)
batch = to_parameter_batch(
[42, 1.5, "a", np.float32(0.25), vector, pa.scalar(7, pa.int16())]
)
assert batch.num_rows == 1
assert batch.schema.names == ["1", "2", "3", "4", "5", "6"]
assert batch.schema.types == [
pa.int64(),
pa.float64(),
pa.string(),
pa.float32(),
pa.list_(pa.float32(), 3),
pa.int16(),
]
sent = batch.column(4).flatten().to_numpy()
assert sent.dtype == np.float32
assert sent.tobytes() == vector.tobytes()
def test_named_parameters_and_batches():
half = pa.array(np.array([1.0, 2.0], dtype=np.float16))
batch = to_parameter_batch({"vector": half, "k": 3})
assert batch.schema.names == ["vector", "k"]
assert batch.schema.field("vector").type == pa.list_(pa.float16(), 2)
row = pa.record_batch({"id": [1]})
assert to_parameter_batch(row) is row
assert to_parameter_batch(pa.table({"id": [1]})).equals(row)
assert to_parameter_batch(None) is None
@pytest.mark.parametrize(
"parameters",
[pa.record_batch({"id": [1, 2]}), pa.table({"id": pa.array([], pa.int64())})],
)
def test_parameters_are_exactly_one_row(parameters):
with pytest.raises(ValueError, match="exactly one row"):
to_parameter_batch(parameters)
def test_parameters_must_be_a_collection():
with pytest.raises(TypeError, match="query parameters"):
to_parameter_batch("SELECT 1")
@pytest.mark.asyncio
async def test_async_execute_query_sends_parameters_as_a_batch():
native = FakeNativeConnection()
connection = AsyncConnection(native)
await connection.execute_query("SELECT $1", parameters=[1])
await connection.execute_query_async("SELECT 1")
sent, unparameterized = native.parameters
assert sent.equals(pa.record_batch({"1": [1]}))
assert unparameterized is None
def test_parameters_reach_the_server_as_the_values_that_were_sent():
"""A round trip through the bindings against a Flight server that answers
each exchange with the parameter row it received."""
flight = pytest.importorskip("pyarrow.flight")
class EchoServer(flight.FlightServerBase):
def __init__(self):
super().__init__("grpc://127.0.0.1:0")
self.commands = []
self.lock = threading.Lock()
def do_exchange(self, context, descriptor, reader, writer):
parameters = reader.read_all()
with self.lock:
self.commands.append(descriptor.command)
writer.begin(parameters.schema)
writer.write_table(parameters)
writer.close()
vector = np.array([0.1, 1e-45, -0.0], dtype=np.float32)
with EchoServer() as server:
db = lancedb.connect(
"db://analytics",
api_key="test-key",
host_override="http://localhost:10024",
sql_host_override=f"grpc://127.0.0.1:{server.port}",
)
echoed = db.execute_query(
"SELECT echo", parameters={"vector": vector, "k": 10}
).read_all()
(command,) = server.commands
assert b"CommandStatementQuery" in command
assert b"SELECT echo" in command
assert echoed.schema.field("vector").type == pa.list_(pa.float32(), 3)
received = echoed.column("vector").combine_chunks().flatten().to_numpy()
assert received.view(np.uint32).tolist() == vector.view(np.uint32).tolist()
assert echoed.column("k").to_pylist() == [10]
python/python/tests/test_table.py
+755 -45
View File
@@ -4,6 +4,7 @@
import ctypes
import gc
import json
import os
import sys
import threading
@@ -11,12 +12,13 @@ import warnings
import weakref
from concurrent.futures import ThreadPoolExecutor
from datetime import date, datetime, timedelta
from decimal import Decimal
from decimal import Decimal, localcontext
from time import sleep
from typing import List
from unittest.mock import patch
import lancedb
from lancedb import table as table_module
from lancedb.dependencies import _PANDAS_AVAILABLE
from lancedb.index import BTree, FTS, HnswFlat, HnswPq, HnswSq, IvfPq
import numpy as np
@@ -64,15 +66,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)
@@ -103,6 +113,32 @@ def test_basic(mem_db: DBConnection):
assert table.to_arrow() == expected_data
@pytest.mark.parametrize("enable_v2", [False, True])
def test_migrate_v2_manifest_paths(tmp_path, enable_v2):
db = lancedb.connect(
tmp_path,
storage_options={"new_table_enable_v2_manifest_paths": str(enable_v2).lower()},
)
table = db.create_table("calls", [{"id": 1, "vector": [1.0, 1.0]}])
table.add([{"id": 2, "vector": [2.0, 2.0]}])
expected_data = table.to_arrow()
expected_versions = table.list_versions()
assert table.uses_v2_manifest_paths() == enable_v2
# Migration is also safe to repeat on a table already using v2 paths.
for _ in range(2):
table.migrate_v2_manifest_paths()
assert table.uses_v2_manifest_paths()
reopened = db.open_table("calls")
assert reopened.uses_v2_manifest_paths()
assert reopened.to_arrow() == expected_data
assert reopened.list_versions() == expected_versions
manifests = list((tmp_path / "calls.lance" / "_versions").glob("*.manifest"))
assert len(manifests) == len(expected_versions)
assert all(len(path.stem) == 20 and path.stem.isdigit() for path in manifests)
def test_search_preserves_nulls_from_sliced_arrow_table(mem_db: DBConnection):
data = pa.table(
{
@@ -147,7 +183,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)
@@ -953,6 +993,63 @@ def test_polars(mem_db: DBConnection):
assert len(filtered_result) == 2
@pytest.mark.asyncio
async def test_list_versions_timestamp_precision():
# 2026-10-03T01:03:49.274Z in nanoseconds. Float math turns .274000 into
# .273999, so make sure the conversion is exact.
ts_nanos = 1790989429274000000
class FakeInner:
async def list_versions(self):
return [{"version": 1, "timestamp": ts_nanos, "metadata": {}}]
table = table_module.AsyncTable(FakeInner())
versions = await table.list_versions()
expected = datetime.fromtimestamp(ts_nanos // 1_000_000_000)
assert versions[0]["timestamp"] == expected + timedelta(microseconds=274000)
assert versions[0]["timestamp"].microsecond == 274000
def test_list_versions_timestamp_precision_sync(mem_db: DBConnection):
# The sync LanceTable.list_versions delegates to AsyncTable.list_versions,
# which is the path the parity report hit.
ts_nanos = 1790989429274000000
class FakeInner:
async def list_versions(self):
return [{"version": 1, "timestamp": ts_nanos, "metadata": {}}]
table = mem_db.create_table("ts_precision", data=[{"id": 1}])
table._table = table_module.AsyncTable(FakeInner())
versions = table.list_versions()
expected = datetime.fromtimestamp(ts_nanos // 1_000_000_000)
assert versions[0]["timestamp"] == expected + timedelta(microseconds=274000)
assert versions[0]["timestamp"].microsecond == 274000
def test_list_versions_timestamp_precision_remote():
# RemoteTable.list_versions also delegates to AsyncTable.list_versions.
from lancedb.remote.table import RemoteTable
ts_nanos = 1790989429274000000
class FakeInner:
def name(self):
return "ts_precision"
async def list_versions(self):
return [{"version": 1, "timestamp": ts_nanos, "metadata": {}}]
table = RemoteTable(table_module.AsyncTable(FakeInner()), "dev")
versions = table.list_versions()
expected = datetime.fromtimestamp(ts_nanos // 1_000_000_000)
assert versions[0]["timestamp"] == expected + timedelta(microseconds=274000)
assert versions[0]["timestamp"].microsecond == 274000
def test_versioning(mem_db: DBConnection):
table = mem_db.create_table(
"test",
@@ -2515,6 +2612,51 @@ def test_update_with_arrow_scalar(mem_db: DBConnection):
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0, 3.0, 4.0]]
@pytest.mark.parametrize("use_arrow_scalar", [False, True])
@pytest.mark.parametrize(
"decimal_type, value",
[
(pa.decimal128(10, 2), "9.99"),
(pa.decimal128(10, 2), "-9.99"),
(pa.decimal128(10, 2), "0.00"),
(pa.decimal128(10, 2), "1.23E+5"),
(pa.decimal128(38, 18), "1E-18"),
(pa.decimal128(38, 18), "12345678901234567.123456789012345678"),
(pa.decimal128(38, 18), "12345678901234567890.123456789012345678"),
(pa.decimal128(38, 18), "-12345678901234567890.123456789012345678"),
(
pa.decimal256(76, 38),
"12345678901234567890123456789012345678."
"12345678901234567890123456789012345678",
),
(
pa.decimal256(76, 38),
"-12345678901234567890123456789012345678."
"12345678901234567890123456789012345678",
),
],
)
def test_update_decimal(mem_db: DBConnection, decimal_type, value, use_arrow_scalar):
expected = Decimal(value)
table = mem_db.create_table(
"my_table",
pa.table(
{
"id": [1, 2],
"price": pa.array([Decimal("1"), expected], type=decimal_type),
}
),
)
update = table.to_arrow()["price"][1] if use_arrow_scalar else expected
with localcontext() as context:
context.prec = 6
result = table.update(where="id = 1", values={"price": update})
assert result.rows_updated == 1
actual = table.to_arrow()
assert actual.schema.field("price").type == decimal_type
assert actual["price"].to_pylist() == [expected, expected]
def test_update_types(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
@@ -2682,6 +2824,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 +3109,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 +3146,406 @@ 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")
def test_add_python_objects_to_json_column(mem_db: DBConnection):
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("payload", pa.json_()),
pa.field("metadata", pa.struct([("label", pa.string())])),
pa.field("tags", pa.list_(pa.string())),
]
)
table = mem_db.create_table("json_python_objects_add", schema=schema)
table.add(
[
{
"id": 1,
"payload": {"foo": "bar", "count": 2},
"metadata": {"label": "dict"},
"tags": ["a", "b"],
},
{
"id": 2,
"payload": {
"name": "alice",
"tags": ["x", "y"],
"nested": {"enabled": True},
},
"metadata": {"label": "nested"},
"tags": ["c"],
},
{
"id": 3,
"payload": ["x", "y", "z"],
"metadata": {"label": "list"},
"tags": ["d", "e"],
},
{
"id": 4,
"payload": '{"foo": "bar"}',
"metadata": {"label": "string"},
"tags": ["f"],
},
{
"id": 5,
"payload": None,
"metadata": {"label": "null"},
"tags": ["g"],
},
]
)
rows = {row["id"]: row for row in table.to_arrow().to_pylist()}
assert json.loads(rows[1]["payload"]) == {"foo": "bar", "count": 2}
assert json.loads(rows[2]["payload"]) == {
"name": "alice",
"tags": ["x", "y"],
"nested": {"enabled": True},
}
assert json.loads(rows[3]["payload"]) == ["x", "y", "z"]
assert json.loads(rows[4]["payload"]) == {"foo": "bar"}
assert rows[5]["payload"] is None
assert rows[1]["metadata"] == {"label": "dict"}
assert rows[1]["tags"] == ["a", "b"]
matched = table.search().where("json_extract(payload, '$.count') = '2'").to_list()
assert [row["id"] for row in matched] == [1]
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
def test_merge_insert_python_objects_to_json_column(mem_db: DBConnection):
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("payload", pa.json_()),
pa.field("metadata", pa.struct([("label", pa.string())])),
pa.field("tags", pa.list_(pa.string())),
]
)
table = mem_db.create_table("json_python_objects_merge", schema=schema)
table.add(
[
{
"id": 1,
"payload": {"old": True},
"metadata": {"label": "old"},
"tags": ["old"],
}
]
)
table.merge_insert(
"id"
).when_matched_update_all().when_not_matched_insert_all().execute(
[
{
"id": 1,
"payload": {"foo": "bar", "count": 2},
"metadata": {"label": "dict"},
"tags": ["a", "b"],
},
{
"id": 2,
"payload": {
"name": "alice",
"tags": ["x", "y"],
"nested": {"enabled": True},
},
"metadata": {"label": "nested"},
"tags": ["c"],
},
{
"id": 3,
"payload": ["x", "y", "z"],
"metadata": {"label": "list"},
"tags": ["d", "e"],
},
{
"id": 4,
"payload": '{"foo": "bar"}',
"metadata": {"label": "string"},
"tags": ["f"],
},
{
"id": 5,
"payload": None,
"metadata": {"label": "null"},
"tags": ["g"],
},
]
)
rows = {row["id"]: row for row in table.to_arrow().to_pylist()}
assert json.loads(rows[1]["payload"]) == {"foo": "bar", "count": 2}
assert json.loads(rows[2]["payload"]) == {
"name": "alice",
"tags": ["x", "y"],
"nested": {"enabled": True},
}
assert json.loads(rows[3]["payload"]) == ["x", "y", "z"]
assert json.loads(rows[4]["payload"]) == {"foo": "bar"}
assert rows[5]["payload"] is None
assert rows[1]["metadata"] == {"label": "dict"}
assert rows[1]["tags"] == ["a", "b"]
matched = table.search().where("json_extract(payload, '$.count') = '2'").to_list()
assert [row["id"] for row in matched] == [1]
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
def test_add_python_objects_to_nested_json_fields(mem_db: DBConnection):
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field(
"info",
pa.struct([pa.field("payload", pa.json_())]),
),
pa.field("documents", pa.list_(pa.field("item", pa.json_()))),
]
)
table = mem_db.create_table("nested_json_python_objects_add", schema=schema)
table.add(
[
{
"id": 1,
"info": {"payload": {"kind": "struct", "value": 1}},
"documents": [{"kind": "list", "value": 2}],
}
]
)
row = table.to_arrow().to_pylist()[0]
assert json.loads(row["info"]["payload"]) == {"kind": "struct", "value": 1}
assert json.loads(row["documents"][0]) == {"kind": "list", "value": 2}
assert [
row["id"]
for row in table.search()
.where("json_extract(info.payload, '$.value') = '1'")
.to_list()
] == [1]
assert [
row["id"]
for row in table.search()
.where("json_extract(documents[1], '$.value') = '2'")
.to_list()
] == [1]
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
def test_json_serialization_plan_skips_non_json_branches():
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("vector", pa.list_(pa.float32(), 3)),
pa.field("payload", pa.json_()),
pa.field(
"metadata",
pa.struct(
[
pa.field("label", pa.string()),
pa.field("payload", pa.json_()),
]
),
),
pa.field(
"documents",
pa.list_(
pa.struct(
[
pa.field("title", pa.string()),
pa.field("payload", pa.json_()),
]
)
),
),
]
)
plans = table_module._json_serialization_plans(schema)
assert set(plans) == {"payload", "metadata", "documents"}
metadata_children = plans["metadata"].children
assert metadata_children is not None
assert set(metadata_children) == {"payload"}
documents_item = plans["documents"].item
assert documents_item is not None
assert documents_item.children is not None
assert set(documents_item.children) == {"payload"}
@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 +3871,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
@@ -3313,6 +3879,24 @@ def test_empty_query(mem_db: DBConnection):
assert df.num_rows == 42
@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)])
@pytest.mark.parametrize("vector_column_name", [None, "vector"])
@pytest.mark.parametrize("query_type", ["auto", "vector"])
@pytest.mark.parametrize("multiple_vector_columns", [False, True])
def test_search_empty_vector(
mem_db, query, vector_column_name, query_type, multiple_vector_columns
):
fields = [pa.field("vector", pa.list_(pa.float32(), 8))]
if multiple_vector_columns:
fields.append(pa.field("vec2", pa.list_(pa.float32(), 4)))
table = mem_db.create_table("empty_vector_query", schema=pa.schema(fields))
with pytest.raises(ValueError, match="^Query vector must not be empty$"):
table.search(
query, vector_column_name=vector_column_name, query_type=query_type
).limit(3).to_arrow()
def test_search_with_schema_inf_single_vector(mem_db: DBConnection):
class MyTable(LanceModel):
text: str
@@ -3706,14 +4290,74 @@ async def test_add_columns_with_schema(mem_db_async: AsyncConnection):
)
def test_alter_columns(mem_db: DBConnection):
@pytest.mark.parametrize("rename_key", ["rename", "name"])
def test_alter_columns(mem_db: DBConnection, rename_key):
data = pa.table({"id": [0, 1]})
table = mem_db.create_table("my_table", data=data)
alter_columns_res = table.alter_columns({"path": "id", "rename": "new_id"})
alter_columns_res = table.alter_columns({"path": "id", rename_key: "new_id"})
assert alter_columns_res.version == 2
assert table.to_arrow().column_names == ["new_id"]
INVALID_COLUMN_ALTERATIONS = [
(({"path": "id"},), "One of rename, nullable or data_type"),
(
({"path": "id", "nulable": False},), # spellchecker:disable-line
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
(
(
{
"path": "id",
"rename": "new_id",
"nulable": False, # spellchecker:disable-line
},
),
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
(
({"path": "id", "rename": "new_id"}, {"path": "id"}),
"One of rename, nullable or data_type",
),
]
def test_alter_columns_nullable_false(mem_db: DBConnection):
table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]}))
result = table.alter_columns({"path": "id", "nullable": False})
assert result.version == 2
assert not table.schema.field("id").nullable
@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS)
def test_alter_columns_rejects_invalid(mem_db: DBConnection, alterations, match):
table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]}))
initial_version = table.version
initial_schema = table.schema
with pytest.raises(ValueError, match=match):
table.alter_columns(*alterations)
assert table.version == initial_version
assert table.schema == initial_schema
@pytest.mark.asyncio
@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS)
async def test_alter_columns_rejects_invalid_async(
mem_db_async: AsyncConnection, alterations, match
):
table = await mem_db_async.create_table("my_table", data=pa.table({"id": [0, 1]}))
initial_version = await table.version()
initial_schema = await table.schema()
with pytest.raises(ValueError, match=match):
await table.alter_columns(*alterations)
assert await table.version() == initial_version
assert await table.schema() == initial_schema
def test_update_field_metadata(mem_db: DBConnection):
data = pa.table({"id": [0, 1], "category": ["a", "b"]})
table = mem_db.create_table("my_table", data=data)
@@ -3739,10 +4383,11 @@ def test_update_field_metadata(mem_db: DBConnection):
@pytest.mark.asyncio
async def test_alter_columns_async(mem_db_async: AsyncConnection):
@pytest.mark.parametrize("rename_key", ["rename", "name"])
async def test_alter_columns_async(mem_db_async: AsyncConnection, rename_key):
data = pa.table({"id": [0, 1]})
table = await mem_db_async.create_table("my_table", data=data)
alter_columns_res = await table.alter_columns({"path": "id", "rename": "new_id"})
alter_columns_res = await table.alter_columns({"path": "id", rename_key: "new_id"})
assert alter_columns_res.version == 2
assert (await table.to_arrow()).column_names == ["new_id"]
alter_columns_res = await table.alter_columns(
@@ -3869,7 +4514,8 @@ async def test_optimize(mem_db_async: AsyncConnection):
assert stats.prune.bytes_removed == 0
assert stats.prune.old_versions_removed == 0
stats = await table.optimize(cleanup_older_than=timedelta(seconds=0))
with pytest.warns(UserWarning, match="concurrent"):
stats = await table.optimize(cleanup_older_than=timedelta(seconds=0))
assert stats.prune.bytes_removed > 0
assert stats.prune.old_versions_removed == 3
@@ -3897,12 +4543,35 @@ async def test_optimize_delete_unverified(tmp_db_async: AsyncConnection, tmp_pat
stats = await table.optimize(delete_unverified=False)
assert stats.prune.old_versions_removed == 0
stats = await table.optimize(
cleanup_older_than=timedelta(seconds=0), delete_unverified=True
)
with pytest.warns(UserWarning, match="concurrent"):
stats = await table.optimize(
cleanup_older_than=timedelta(seconds=0), delete_unverified=True
)
assert stats.prune.old_versions_removed == 2
@pytest.mark.asyncio
async def test_optimize_warns_on_zero_cleanup(mem_db_async: AsyncConnection):
table = await mem_db_async.create_table("test", data=[{"x": [1]}])
with pytest.warns(UserWarning, match="concurrent"):
await table.optimize(cleanup_older_than=timedelta(0))
with warnings.catch_warnings():
warnings.simplefilter("error")
await table.optimize(cleanup_older_than=timedelta(days=1))
await table.optimize()
def test_optimize_warns_on_zero_cleanup_sync(mem_db: DBConnection):
table = mem_db.create_table("test", data=[{"x": [1]}])
with warnings.catch_warnings(record=True) as seen:
warnings.simplefilter("default")
table.optimize(cleanup_older_than=timedelta(0))
table.optimize(cleanup_older_than=timedelta(0))
assert [w.filename for w in seen] == [__file__, __file__]
assert all("concurrent" in str(w.message) for w in seen)
def test_replace_field_metadata(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("my_table", data=[{"x": 0}])
@@ -3917,22 +4586,25 @@ def test_stats(mem_db: DBConnection):
"my_table",
data=[{"text": "foo", "id": 0}, {"text": "bar", "id": 1}],
)
assert len(table) == 2
table.add([{"text": "baz", "id": 1}])
assert len(table) == 3
stats = table.stats()
print(f"{stats=}")
assert stats == {
# Full on-disk size of the data file, footer and metadata included.
"total_bytes": 633,
"num_rows": 2,
# Full on-disk size of the data files, footer and metadata included:
# 637 bytes for the two-row fragment plus 415 for the one-row fragment.
"total_bytes": 1052,
"num_rows": 3,
"num_deleted_rows": 0,
"num_indices": 0,
"fragment_stats": {
"num_fragments": 1,
"num_small_fragments": 1,
"num_fragments": 2,
"num_small_fragments": 2,
"lengths": {
"min": 2,
"min": 1,
"max": 2,
"mean": 2,
"p25": 2,
"mean": 1,
"p25": 1,
"p50": 2,
"p75": 2,
"p99": 2,
@@ -3940,12 +4612,19 @@ def test_stats(mem_db: DBConnection):
},
}
# Both rows with id = 1 are deleted, but that empties the second fragment,
# which is dropped outright rather than kept with a deletion file, so only
# the row marked in the surviving fragment is counted.
table.delete("id = 1")
# Index files count toward total_bytes too (only deletion files and
# manifests are excluded).
table.create_index("id", config=BTree())
stats_with_index = table.stats()
assert stats_with_index["num_indices"] == 1
assert stats_with_index["total_bytes"] > stats["total_bytes"]
assert stats_with_index["num_rows"] == 1
assert stats_with_index["num_deleted_rows"] == 1
assert stats_with_index["fragment_stats"]["num_fragments"] == 1
def test_create_table_empty_list_with_schema(mem_db: DBConnection):
@@ -4087,6 +4766,29 @@ def test_computed_column_rejects_transforms_and_computed_together(tmp_path):
table.add_columns({"a": "x + 1"}, computed={"b": "x * 2"})
def test_computed_column_blob_projection_inherits_semantics(tmp_path):
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
db = lancedb.connect(tmp_path)
table = db.create_table("computed_column_blob", schema=schema)
table.add(
[
{"id": 1, "image": b"hello"},
{"id": 2, "image": b""},
{"id": 3, "image": None},
]
)
table.add_columns(computed={"image_copy": "image", "second_copy": "image_copy"})
assert table.refresh_column("image_copy").rows_filled == 2
assert table.refresh_column("second_copy").rows_filled == 2
assert table.blob_columns() == ["image", "image_copy", "second_copy"]
hits = table.search().with_row_id(True).limit(10).to_arrow()
rows = sorted(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
copied = table.fetch_blobs("second_copy", [row_id for _, row_id in rows])
assert copied.to_pylist() == [b"hello", b"", None]
@pytest.mark.asyncio
async def test_computed_column_async(tmp_path):
db = await lancedb.connect_async(tmp_path)
@@ -4098,6 +4800,13 @@ async def test_computed_column_async(tmp_path):
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
def test_function_errors_are_remote_only(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("t", [{"x": 1}])
with pytest.raises(NotImplementedError, match="LanceDB Cloud and Enterprise"):
table.function_errors()
def test_refresh_column_async_returns_job(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("computed_job", [{"x": 1}, {"x": 2}])
@@ -4111,13 +4820,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.
@@ -4136,6 +4846,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/python/tests/test_typesafe_reranker.py
+338
View File
@@ -0,0 +1,338 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import json
from threading import Barrier, Event, Lock
from types import SimpleNamespace
from unittest.mock import Mock
import lancedb
import pyarrow as pa
import pytest
from lancedb.index import FTS
from lancedb.rerankers import TypeSafeReranker
def response(questions, probability=0.5):
return SimpleNamespace(
answers={key: SimpleNamespace(noul=probability) for key in reversed(questions)}
)
def vector_results(documents):
return pa.table(
{
"text": pa.array(documents, type=pa.string()),
"_rowid": pa.array(range(len(documents)), type=pa.uint64()),
"_distance": pa.array(range(len(documents)), type=pa.float32()),
}
)
def fake_reranker(**kwargs):
reranker = TypeSafeReranker(**kwargs)
reranker._client = Mock()
reranker._client.system_one.side_effect = lambda **kw: response(kw["questions"])
return reranker
@pytest.mark.parametrize("batch_size", [0, -1, 1.5, 40.0, True, False, "40", None])
def test_invalid_batch_size(batch_size):
with pytest.raises(ValueError, match="batch_size must be a positive integer"):
TypeSafeReranker(batch_size=batch_size)
@pytest.mark.parametrize("batch_size", [1, 7, 40])
@pytest.mark.parametrize("count", [0, 1, 39, 40, 41, 80, 81])
def test_request_counts_and_boundaries(batch_size, count):
reranker = fake_reranker(batch_size=batch_size)
docs = [f"candidate {i}" for i in range(count)]
result = reranker.rerank_vector("query", vector_results(docs))
calls = reranker._client.system_one.call_args_list
assert len(calls) == (count + batch_size - 1) // batch_size
assert sorted(len(call.kwargs["questions"]) for call in calls) == sorted(
min(batch_size, count - start) for start in range(0, count, batch_size)
)
assert result["_rowid"].to_pylist() == list(range(count))
@pytest.mark.parametrize("batch_size", [1, 40])
@pytest.mark.parametrize(
"criteria", [{}, {"true": " state.document fits\nexactly", "false": "No!"}]
)
def test_payload_isolation_and_custom_prompts(batch_size, criteria):
instructions = "Read state.document and state.query.\n Preserve this text. "
reranker = fake_reranker(
batch_size=batch_size,
instructions=instructions,
criteria=criteria,
model_name="jev-1.13.0",
column="body",
)
docs = [f"UNIQUE_CANDIDATE_{i}" for i in range(41)]
table = vector_results(docs).rename_columns(["body", "_rowid", "_distance"])
reranker.rerank_vector("shared query", table)
seen = []
for call in reranker._client.system_one.call_args_list:
state, questions = call.kwargs["state"], call.kwargs["questions"]
assert call.kwargs["model"] == "jev-1.13.0"
if batch_size == 1:
doc = state["document"]
assert state == {"query": "shared query", "document": doc}
assert list(questions) == ["relevance"]
expected_instructions = instructions
else:
assert state == {"query": "shared query"}
for question in questions.values():
if batch_size > 1:
doc = question["instructions"]["document"]
expected_instructions = {"question": instructions, "document": doc}
expected = {"type": "noul", "instructions": expected_instructions}
if criteria:
expected["criteria"] = criteria
# Exact equality rules out other candidates anywhere in this question.
assert question == expected
seen.append(doc)
assert sorted(seen) == sorted(docs)
assert reranker.instructions == instructions
assert reranker.criteria == criteria
def test_default_payload():
reranker = fake_reranker()
reranker.rerank_vector("query", vector_results(["document"]))
reranker._client.system_one.assert_called_once_with(
state={"query": "query", "document": "document"},
questions={
"relevance": {
"type": "noul",
"instructions": reranker.instructions,
"criteria": reranker.criteria,
}
},
model="jev-latest",
)
def test_response_mapping_with_duplicates_and_nulls():
reranker = fake_reranker(batch_size=2, max_concurrency=2, return_score="all")
second_batch_scored = Event()
def score(**kw):
questions = kw["questions"]
documents = [
question["instructions"]["document"] for question in questions.values()
]
if documents == ["same", "same"]:
# Prepare the later batch first; result order must still match input.
assert second_batch_scored.wait(timeout=10)
scores = [0.0, 0.5]
else:
assert documents == ["different", "same"]
scores = [0.75, 1.0]
result = SimpleNamespace(
answers={
key: SimpleNamespace(noul=value)
for key, value in reversed(list(zip(questions, scores)))
}
)
second_batch_scored.set()
return result
reranker._client.system_one.side_effect = score
result = reranker.rerank_vector(
"query", vector_results(["same", None, "same", "different", "same"])
)
assert result["_rowid"].to_pylist() == [4, 3, 2, 0, 1]
assert result["_distance"].to_pylist() == [4, 3, 2, 0, 1]
assert result["_relevance_score"].to_pylist() == [1, 0.75, 0.5, 0, 0]
assert reranker._client.system_one.call_count == 2
@pytest.mark.parametrize("batch_size", [1, 40])
@pytest.mark.parametrize("problem", ["missing", "unexpected", "replaced", "no_noul"])
def test_invalid_answers(batch_size, problem):
reranker = fake_reranker(batch_size=batch_size)
def invalid(**kw):
result = response(kw["questions"])
key = next(iter(result.answers))
if problem in ("missing", "replaced"):
del result.answers[key]
if problem in ("unexpected", "replaced"):
result.answers["unknown"] = SimpleNamespace(noul=0.5)
if problem == "no_noul":
result.answers[key] = SimpleNamespace()
return result
reranker._client.system_one.side_effect = invalid
with pytest.raises(ValueError, match="invalid relevance|answer IDs"):
reranker.rerank_vector("query", vector_results(["one", "two"]))
@pytest.mark.parametrize("batch_size", [1, 40])
@pytest.mark.parametrize(
"probability",
[None, True, False, "0.5", -0.01, 1.01, float("nan"), float("inf"), -float("inf")],
)
def test_invalid_probabilities(batch_size, probability):
reranker = fake_reranker(batch_size=batch_size)
reranker._client.system_one.side_effect = lambda **kw: response(
kw["questions"], probability
)
with pytest.raises(ValueError, match="invalid relevance probability"):
reranker.rerank_vector("query", vector_results(["one", "two"]))
@pytest.mark.parametrize("batch_size", [1, 40])
def test_api_failure_propagates(batch_size):
reranker = fake_reranker(batch_size=batch_size)
error = RuntimeError("SDK exhausted retries")
reranker._client.system_one.side_effect = error
with pytest.raises(RuntimeError) as caught:
reranker.rerank_vector("query", vector_results(["one"]))
assert caught.value is error
assert reranker._client.system_one.call_count == 1
@pytest.mark.parametrize("batch_size", [1, 40])
def test_concurrency_limits_requests(batch_size):
reranker = fake_reranker(batch_size=batch_size, max_concurrency=2)
barrier, lock = Barrier(2), Lock()
active = peak = 0
def blocking(**kw):
nonlocal active, peak
with lock:
active += 1
peak = max(peak, active)
barrier.wait(timeout=10)
with lock:
active -= 1
return response(kw["questions"])
reranker._client.system_one.side_effect = blocking
reranker.rerank_vector("query", vector_results(["doc"] * (6 * batch_size)))
assert peak == 2
assert active == 0
assert reranker._client.system_one.call_count == 6
@pytest.mark.parametrize("batch_size", [1, 40])
@pytest.mark.parametrize("return_score", ["relevance", "all"])
@pytest.mark.parametrize("method", ["vector", "fts", "hybrid"])
@pytest.mark.parametrize("docs", [[], [None, None], [None, "", None]])
def test_empty_and_null_inputs(batch_size, return_score, method, docs):
reranker = fake_reranker(batch_size=batch_size, return_score=return_score)
vector = vector_results(docs)
fts = vector.rename_columns(["text", "_rowid", "_score"])
args = (
(vector, fts)
if method == "hybrid"
else (vector if method == "vector" else fts,)
)
result = getattr(reranker, f"rerank_{method}")("query", *args)
assert len(result) == len(docs)
assert result["_relevance_score"].type == pa.float32()
assert sorted(result["_relevance_score"].to_pylist()) == sorted(
0 if doc is None else 0.5 for doc in docs
)
assert reranker._client.system_one.call_count == docs.count("")
if return_score == "relevance":
assert "_distance" not in result.column_names
assert "_score" not in result.column_names
@pytest.mark.parametrize("batch_size", [1, 40])
@pytest.mark.parametrize("return_score", ["relevance", "all"])
def test_hybrid_deduplicates_by_row_id_and_keeps_scores(batch_size, return_score):
reranker = fake_reranker(batch_size=batch_size, return_score=return_score)
vector = vector_results(["same", "same", "third"])
fts = vector.slice(1).rename_columns(["text", "_rowid", "_score"])
result = reranker.rerank_hybrid("query", vector.slice(0, 2), fts)
assert result["_rowid"].to_pylist() == [0, 1, 2]
assert (
sum(
len(call.kwargs["questions"])
for call in reranker._client.system_one.call_args_list
)
== 3
)
assert reranker._client.system_one.call_count == (3 if batch_size == 1 else 1)
if return_score == "all":
assert result["_distance"].to_pylist() == [0, 1, None]
assert result["_score"].to_pylist() == [None, 1, 2]
else:
assert result.column_names == ["text", "_rowid", "_relevance_score"]
@pytest.mark.asyncio
@pytest.mark.parametrize("batch_size", [1, 40])
async def test_async_search_integrations(tmp_path, batch_size):
db = await lancedb.connect_async(tmp_path)
table = await db.create_table(
"docs",
[
{"text": "cat naps", "vector": [1.0, 0.0]},
{"text": "cat plays", "vector": [0.0, 1.0]},
],
)
await table.create_index("text", config=FTS())
reranker = fake_reranker(batch_size=batch_size)
queries = [
table.query().nearest_to([1.0, 0.0]).rerank(reranker, query_string="cat"),
table.query().nearest_to_text("cat").rerank(reranker),
table.query().nearest_to([1.0, 0.0]).nearest_to_text("cat").rerank(reranker),
]
for query in queries:
reranker._client.system_one.reset_mock()
result = await query.with_row_id().to_arrow()
assert sorted(result["_rowid"].to_pylist()) == [0, 1]
assert result["_relevance_score"].to_pylist() == [0.5, 0.5]
assert "_distance" not in result.column_names
assert "_score" not in result.column_names
assert reranker._client.system_one.call_count == (2 if batch_size == 1 else 1)
@pytest.mark.parametrize("batch_size", [1, 40])
@pytest.mark.parametrize("retry_once", [False, True])
def test_sdk_transport_and_retries(batch_size, retry_once):
sdk = pytest.importorskip("typesafe_sdk")
httpx = pytest.importorskip("httpx2")
payloads = []
def handle(request):
payload = json.loads(request.content)
payloads.append(payload)
if retry_once and len(payloads) == 1:
return httpx.Response(503, headers={"retry-after-ms": "0"})
return httpx.Response(
200,
json={
"model": "jev-1.13.0",
"usage": {},
"answers": {
key: {"type": "noul", "noul": 0.5}
for key in reversed(payload["questions"])
},
},
)
with sdk.TypeSafeClient(
api_key="test-key", transport=httpx.MockTransport(handle)
) as client:
reranker = TypeSafeReranker(batch_size=batch_size, max_concurrency=1)
reranker._client = client
result = reranker.rerank_vector("query", vector_results(["document"] * 80))
assert len(payloads) == (80 if batch_size == 1 else 2) + int(retry_once)
assert result["_relevance_score"].to_pylist() == [0.5] * 80
if retry_once:
assert payloads[0] == payloads[1]
for payload in payloads:
assert len(payload["questions"]) == batch_size
if batch_size == 40:
assert payload["state"] == {"query": "query"}
assert all(
question["instructions"]["document"] == "document"
for question in payload["questions"].values()
)
python/python/tests/test_util.py
+274
View File
@@ -4,9 +4,11 @@
import os
import pathlib
from decimal import Decimal, localcontext
from typing import Optional
import lance
from lance.blob import BlobType as LanceBlobType
from lancedb.conftest import MockTextEmbeddingFunction
from lancedb.embeddings.base import EmbeddingFunctionConfig
from lancedb.embeddings.registry import EmbeddingFunctionRegistry
@@ -156,6 +158,41 @@ def test_value_to_sql_string(tmp_path):
assert table.to_pandas().query("search == @value")["replace"].item() == value
@pytest.mark.parametrize(
"value, expected",
[
("9.99", "'9.99'"),
("-9.99", "'-9.99'"),
("0", "'0'"),
("-0.00", "'-0.00'"),
("1.23E+5", "'123000'"),
("1E-18", "'0.000000000000000001'"),
(
"12345678901234567890.123456789012345678",
"'12345678901234567890.123456789012345678'",
),
(
"12345678901234567890123456789012345678."
"12345678901234567890123456789012345678",
"'12345678901234567890123456789012345678."
"12345678901234567890123456789012345678'",
),
],
)
def test_value_to_sql_decimal(value, expected):
with localcontext() as context:
context.prec = 6
assert value_to_sql(Decimal(value)) == expected
@pytest.mark.parametrize("value", ["NaN", "sNaN", "NaN123456", "Infinity", "-Infinity"])
def test_value_to_sql_decimal_non_finite(value):
with pytest.raises(
ValueError, match="^Non-finite Decimal values cannot be converted to SQL$"
):
value_to_sql(Decimal(value))
def test_value_to_sql_dict():
# Simple flat struct
assert value_to_sql({"a": 1, "b": "hello"}) == "named_struct('a', 1, 'b', 'hello')"
@@ -212,6 +249,35 @@ def test_value_to_sql_numpy_scalars():
assert value_to_sql(np.bool_(False)) == "FALSE"
def test_value_to_sql_nan_and_infinity(tmp_path):
# str(float("nan")) is "nan", which SQL reads as a column named nan, so
# table.update(values={"x": float("nan")}) failed with "No field named nan".
import math
import numpy as np
db = lancedb.connect(tmp_path)
table = db.create_table(
"test",
pa.table(
{
"id": [1, 2, 3, 4],
"f64": [1.0, 2.0, 3.0, 4.0],
"f32": pa.array([1.0, 2.0, 3.0, 4.0], pa.float32()),
}
),
)
table.update(where="id = 1", values={"f64": float("nan"), "f32": np.nan})
table.update(where="id = 2", values={"f64": math.inf, "f32": np.float32("inf")})
table.update(where="id = 3", values={"f64": -math.inf, "f32": -math.inf})
rows = {row["id"]: row for row in table.to_arrow().to_pylist()}
assert math.isnan(rows[1]["f64"]) and math.isnan(rows[1]["f32"])
assert rows[2]["f64"] == math.inf and rows[2]["f32"] == math.inf
assert rows[3]["f64"] == -math.inf and rows[3]["f32"] == -math.inf
assert rows[4]["f64"] == 4.0 and rows[4]["f32"] == 4.0
def test_append_vector_columns():
registry = EmbeddingFunctionRegistry.get_instance()
registry.register("test")(MockTextEmbeddingFunction)
@@ -907,6 +973,214 @@ def test_cast_to_target_schema():
assert output == expected
def test_cast_to_target_schema_coerces_binary_to_blob_v2():
data = pa.table({"image": pa.array([b"hello", None], type=pa.binary())})
target = pa.schema([lancedb.blob("image")])
output = _cast_to_target_schema(data.to_reader(), target).read_all()
image = output["image"].chunk(0)
assert type(image.type) is lancedb.BlobType
assert image.storage.to_pylist() == [
{"data": b"hello", "uri": None, "position": None, "size": None},
None,
]
def test_cast_to_target_schema_coerces_binary_to_metadata_blob_struct():
storage = lancedb.blob("image").type.storage_type
target = pa.schema(
[
pa.field(
"image",
storage,
metadata={
b"ARROW:extension:name": b"lance.blob.v2",
b"ARROW:extension:metadata": b"",
},
)
]
)
data = pa.table({"image": pa.array([b"hello", None], type=pa.binary())})
output = _cast_to_target_schema(data.to_reader(), target).read_all()
image = output["image"].chunk(0)
assert not isinstance(image.type, pa.ExtensionType)
assert image.to_pylist() == [
{"data": b"hello", "uri": None, "position": None, "size": None},
None,
]
def test_cast_to_target_schema_labels_json_text_with_table_extension_metadata():
# A table created from pa.json_() stores the empty ARROW:extension:metadata that
# pyarrow exports next to the name. Lance keeps that key when it turns the
# arrow.json label back into lance.json on write and then requires the field to
# match the stored one exactly, so the label has to mirror the table's metadata.
target = pa.schema(
[
pa.field(
"payload",
pa.large_binary(),
metadata={
b"ARROW:extension:name": b"lance.json",
b"ARROW:extension:metadata": b"",
},
),
pa.field(
"docs",
pa.list_(
pa.field(
"item",
pa.large_binary(),
metadata={b"ARROW:extension:name": b"lance.json"},
)
),
),
]
)
data = pa.table(
{
"payload": pa.array(['{"a": 1}', None], type=pa.string()),
"docs": pa.array([['{"b": 2}'], []], type=pa.list_(pa.string())),
}
)
output = _cast_to_target_schema(data.to_reader(), target).read_all()
payload = output.schema.field("payload")
assert payload.type == pa.string()
assert payload.metadata == {
b"ARROW:extension:name": b"arrow.json",
b"ARROW:extension:metadata": b"",
}
docs_item = output.schema.field("docs").type.value_field
assert docs_item.type == pa.string()
assert docs_item.metadata == {b"ARROW:extension:name": b"arrow.json"}
assert output["payload"].to_pylist() == ['{"a": 1}', None]
assert output["docs"].to_pylist() == [['{"b": 2}'], []]
def test_cast_to_target_schema_coerces_nested_binary_blob():
data = pa.table(
{
"info": pa.array(
[{"blob": b"hello"}, {"blob": None}],
type=pa.struct([pa.field("blob", pa.binary())]),
)
}
)
target = pa.schema([pa.field("info", pa.struct([lancedb.blob("blob")]))])
output = _cast_to_target_schema(data.to_reader(), target).read_all()
blob = output["info"].chunk(0).field("blob")
assert type(blob.type) is lancedb.BlobType
assert blob.storage.to_pylist() == [
{"data": b"hello", "uri": None, "position": None, "size": None},
None,
]
def test_cast_to_target_schema_coerces_list_binary_blob_with_inferred_child_name():
data = pa.table(
{"images": pa.array([[b"a", b"b"], None], type=pa.list_(pa.binary()))}
)
target = pa.schema([pa.field("images", pa.list_(lancedb.blob("image")))])
output = _cast_to_target_schema(data.to_reader(), target).read_all()
images = output["images"].chunk(0)
assert images.type.value_field.name == "image"
assert type(images.type.value_type) is lancedb.BlobType
assert images.to_pylist()[1] is None
assert images.values.storage.to_pylist() == [
{"data": b"a", "uri": None, "position": None, "size": None},
{"data": b"b", "uri": None, "position": None, "size": None},
]
def test_list_blob_coercion_preserves_null_slots_with_nonzero_extent():
child = pa.field("image", pa.binary())
source = pa.ListArray.from_arrays(
pa.array([0, 2, 4], type=pa.int32()),
pa.array([b"a", b"b", b"dead", b"beef"], type=pa.binary()),
mask=pa.array([False, True]),
).cast(pa.list_(child))
target = pa.schema([pa.field("images", pa.list_(lancedb.blob("image")))])
output = _cast_to_target_schema(
pa.table({"images": source}).to_reader(), target
).read_all()
images = output["images"].chunk(0)
assert images.to_pylist()[1] is None
assert [b["data"] for b in images.to_pylist()[0]] == [b"a", b"b"]
def test_fixed_size_list_blob_coercion_keeps_null_rows():
child = pa.field("frame", pa.binary())
source = (
pa.FixedSizeListArray.from_arrays(
pa.array([b"a", b"b", b"c", b"d"], type=pa.binary()), 2
)
.take(pa.array([0, None], type=pa.int32()))
.cast(pa.list_(child, 2))
)
target = pa.schema([pa.field("frames", pa.list_(lancedb.blob("frame"), 2))])
output = _cast_to_target_schema(
pa.table({"frames": source}).to_reader(), target
).read_all()
frames = output["frames"].chunk(0)
assert frames.to_pylist()[1] is None
assert [b["data"] for b in frames.to_pylist()[0]] == [b"a", b"b"]
def test_cast_to_target_schema_accepts_pylance_blob_v2():
target_type = lancedb.BlobType()
source = lance.blob_array([b"hello", None])
assert type(source.type) is LanceBlobType
assert type(source.type) is type(target_type)
data = pa.table({"image": source})
target = pa.schema([pa.field("image", target_type)])
output = _cast_to_target_schema(data.to_reader(), target).read_all()
image = output["image"].chunk(0)
assert type(image.type) is LanceBlobType
assert image.type == target_type
assert image.storage.to_pylist() == [
{"data": b"hello", "uri": None, "position": None, "size": None},
None,
]
def test_cast_to_target_schema_rejects_different_blob_v2_class():
class OtherBlobType(pa.ExtensionType):
def __init__(self):
super().__init__(lancedb.BlobType().storage_type, "lance.blob.v2")
def __arrow_ext_serialize__(self) -> bytes:
return b""
@classmethod
def __arrow_ext_deserialize__(
cls, storage_type: pa.DataType, serialized: bytes
) -> "OtherBlobType":
return cls()
storage = lance.blob_array([b"hello"]).storage
source = pa.ExtensionArray.from_storage(OtherBlobType(), storage)
data = pa.table({"image": source})
target = pa.schema([lancedb.blob("image")])
with pytest.raises(pa.ArrowTypeError, match="different extension type"):
_cast_to_target_schema(data.to_reader(), target).read_all()
def test_sanitize_data_stream():
# Make sure we don't collect the whole stream when running sanitize_data
schema = pa.schema({"a": pa.int32()})
python/python/tests/test_wal_hybrid.py
+219
View File
@@ -0,0 +1,219 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""WAL-PK-FUSION: delete this file with `lancedb/_wal_hybrid.py`."""
from unittest import mock
import lancedb
import pyarrow as pa
import pytest
from lancedb._wal_hybrid import WAL_ROW_ID_UNSUPPORTED, with_pk_fallback
from lancedb.query import LanceHybridQueryBuilder
from lancedb.rerankers import RRFReranker
def _wal_table(schema: pa.Schema) -> mock.Mock:
"""A MemWAL-backed table a previous query has already been refused
`_rowid` on, so the fusion goes straight to the primary key."""
table = mock.Mock()
table._hybrid_pk_fusion_learned.return_value = True
table.schema = schema
return table
_PK_SCHEMA = pa.schema(
[
pa.field(
"id",
pa.string(),
metadata={b"lance-schema:unenforced-primary-key:position": b"1"},
),
pa.field("text", pa.string()),
]
)
def _wal_hybrid(**select):
builder = (
LanceHybridQueryBuilder(_wal_table(_PK_SCHEMA)).vector([0.1, 0.2]).text("puppy")
)
if "columns" in select:
builder = builder.select(select["columns"])
return builder
def test_wal_hybrid_projects_the_key_instead_of_asking_for_row_ids():
builder = _wal_hybrid(columns=["text"])
builder._create_query_builders()
assert builder._pk_fusion.pk_columns == ["id"]
assert builder._pk_fusion.injected == ["id"]
for leg in (builder._vector_query, builder._fts_query):
assert leg._columns == ["text", "id"]
# A MemWAL table rejects with_row_id before it plans.
assert leg._with_row_id is not True
def test_wal_hybrid_selecting_the_key_injects_nothing():
builder = _wal_hybrid(columns=["text", "id"])
builder._create_query_builders()
assert builder._pk_fusion.injected == []
def test_wal_hybrid_without_a_projection_leaves_it_alone():
builder = _wal_hybrid()
builder._create_query_builders()
# Select-all already returns the key.
assert builder._pk_fusion.injected == []
assert builder._vector_query._columns is None
def test_base_table_hybrid_still_joins_on_row_ids():
table = mock.Mock()
table._hybrid_pk_fusion_learned.return_value = False
builder = LanceHybridQueryBuilder(table).vector([0.1, 0.2]).text("puppy")
builder._create_query_builders()
# `to_arrow` is what turns row ids on for this path; all that matters here
# is that no key was resolved and no projection was touched.
assert builder._pk_fusion is None
assert builder._vector_query._columns is None
def test_wal_hybrid_with_use_lsm_false_joins_on_row_ids():
builder = _wal_hybrid().use_lsm(False)
builder._create_query_builders()
assert builder._pk_fusion is None
def test_wal_hybrid_refuses_an_explicit_with_row_id():
builder = _wal_hybrid().with_row_id(True)
with pytest.raises(NotImplementedError, match="use_lsm"):
builder._create_query_builders()
def test_wal_hybrid_return_score_all_is_not_refused():
"""`return_score="all"` turns row ids on for the reranker, not the caller."""
builder = _wal_hybrid().rerank(RRFReranker(return_score="all"))
builder._create_query_builders()
assert builder._pk_fusion.pk_columns == ["id"]
def test_wal_hybrid_fuses_on_the_key_and_drops_it_again():
builder = _wal_hybrid(columns=["text"])
builder._create_query_builders()
vector_results = pa.table(
{"text": ["a", "b"], "id": ["a", "b"], "_distance": [0.1, 0.2]}
)
fts_results = pa.table({"text": ["b", "c"], "id": ["b", "c"], "_score": [1.0, 2.0]})
vector_results, fts_results = builder._pk_fusion.stamp(vector_results, fts_results)
combined = LanceHybridQueryBuilder._combine_hybrid_results(
fts_results=fts_results,
vector_results=vector_results,
norm="score",
fts_query="puppy",
reranker=RRFReranker(),
limit=10,
with_row_ids=True,
)
results = builder._finish_hybrid_results(combined)
assert "_rowid" not in results.column_names, "the surrogate is internal"
assert "id" not in results.column_names, "the caller selected only `text`"
assert sorted(results.column("text").to_pylist()) == ["a", "b", "c"]
@pytest.mark.parametrize(
"row_id_first", [True, False], ids=["row_id_first", "rerank_first"]
)
def test_wal_hybrid_refuses_row_id_even_with_return_score_all(row_id_first: bool):
"""`return_score="all"` sets the same flag the caller's request does, so
asking for both must still be refused — in either call order."""
builder = _wal_hybrid()
reranker = RRFReranker(return_score="all")
if row_id_first:
builder = builder.with_row_id(True).rerank(reranker)
else:
builder = builder.rerank(reranker).with_row_id(True)
with pytest.raises(NotImplementedError, match="use_lsm"):
builder._create_query_builders()
def test_wal_hybrid_refuses_a_blob_projection_with_return_score_all():
"""A projected blob needs a real row id to fetch through, which the
reranker turning `_with_row_id` on does not provide."""
schema = pa.schema(
[
pa.field(
"id",
pa.string(),
metadata={b"lance-schema:unenforced-primary-key:position": b"1"},
),
pa.field("text", pa.string()),
lancedb.blob("blob"),
]
)
builder = (
LanceHybridQueryBuilder(_wal_table(schema))
.vector([0.1, 0.2])
.text("puppy")
.select(["blob"])
.rerank(RRFReranker(return_score="all"))
)
with pytest.raises(NotImplementedError, match="blob"):
builder._create_query_builders()
_REFUSAL = RuntimeError(
"Bad request: the MemWAL LSM scanner does not support with_row_id"
)
def test_fallback_retries_on_the_key_and_remembers_it():
table = mock.Mock()
attempts = []
def run():
attempts.append(len(attempts))
if len(attempts) == 1:
raise _REFUSAL
return "fused"
result = with_pk_fallback(
run, table, fused=lambda: False, caller_requested_row_id=False
)
assert result == "fused"
assert len(attempts) == 2
table._note_hybrid_pk_fusion.assert_called_once()
def test_fallback_explains_the_refusal_to_a_caller_who_asked_for_row_ids():
def run():
raise _REFUSAL
with pytest.raises(NotImplementedError, match="use_lsm") as info:
with_pk_fallback(
run, mock.Mock(), fused=lambda: False, caller_requested_row_id=True
)
assert str(info.value) == WAL_ROW_ID_UNSUPPORTED
@pytest.mark.parametrize(
"error, fused",
[(ValueError("unrelated"), False), (_REFUSAL, True)],
ids=["other_error", "already_fused"],
)
def test_fallback_reraises_what_it_cannot_retry(error, fused):
def run():
raise error
table = mock.Mock()
with pytest.raises(type(error)):
with_pk_fallback(run, table, fused=lambda: fused, caller_requested_row_id=False)
table._note_hybrid_pk_fusion.assert_not_called()