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

Merge origin/main into gatekeeper/fix-3759-1

Browse Source
This commit is contained in:
Gatefixer
2026-08-28 01:19:19 +00:00
parent 43465f06d5 c94d9a2a16
commit 05fce57841
48 changed files with 2772 additions and 348 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/Cargo.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb-python"
version = "0.38.0-beta.10"
version = "0.38.0-beta.12"
publish = false
edition.workspace = true
description = "Python bindings for LanceDB"
python/python/lancedb/__init__.py
+15 -2
View File
@@ -6,7 +6,7 @@ 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")
@@ -20,7 +20,7 @@ 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 .functions import (
FunctionArtifactRequest as FunctionArtifactRequest,
@@ -49,6 +49,19 @@ from .namespace import (
)
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]]
) -> None:
python/python/lancedb/_blob.py
+9 -5
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(
@@ -270,7 +272,8 @@ def _iter_projection_pairs(
if isinstance(expr, str):
yield name, expr
elif isinstance(expr, Expr):
yield name, expr.to_sql()
source = expr._column_name()
yield name, source if source is not None else expr.to_sql()
return
for column in projection:
if isinstance(column, str):
@@ -280,7 +283,8 @@ def _iter_projection_pairs(
if isinstance(expr, str):
yield name, expr
elif isinstance(expr, Expr):
yield name, expr.to_sql()
source = expr._column_name()
yield name, source if source is not None else expr.to_sql()
def _set_blob_column(tbl: pa.Table, output_name: str, blobs: pa.Array) -> pa.Table:
python/python/lancedb/_lancedb.pyi
+2
View File
@@ -87,6 +87,7 @@ class PyExpr:
def contains(self, substr: "PyExpr") -> "PyExpr": ...
def isin(self, values: List["PyExpr"]) -> "PyExpr": ...
def cast(self, data_type: pa.DataType) -> "PyExpr": ...
def column_name(self) -> Optional[str]: ...
def to_sql(self) -> str: ...
def expr_col(name: str) -> PyExpr: ...
@@ -608,6 +609,7 @@ class PyQueryRequest:
filter: Optional[Union[str, bytes]]
full_text_search: Optional[FullTextQuery]
select: Optional[Union[str, List[str]]]
select_source_columns: Optional[Dict[str, str]]
fast_search: Optional[bool]
with_row_id: Optional[bool]
use_lsm: Optional[bool]
python/python/lancedb/expr.py
+5 -1
View File
@@ -249,6 +249,10 @@ class Expr:
# ── utilities ────────────────────────────────────────────────────────────
def _column_name(self) -> str | None:
"""Return the source name when this is a bare column expression."""
return self._inner.column_name()
def to_sql(self) -> str:
"""Render the expression as a SQL string (useful for debugging)."""
return self._inner.to_sql()
@@ -312,7 +316,7 @@ def func(name: str, *args: ExprLike) -> Expr:
--------
>>> from lancedb.expr import col, func
>>> func("lower", col("name"))
Expr(lower(name))
Expr(lower(`name`))
"""
inner_args = [_coerce(a)._inner for a in args]
return Expr(expr_func(name, inner_args))
python/python/lancedb/query.py
+12 -4
View File
@@ -167,6 +167,12 @@ def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
return {"columns": projection}
def _query_request_projection(req: "PyQueryRequest") -> QueryProjection:
if req.select_source_columns is not None:
return req.select_source_columns
return req.select
def _scanner_kwargs_for_query(
query: Query,
blob_mode: BlobMode,
@@ -2799,15 +2805,16 @@ class AsyncQueryBase(object):
req = self._inner.to_query_request()
schema = await self._table.schema()
projection = _query_request_projection(req)
self._blob_auto_row_id = blob_auto_row_id_for_scan(
schema,
req.select,
projection,
with_row_id=self._with_row_id,
)
if not self._blob_auto_row_id:
self._blob_paths = ()
return
self._blob_paths = tuple(blob_v2_projection_sources(schema, req.select).keys())
self._blob_paths = tuple(blob_v2_projection_sources(schema, projection).keys())
self._inner.with_row_id()
def select(self, columns: Union[List[str], dict[str, str]]) -> Self:
@@ -3894,14 +3901,15 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
blob_paths: tuple[str, ...] = ()
if self._table is not None:
schema = await self._table.schema()
projection = _query_request_projection(req)
blob_auto_row_id = blob_auto_row_id_for_scan(
schema,
req.select,
projection,
with_row_id=self._with_row_id,
)
if blob_auto_row_id:
blob_paths = tuple(
blob_v2_projection_sources(schema, req.select).keys()
blob_v2_projection_sources(schema, projection).keys()
)
self._blob_auto_row_id = blob_auto_row_id
self._blob_paths = blob_paths
python/python/lancedb/remote/table.py
+7 -4
View File
@@ -36,6 +36,7 @@ from lancedb._lancedb import (
UpdateResult,
)
from lancedb.embeddings.base import EmbeddingFunctionConfig
from lancedb.expr import Expr
from lancedb.index import (
FTS,
BTree,
@@ -863,7 +864,7 @@ class RemoteTable(Table):
def update(
self,
where: Optional[str] = None,
where: Optional[Union[str, Expr]] = None,
values: Optional[dict] = None,
*,
values_sql: Optional[Dict[str, str]] = None,
@@ -874,9 +875,11 @@ class RemoteTable(Table):
Parameters
----------
where: str, optional
The SQL where clause to use when updating rows. For example, 'x = 2'
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
error.
values: dict, optional
The values to update. The keys are the column names and the values
are the values to set.
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/table.py
+279 -105
View File
@@ -104,7 +104,12 @@ from .util import (
value_to_sql,
)
from .index import lang_mapping
from .schema import blob_v2_column_paths, schema_has_blob_field
from .schema import (
blob_v2_column_paths,
is_blob_v2_field,
row_addressable_blob_v2_paths,
schema_has_blob_field,
)
def _should_push_down_query_table(
@@ -426,6 +431,7 @@ def _cast_to_target_schema(
def gen():
for batch in reader:
batch = _coerce_blob_write_columns(batch, reordered_schema)
# Table but not RecordBatch has cast.
cast_batches = (
pa.Table.from_batches([batch]).cast(reordered_schema).to_batches()
@@ -438,6 +444,185 @@ def _cast_to_target_schema(
return pa.RecordBatchReader.from_batches(reordered_schema, gen())
def _coerce_blob_write_columns(
batch: pa.RecordBatch, target_schema: pa.Schema
) -> pa.RecordBatch:
"""Materialize blob storage structs before the stream leaves Python.
merge_insert requires its source reader to already match the table's
physical schema. Unlike add and insert, it does not pass through
LanceDB's Rust blob coercion, so preserving binary input here would
reach Lance as binary and fail the schema check.
"""
columns = []
fields = []
changed = False
for field, column in zip(batch.schema, batch.columns):
target_field = target_schema.field(field.name)
coerced = _coerce_blob_value(column, target_field)
if coerced is not column:
column = coerced
field = pa.field(
field.name,
coerced.type,
field.nullable,
target_field.metadata,
)
changed = True
columns.append(column)
fields.append(field)
if not changed:
return batch
return pa.RecordBatch.from_arrays(
columns, schema=pa.schema(fields, metadata=batch.schema.metadata)
)
def _coerce_blob_value(column: pa.Array, target_field: pa.Field) -> pa.Array:
if is_blob_v2_field(target_field) and _can_coerce_to_blob(column.type):
return _coerce_value_to_blob(column, target_field)
target_type = target_field.type
if pa.types.is_struct(target_type) and pa.types.is_struct(column.type):
children = []
fields = []
changed = False
for source_field in column.type:
source_column = column.field(source_field.name)
nested_target = next(
(field for field in target_type if field.name == source_field.name),
None,
)
if nested_target is None:
children.append(source_column)
fields.append(source_field)
continue
coerced = _coerce_blob_value(source_column, nested_target)
if coerced is not source_column:
changed = True
child_array, child_type = _physical_array_and_type(coerced)
children.append(child_array)
fields.append(
pa.field(
source_field.name,
child_type,
source_field.nullable,
nested_target.metadata,
)
)
if not changed:
return column
return pa.StructArray.from_arrays(
children,
fields=fields,
mask=column.is_null() if column.null_count else None,
)
if _is_list_like(target_type) and _is_list_like(column.type):
return _coerce_blob_list_values(column, target_type.value_field)
return column
def _coerce_blob_list_values(
column: pa.Array, target_value_field: pa.Field
) -> pa.Array:
"""Coerce blob values inside a list column, preserving offsets and nulls.
Works on the raw child values window instead of ``pc.list_flatten`` because
flatten drops values spanned by null slots, which would misalign offsets.
"""
mask = column.is_null() if column.null_count else None
if pa.types.is_fixed_size_list(column.type):
list_size = column.type.list_size
values = column.values.slice(column.offset * list_size, len(column) * list_size)
coerced = _coerce_blob_value(values, target_value_field)
if coerced is values:
return column
physical_values, _ = _physical_array_and_type(coerced)
return pa.FixedSizeListArray.from_arrays(physical_values, list_size, mask=mask)
offsets = column.offsets
first_offset = offsets[0].as_py()
values = column.values.slice(
first_offset,
offsets[-1].as_py() - first_offset,
)
coerced = _coerce_blob_value(values, target_value_field)
if coerced is values:
return column
physical_values, _ = _physical_array_and_type(coerced)
if first_offset:
offsets = pc.subtract(offsets, pa.scalar(first_offset, offsets.type))
if pa.types.is_large_list(column.type):
return pa.LargeListArray.from_arrays(offsets, physical_values, mask=mask)
return pa.ListArray.from_arrays(offsets, physical_values, mask=mask)
def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array:
if _is_string_like(values.type):
carrier_name = "uri"
carrier = values
elif pa.types.is_null(values.type):
carrier_name = None
carrier = None
elif pa.types.is_large_binary(values.type):
carrier_name = "data"
carrier = values
else:
carrier_name = "data"
carrier = values.cast(pa.large_binary())
length = len(values)
storage_type = target_field.type
if isinstance(storage_type, pa.ExtensionType):
storage_type = storage_type.storage_type
storage_fields = list(storage_type)
children = []
for storage_field in storage_fields:
if storage_field.name == carrier_name:
children.append(carrier.cast(storage_field.type))
else:
children.append(pa.nulls(length, type=storage_field.type))
storage = pa.StructArray.from_arrays(
children,
fields=storage_fields,
mask=values.is_null() if values.null_count else None,
)
if isinstance(target_field.type, pa.ExtensionType):
return pa.ExtensionArray.from_storage(target_field.type, storage)
return storage
def _physical_array_and_type(array: pa.Array) -> tuple[pa.Array, pa.DataType]:
if isinstance(array.type, pa.ExtensionType):
return array.storage, array.type.storage_type
return array, array.type
def _can_coerce_to_blob(data_type: pa.DataType) -> bool:
return (
_is_binary_like(data_type)
or _is_string_like(data_type)
or pa.types.is_null(data_type)
)
def _is_binary_like(data_type: pa.DataType) -> bool:
return (
pa.types.is_binary(data_type)
or pa.types.is_large_binary(data_type)
or pa.types.is_binary_view(data_type)
)
def _is_string_like(data_type: pa.DataType) -> bool:
predicates = ("is_string", "is_large_string", "is_string_view")
return any(
predicate(data_type)
for name in predicates
if (predicate := getattr(pa.types, name, None)) is not None
)
def _field_extension_name(field: pa.Field) -> Optional[str]:
extension_name = getattr(field.type, "extension_name", None)
if extension_name is not None:
@@ -633,29 +818,6 @@ def _prepare_extension_list(data: DATA, target_schema: pa.Schema) -> DATA:
return pa.Table.from_pylist(prepared_data, schema=insert_schema)
def _is_blob_source_field(field: pa.Field) -> bool:
if _field_extension_name(field) == _BLOB_EXTENSION_NAME:
return True
predicates = (
"is_binary",
"is_large_binary",
"is_binary_view",
"is_string",
"is_large_string",
"is_string_view",
)
if any(
predicate(field.type)
for name in predicates
if (predicate := getattr(pa.types, name, None)) is not None
):
return True
return pa.types.is_struct(field.type) and any(
child.name in {"data", "uri"} for child in field.type
)
def _align_field_types(
fields: List[pa.Field],
target_fields: List[pa.Field],
@@ -668,71 +830,71 @@ def _align_field_types(
target_field = next((f for f in target_fields if f.name == field.name), None)
if target_field is None:
raise ValueError(f"Field '{field.name}' not found in target schema")
target_extension_name = _field_extension_name(target_field)
# Preserve accepted blob carriers so Lance can construct the declared
# blob struct after optional Python preprocessing.
if target_extension_name == _BLOB_EXTENSION_NAME and _is_blob_source_field(
field
):
new_fields.append(field)
continue
# Preserve arrow.json input until it reaches Lance. LanceDB exposes stored
# JSON columns as lance.json (JSONB-backed LargeBinary), but casting the
# input to that storage type here merely relabels the raw JSON bytes as
# JSONB. Lance must see arrow.json so it can perform the JSONB encoding.
if (
_field_extension_name(field) == "arrow.json"
and target_extension_name in _JSON_EXTENSION_NAMES
):
new_fields.append(field)
continue
if pa.types.is_struct(target_field.type):
if pa.types.is_struct(field.type):
new_type = pa.struct(
_align_field_types(
field.type.fields,
target_field.type.fields,
)
new_fields.append(_align_field(field, target_field))
return new_fields
def _align_list_value_field(
value_field: pa.Field, target_value_field: pa.Field
) -> pa.Field:
# A list has exactly one child, so the inferred child name ("item") aligns
# positionally and adopts the table's child name; pa.Table.cast renames it.
return _align_field(value_field, target_value_field).with_name(
target_value_field.name
)
def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
# Preserve arrow.json input until it reaches Lance. LanceDB exposes stored
# JSON columns as lance.json (JSONB-backed LargeBinary), but casting the
# input to that storage type here merely relabels the raw JSON bytes as
# JSONB. Lance must see arrow.json so it can perform the JSONB encoding.
if (
_field_extension_name(field) == "arrow.json"
and _field_extension_name(target_field) == "lance.json"
):
return field
if pa.types.is_struct(target_field.type):
if pa.types.is_struct(field.type):
new_type = pa.struct(
_align_field_types(
field.type.fields,
target_field.type.fields,
)
else:
new_type = target_field.type
elif pa.types.is_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.list_(
_align_field_types(
[field.type.value_field],
[target_field.type.value_field],
)[0]
)
else:
new_type = target_field.type
elif pa.types.is_large_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.large_list(
_align_field_types(
[field.type.value_field],
[target_field.type.value_field],
)[0]
)
else:
new_type = target_field.type
elif pa.types.is_fixed_size_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.list_(
_align_field_types(
[field.type.value_field],
[target_field.type.value_field],
)[0],
target_field.type.list_size,
)
else:
new_type = target_field.type
)
else:
new_type = target_field.type
new_fields.append(
pa.field(field.name, new_type, field.nullable, target_field.metadata)
)
return new_fields
elif pa.types.is_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.list_(
_align_list_value_field(
field.type.value_field, target_field.type.value_field
)
)
else:
new_type = target_field.type
elif pa.types.is_large_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.large_list(
_align_list_value_field(
field.type.value_field, target_field.type.value_field
)
)
else:
new_type = target_field.type
elif pa.types.is_fixed_size_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.list_(
_align_list_value_field(
field.type.value_field, target_field.type.value_field
),
target_field.type.list_size,
)
else:
new_type = target_field.type
else:
new_type = target_field.type
return pa.field(field.name, new_type, field.nullable, target_field.metadata)
def _infer_subschema(
@@ -801,7 +963,7 @@ def sanitize_create_table(
schema = data.schema
else:
if schema is not None:
data = pa.Table.from_pylist([], schema)
data = pa.Table.from_batches([], schema=schema)
if schema is None:
if data is None:
raise ValueError("Either data or schema must be provided")
@@ -1956,7 +2118,7 @@ class Table(ABC):
@abstractmethod
def update(
self,
where: Optional[str] = None,
where: Optional[Union[str, Expr]] = None,
values: Optional[dict] = None,
*,
values_sql: Optional[Dict[str, str]] = None,
@@ -1971,9 +2133,11 @@ class Table(ABC):
Parameters
----------
where: str, optional
The SQL where clause to use when updating rows. For example, 'x = 2'
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
error.
values: dict, optional
The values to update. The keys are the column names and the values
are the values to set.
@@ -1991,6 +2155,7 @@ class Table(ABC):
Examples
--------
>>> import lancedb
>>> from lancedb.expr import col
>>> import pandas as pd
>>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]})
>>> db = lancedb.connect("./.lancedb")
@@ -2000,7 +2165,7 @@ class Table(ABC):
0 1 [1.0, 2.0]
1 2 [3.0, 4.0]
2 3 [5.0, 6.0]
>>> table.update(where="x = 2", values={"vector": [10.0, 10]})
>>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]})
UpdateResult(rows_updated=1, version=2)
>>> table.to_pandas()
x vector
@@ -2907,7 +3072,7 @@ class LanceTable(Table):
arrow_tbl = self.to_arrow()
if blob_mode == "descriptions":
arrow_tbl = strip_auto_row_ids(
arrow_tbl, blob_v2_column_paths(self.schema)
arrow_tbl, row_addressable_blob_v2_paths(self.schema)
)
return arrow_tbl.to_pandas(**kwargs)
@@ -4053,7 +4218,7 @@ class LanceTable(Table):
def update(
self,
where: Optional[str] = None,
where: Optional[Union[str, Expr]] = None,
values: Optional[dict] = None,
*,
values_sql: Optional[Dict[str, str]] = None,
@@ -4064,9 +4229,11 @@ class LanceTable(Table):
Parameters
----------
where: str, optional
The SQL where clause to use when updating rows. For example, 'x = 2'
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
error.
values: dict, optional
The values to update. The keys are the column names and the values
are the values to set.
@@ -4084,6 +4251,7 @@ class LanceTable(Table):
Examples
--------
>>> import lancedb
>>> from lancedb.expr import col
>>> import pandas as pd
>>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]})
>>> db = lancedb.connect("./.lancedb")
@@ -4093,7 +4261,7 @@ class LanceTable(Table):
0 1 [1.0, 2.0]
1 2 [3.0, 4.0]
2 3 [5.0, 6.0]
>>> table.update(where="x = 2", values={"vector": [10.0, 10]})
>>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]})
UpdateResult(rows_updated=1, version=2)
>>> table.to_pandas()
x vector
@@ -5308,7 +5476,9 @@ class AsyncTable:
if blob_mode == "descriptions" or not schema_has_blob_field(schema):
arrow_tbl = await self.to_arrow()
if blob_mode == "descriptions":
arrow_tbl = strip_auto_row_ids(arrow_tbl, blob_v2_column_paths(schema))
arrow_tbl = strip_auto_row_ids(
arrow_tbl, row_addressable_blob_v2_paths(schema)
)
return arrow_tbl.to_pandas(**kwargs)
if blob_mode == "lazy" and get_uri_scheme(await self.uri()) == "memory":
@@ -6210,7 +6380,7 @@ class AsyncTable:
self,
updates: Optional[Dict[str, Any]] = None,
*,
where: Optional[str] = None,
where: Optional[Union[str, Expr]] = None,
updates_sql: Optional[Dict[str, str]] = None,
) -> UpdateResult:
"""
@@ -6225,9 +6395,11 @@ class AsyncTable:
The updates to apply. The keys should be the name of the column to
update. The values should be the new values to assign. This is
required unless updates_sql is supplied.
where: str, optional
An SQL filter that controls which rows are updated. For example, 'x = 2'
or 'x IN (1, 2, 3)'. Only rows that satisfy this filter will be udpated.
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. Only rows that satisfy this filter will
be updated.
updates_sql: dict, optional
The updates to apply, expressed as SQL expression strings. The keys should
be column names. The values should be SQL expressions. These can be SQL
@@ -6245,13 +6417,14 @@ class AsyncTable:
--------
>>> import asyncio
>>> import lancedb
>>> from lancedb.expr import col
>>> import pandas as pd
>>> async def demo_update():
... data = pd.DataFrame({"x": [1, 2], "vector": [[1, 2], [3, 4]]})
... db = await lancedb.connect_async("./.lancedb")
... table = await db.create_table("my_table", data)
... # x is [1, 2], vector is [[1, 2], [3, 4]]
... await table.update({"vector": [10, 10]}, where="x = 2")
... await table.update({"vector": [10, 10]}, where=col("x") == 2)
... # x is [1, 2], vector is [[1, 2], [10, 10]]
... await table.update(updates_sql={"x": "x + 1"})
... # x is [2, 3], vector is [[1, 2], [10, 10]]
@@ -6265,7 +6438,8 @@ class AsyncTable:
if updates is not None:
updates_sql = {k: value_to_sql(v) for k, v in updates.items()}
return await self._inner.update(updates_sql, where)
predicate = where.to_sql() if isinstance(where, Expr) else where
return await self._inner.update(updates_sql, predicate)
async def add_columns(
self,
python/python/tests/test_blob.py
+552 -1
View File
@@ -2,17 +2,41 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import io
import subprocess
import sys
import textwrap
import lance
import pyarrow as pa
import pyarrow.compute as pc
import pytest
from lance.blob import BlobType as LanceBlobType
import lancedb
from lancedb._blob import read_row_ids_from_hits, stash_auto_row_ids
from lancedb._blob import (
blob_v2_projection_sources,
read_row_ids_from_hits,
stash_auto_row_ids,
)
from lancedb.expr import col
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")])
@@ -46,6 +70,181 @@ 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():
@@ -70,6 +269,14 @@ def test_blob_v2_column_paths_include_list_children():
]
def test_blob_v2_projection_sources_use_typed_column_name():
schema = pa.schema([lancedb.blob("blob")])
assert blob_v2_projection_sources(schema, {"blob_alias": col("blob")}) == {
"blob_alias": "blob"
}
def _legacy_v1_table(name):
db = lancedb.connect("memory:///")
schema = pa.schema(
@@ -166,6 +373,20 @@ async def test_async_table_to_pandas_descriptions_mode_omits_row_id():
assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}
@pytest.mark.asyncio
async def test_async_typed_blob_projection_preserves_source_column():
db = await lancedb.connect_async("memory:///typed_blob_projection")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
table = await db.create_table("typed_blob_projection", schema=schema)
await table.add([{"id": 1, "blob": b"alpha"}])
hits = await table.query().select({"blob_alias": col("blob")}).to_arrow()
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
blobs = await table.fetch_blobs("blob", hits)
assert blobs.to_pylist() == [b"alpha"]
def test_fetch_blobs_round_trip():
table = _blob_table(
"round_trip",
@@ -176,6 +397,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()
@@ -477,6 +984,50 @@ async def test_blob_v2_hybrid_fetch_blobs_async():
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
@pytest.mark.asyncio
async def test_async_hybrid_typed_blob_projection_preserves_source_column():
db = await lancedb.connect_async("memory:///hybrid_typed_blob")
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("text", pa.utf8()),
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
lancedb.blob("blob"),
]
)
table = await db.create_table("hybrid_typed_blob", schema=schema)
await table.add(
[
{
"id": 1,
"text": "hello alpha",
"vector": [1.0, 0.0],
"blob": b"alpha",
},
{
"id": 2,
"text": "hello beta",
"vector": [0.9, 0.1],
"blob": b"beta",
},
]
)
await table.create_index("text", config=FTS(with_position=False))
hits = await (
table.query()
.nearest_to([1.0, 0.0])
.nearest_to_text("hello")
.select({"blob_alias": col("blob")})
.limit(2)
.to_arrow()
)
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
blobs = await table.fetch_blobs("blob", hits)
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
def test_blob_file_seek_read_and_read_range():
payload = _identifiable_payload(1024)
table = _blob_table("seek_read", [{"id": 1, "image": payload}])
python/python/tests/test_expr.py
+21 -21
View File
@@ -52,7 +52,7 @@ class TestExprConstruction:
def test_func(self):
e = func("lower", col("name"))
assert isinstance(e, Expr)
assert e.to_sql() == "lower(name)"
assert e.to_sql() == "lower(`name`)"
def test_func_unknown_raises(self):
with pytest.raises(Exception):
@@ -115,7 +115,7 @@ class TestExprOperators:
def test_and_operator(self):
e = (col("age") > lit(18)) & (col("status") == lit("active"))
assert isinstance(e, Expr)
assert e.to_sql() == "((age > 18) AND (status = 'active'))"
assert e.to_sql() == "((age > 18) AND (`status` = 'active'))"
def test_or_operator(self):
e = (col("a") == lit(1)) | (col("b") == lit(2))
@@ -166,7 +166,7 @@ class TestExprOperators:
def test_coerce_plain_str(self):
e = col("name") == "alice"
assert isinstance(e, Expr)
assert e.to_sql() == "(name = 'alice')"
assert e.to_sql() == "(`name` = 'alice')"
def test_reflexive_comparisons(self):
# 10 < col("age") swaps to col("age") > 10
@@ -198,85 +198,85 @@ class TestExprBytesLiteral:
def test_bytes_equality_expr_sql(self):
e = col("data") == lit(b"\xca\xfe")
assert e.to_sql() == "(data = X'CAFE')"
assert e.to_sql() == "(`data` = X'CAFE')"
def test_bytes_ne_expr_sql(self):
e = col("data") != lit(b"\xff")
assert e.to_sql() == "(data <> X'FF')"
assert e.to_sql() == "(`data` <> X'FF')"
def test_bytes_compound_expr_sql(self):
e = (col("data") == lit(b"\x01")) & (col("id") > lit(5))
assert e.to_sql() == "((data = X'01') AND (id > 5))"
assert e.to_sql() == "((`data` = X'01') AND (id > 5))"
def test_bytes_in_function_call(self):
# Regression test: binary literals inside scalar function calls
# used to fail because DataFusion's unparser does not support Binary
# scalars. Now handled via a placeholder-substitution rewrite.
e = func("contains", col("data"), lit(b"\xff"))
assert e.to_sql() == "contains(data, X'FF')"
assert e.to_sql() == "contains(`data`, X'FF')"
def test_bytes_in_not(self):
e = ~(col("data") == lit(b"\xff"))
assert e.to_sql() == "NOT (data = X'FF')"
assert e.to_sql() == "NOT (`data` = X'FF')"
class TestExprStringMethods:
def test_lower(self):
e = col("name").lower()
assert isinstance(e, Expr)
assert e.to_sql() == "lower(name)"
assert e.to_sql() == "lower(`name`)"
def test_upper(self):
e = col("name").upper()
assert isinstance(e, Expr)
assert e.to_sql() == "upper(name)"
assert e.to_sql() == "upper(`name`)"
def test_contains(self):
e = col("text").contains(lit("hello"))
assert isinstance(e, Expr)
assert e.to_sql() == "contains(text, 'hello')"
assert e.to_sql() == "contains(`text`, 'hello')"
def test_contains_with_str_coerce(self):
e = col("text").contains("hello")
assert isinstance(e, Expr)
assert e.to_sql() == "contains(text, 'hello')"
assert e.to_sql() == "contains(`text`, 'hello')"
def test_chained_lower_eq(self):
e = col("name").lower() == lit("alice")
assert isinstance(e, Expr)
assert e.to_sql() == "(lower(name) = 'alice')"
assert e.to_sql() == "(lower(`name`) = 'alice')"
class TestExprCast:
def test_cast_string(self):
e = col("id").cast("string")
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(id AS VARCHAR)"
assert e.to_sql() == "arrow_cast(id, 'Utf8')"
def test_cast_int32(self):
e = col("score").cast("int32")
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(score AS INTEGER)"
assert e.to_sql() == "arrow_cast(score, 'Int32')"
def test_cast_float64(self):
e = col("val").cast("float64")
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(val AS DOUBLE)"
assert e.to_sql() == "arrow_cast(val, 'Float64')"
def test_cast_pyarrow_type(self):
e = col("score").cast(pa.int32())
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(score AS INTEGER)"
assert e.to_sql() == "arrow_cast(score, 'Int32')"
def test_cast_pyarrow_float64(self):
e = col("val").cast(pa.float64())
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(val AS DOUBLE)"
assert e.to_sql() == "arrow_cast(val, 'Float64')"
def test_cast_pyarrow_string(self):
e = col("id").cast(pa.string())
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(id AS VARCHAR)"
assert e.to_sql() == "arrow_cast(id, 'Utf8')"
def test_cast_pyarrow_and_string_equivalent(self):
# pa.int32() and "int32" should produce equivalent SQL
@@ -597,14 +597,14 @@ class TestExprIsin:
def test_isin_strs(self):
assert (
col("status").isin(["active", "pending"]).to_sql()
== "status IN ('active', 'pending')"
== "`status` IN ('active', 'pending')"
)
def test_isin_coerces_and_mixes(self):
assert col("id").isin([lit(1), 2]).to_sql() == "id IN (1, 2)"
def test_isin_empty(self):
assert col("id").isin([]).to_sql() == "id IN ()"
assert col("id").isin([]).to_sql() == "false"
def test_isin_filter(self, simple_table):
result = simple_table.search().where(col("id").isin([1, 3, 5])).to_arrow()
python/python/tests/test_query.py
+15
View File
@@ -675,6 +675,21 @@ def test_distance_range(table: lancedb.table.Table):
assert res["_distance"].to_pylist() == [min_dist, max_dist]
@pytest.mark.parametrize("expression", ["1 - _distance", "1.0 - _distance"])
def test_select_arithmetic_with_distance(table, expression):
result = (
table.search([10, 10])
.select({"similarity": expression, "_distance": "_distance"})
.distance_type("cosine")
.to_arrow()
)
assert result.schema.field("similarity").type == pa.float32()
assert result["similarity"].to_pylist() == pytest.approx(
[1 - distance for distance in result["_distance"].to_pylist()]
)
@pytest.mark.asyncio
async def test_distance_range_async(table_async: AsyncTable):
q = [0, 0]
python/python/tests/test_table.py
+158
View File
@@ -11,6 +11,7 @@ import warnings
import weakref
from concurrent.futures import ThreadPoolExecutor
from datetime import date, datetime, timedelta
from decimal import Decimal
from time import sleep
from typing import List
from unittest.mock import patch
@@ -336,6 +337,21 @@ async def test_update_async(mem_db_async: AsyncConnection):
assert await table.count_rows("id == 10") == 1
@pytest.mark.asyncio
async def test_update_expr_filter_literals_async(mem_db_async: AsyncConnection):
values = ["5", "4.66e-84", "it's"]
table = await mem_db_async.create_table(
"update_expr_literals",
data=[{"field": value, "result": "original"} for value in values],
)
for value in values:
update_res = await table.update({"result": value}, where=col("field") == value)
assert update_res.rows_updated == 1
assert (await table.to_arrow())["result"].to_pylist() == values
def test_create_table(mem_db: DBConnection):
schema = pa.schema(
{
@@ -2392,6 +2408,148 @@ def test_update(mem_db: DBConnection):
assert np.allclose(v, np.array([[1.2, 1.9], [1.1, 1.1]]))
def test_update_expr_filter_literals(mem_db: DBConnection):
values = ["5", "4.66e-84", "it's"]
table = mem_db.create_table(
"update_expr_literals",
data=[{"field": value, "result": "original"} for value in values],
)
for value in values:
update_res = table.update(where=col("field") == value, values={"result": value})
assert update_res.rows_updated == 1
assert table.to_arrow()["result"].to_pylist() == values
def test_update_expr_filter_preserves_typed_semantics(mem_db: DBConnection):
low = Decimal("1.234567890123456789")
high = Decimal("1.234567890123456790")
decimal_schema = pa.schema(
[("val", pa.decimal128(19, 18)), ("result", pa.string())]
)
decimal_table = mem_db.create_table(
"update_expr_decimal",
pa.table(
{"val": [low, high], "result": ["old", "old"]},
schema=decimal_schema,
),
)
predicate = col("val") < lit(high)
assert decimal_table.search().where(predicate).to_arrow().num_rows == 1
result = decimal_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
keyword_table = mem_db.create_table(
"update_expr_keyword", [{"null": 1, "result": "old"}]
)
predicate = col("null") == 1
assert keyword_table.search().where(predicate).to_arrow().num_rows == 1
result = keyword_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
empty_in_table = mem_db.create_table(
"update_expr_empty_in", [{"id": 1, "result": "old"}]
)
predicate = col("id").isin([])
assert empty_in_table.search().where(predicate).to_arrow().num_rows == 0
result = empty_in_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 0
marker = "__lancedb_binary_placeholder_0__"
binary_schema = pa.schema(
[("payload", pa.binary()), ("text", pa.string()), ("result", pa.string())]
)
binary_table = mem_db.create_table(
"update_expr_binary",
pa.table(
{
"payload": [b"\x01", b"\x02"],
"text": ["other", marker],
"result": ["old", "old"],
},
schema=binary_schema,
),
)
predicate = (col("payload") == lit(b"\x01")) | (col("text") == marker)
assert binary_table.search().where(predicate).to_arrow().num_rows == 2
result = binary_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 2
nonfinite_table = mem_db.create_table(
"update_expr_nonfinite",
[{"x": 1.0, "result": "old"}, {"x": 2.0, "result": "old"}],
)
predicate = col("x") < float("inf")
assert nonfinite_table.search().where(predicate).to_arrow().num_rows == 2
result = nonfinite_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 2
float16_table = mem_db.create_table(
"update_expr_float16",
[{"x": 1.0, "result": "old"}, {"x": 3.0, "result": "old"}],
)
predicate = col("x").cast(pa.float16()) < 2.0
assert float16_table.search().where(predicate).to_arrow().num_rows == 1
result = float16_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
string_cast_table = mem_db.create_table(
"update_expr_string_cast",
[{"x": 1, "result": "old"}, {"x": 2, "result": "old"}],
)
predicate = col("x").cast("string") == "1"
assert string_cast_table.search().where(predicate).to_arrow().num_rows == 1
result = string_cast_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
quoted_identifier_schema = pa.schema(
[("payload", pa.binary()), ("odd'name", pa.int64()), ("result", pa.string())]
)
quoted_identifier_table = mem_db.create_table(
"update_expr_quoted_identifier",
pa.table(
{"payload": [b"\x01"], "odd'name": [1], "result": ["old"]},
schema=quoted_identifier_schema,
),
)
predicate = (col("payload") == lit(b"\x01")) & (col("odd'name") == 1)
assert quoted_identifier_table.search().where(predicate).to_arrow().num_rows == 1
result = quoted_identifier_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
decimal256_schema = pa.schema(
[("val", pa.decimal256(40, 2)), ("result", pa.string())]
)
decimal256_table = mem_db.create_table(
"update_expr_decimal256",
pa.table(
{
"val": [Decimal("1.00"), Decimal("3.00")],
"result": ["old", "old"],
},
schema=decimal256_schema,
),
)
predicate = col("val") < lit(Decimal("2.00")).cast(pa.decimal256(40, 2))
assert decimal256_table.search().where(predicate).to_arrow().num_rows == 1
result = decimal256_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
binary_empty_table = mem_db.create_table(
"update_expr_binary_empty",
pa.table(
{"payload": [b"\x01", b"\x02"], "result": ["old", "old"]},
schema=pa.schema([("payload", pa.binary()), ("result", pa.string())]),
),
)
predicate = (col("payload") == lit(b"\x01")).isin([])
assert binary_empty_table.search().where(predicate).to_arrow().num_rows == 0
assert predicate.to_sql() == "false"
result = binary_empty_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 0
def test_update_with_arrow_scalar(mem_db: DBConnection):
schema = pa.schema({"id": pa.int64(), "vector": pa.list_(pa.float32(), 4)})
table = mem_db.create_table("my_table", schema=schema)
python/python/tests/test_util.py
+160
View File
@@ -7,6 +7,7 @@ import pathlib
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
@@ -907,6 +908,165 @@ 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_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/src/expr.rs
+8
View File
@@ -130,6 +130,14 @@ impl PyExpr {
// ── utilities ────────────────────────────────────────────────────────────
/// Return the referenced column name for a bare column expression.
fn column_name(&self) -> Option<String> {
match &self.0 {
DfExpr::Column(column) if column.relation.is_none() => Some(column.name.clone()),
_ => None,
}
}
/// Render the expression as a SQL string (useful for debugging).
fn to_sql(&self) -> PyResult<String> {
lancedb::expr::expr_to_sql_string(&self.0).map_err(|e| PyValueError::new_err(e.to_string()))
python/src/query.rs
+23
View File
@@ -1,6 +1,7 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
@@ -325,6 +326,7 @@ pub struct PyQueryRequest {
pub filter: Option<PyQueryFilter>,
pub full_text_search: Option<PyLanceDB<FtsQuery>>,
pub select: PySelect,
pub select_source_columns: Option<HashMap<String, String>>,
pub fast_search: Option<bool>,
pub with_row_id: Option<bool>,
pub use_lsm: Option<bool>,
@@ -355,6 +357,7 @@ impl From<AnyQuery> for PyQueryRequest {
full_text_search: query_request
.full_text_search
.map(|fts| PyLanceDB(fts.query)),
select_source_columns: PySelect::source_columns(&query_request.select),
select: PySelect(query_request.select),
fast_search: Some(query_request.fast_search),
with_row_id: Some(query_request.with_row_id),
@@ -380,6 +383,7 @@ impl From<AnyQuery> for PyQueryRequest {
offset: vector_query.base.offset,
filter: vector_query.base.filter.map(PyQueryFilter),
full_text_search: None,
select_source_columns: PySelect::source_columns(&vector_query.base.select),
select: PySelect(vector_query.base.select),
fast_search: Some(vector_query.base.fast_search),
with_row_id: Some(vector_query.base.with_row_id),
@@ -412,6 +416,25 @@ impl From<AnyQuery> for PyQueryRequest {
#[derive(Clone)]
pub struct PySelect(Select);
impl PySelect {
fn source_columns(select: &Select) -> Option<HashMap<String, String>> {
match select {
Select::Expr(pairs) => Some(
pairs
.iter()
.filter_map(|(output, expr)| match expr {
lancedb::expr::DfExpr::Column(column) if column.relation.is_none() => {
Some((output.clone(), column.name.clone()))
}
_ => None,
})
.collect(),
),
_ => None,
}
}
}
impl<'py> IntoPyObject<'py> for PySelect {
type Target = PyAny;
type Output = Bound<'py, Self::Target>;