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

Merge remote-tracking branch 'refs/remotes/origin/main' into gatekeeper/fix-2820-1

Browse Source 
# Conflicts:
#	rust/lancedb/src/remote/table.rs
This commit is contained in:
Gatefixer
2026-08-26 05:43:48 +00:00
parent f936f65626 391cac9034
commit 029930b412
98 changed files with 8235 additions and 1295 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.6"
version = "0.38.0-beta.10"
publish = false
edition.workspace = true
description = "Python bindings for LanceDB"
python/pyproject.toml
+4 -1
View File
@@ -101,9 +101,12 @@ azure = ["adlfs>=2024.2.0"]
[tool.maturin]
python-source = "python"
module-name = "lancedb._lancedb"
# uv installs the project as an editable package before `uv run`, so keep that
# bootstrap build consistent with `maturin develop`.
editable-profile = "dev"
[build-system]
requires = ["maturin>=1.9.4"]
requires = ["maturin>=1.10"]
build-backend = "maturin"
[tool.ruff.lint]
python/python/lancedb/__init__.py
+21
View File
@@ -179,6 +179,18 @@ def connect(
... },
... )
For Azure Blob Storage, credentials can be passed directly without setting
environment variables:
>>> azure_storage_options = {
... "account_name": "some-account",
... "account_key": "some-key",
... }
>>> db = lancedb.connect( # doctest: +SKIP
... "az://my-container/my-database",
... storage_options=azure_storage_options,
... )
For tests and temporary data, use an in-memory database:
>>> db = lancedb.connect("memory://")
@@ -465,6 +477,10 @@ async def connect_async(
--------
>>> import lancedb
>>> azure_storage_options = {
... "account_name": "some-account",
... "account_key": "some-key",
... }
>>> async def doctest_example():
... # For a local directory, provide a path to the database
... db = await lancedb.connect_async("~/.lancedb")
@@ -472,6 +488,11 @@ async def connect_async(
... db = await lancedb.connect_async("s3://my-bucket/lancedb",
... storage_options={
... "aws_access_key_id": "***"})
... # Azure credentials can also be passed directly
... db = await lancedb.connect_async(
... "az://my-container/my-database",
... storage_options=azure_storage_options,
... )
... # For tests and temporary data, use an in-memory database
... db = await lancedb.connect_async("memory://")
... # Connect to LanceDB cloud
python/python/lancedb/_lancedb.pyi
+1
View File
@@ -283,6 +283,7 @@ class Table:
mode: Literal["append", "overwrite"],
progress: Optional[Any] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult: ...
async def update(
self, updates: Dict[str, str], where: Optional[str]
python/python/lancedb/functions.py
+197 -99
View File
@@ -4,7 +4,7 @@
"""Canonical Function values exchanged with LanceDB Enterprise services.
These immutable models contain client/wire state only. Catalog persistence,
environment bake, secret resolution, and execution are owned by Sophon.
environment bake, and execution are owned by Sophon.
``RefreshColumnResult`` is also the backend-neutral result of a local
expression-backed refresh job.
"""
@@ -12,17 +12,19 @@ expression-backed refresh job.
from __future__ import annotations
import ast
import builtins
import base64
import functools
import hashlib
import importlib
import inspect
import symtable
import json
import math
import re
import sys
import textwrap
import types
import uuid
from collections.abc import Mapping
from datetime import date, datetime
from typing import (
@@ -226,7 +228,7 @@ class PythonEnvironmentSpec(_RemoteValue):
class PythonRuntimeSpec(_RemoteValue):
"""Remote runtime definition with non-secret environment values.
"""Remote runtime definition with environment values.
V1 supports ``kind="python"``. Newer runtime kinds remain readable, while
their unknown payload fields are intentionally not retained by the client.
@@ -265,7 +267,6 @@ class FunctionVersion(_RemoteValue):
runtime: PythonRuntimeSpec
runtime_digest: str
environment_digest: str
required_secrets: tuple[str, ...] = ()
created_at: str
def __call__(self, **inputs: Any) -> FunctionApplication:
@@ -273,7 +274,7 @@ class FunctionVersion(_RemoteValue):
Every input must be a direct [lancedb.col][lancedb.expr.col]
reference. The returned application is immutable and retains a
named-struct output as one sibling group, so every row's sibling values
named-struct output as one binding, so every row's sibling values
come from one logical Function evaluation. Map result fields to table
columns with
[FunctionApplication.rename][lancedb.functions.FunctionApplication.rename],
@@ -323,22 +324,16 @@ class FunctionVersion(_RemoteValue):
function=FunctionVersionRef(name=self.name, version=self.version),
inputs=tuple(bindings),
output=self.signature.output,
group_id=f"fg_{uuid.uuid4().hex}",
)
class FunctionRegistrationRequest(_RemoteValue):
"""Stable remote registration envelope produced by :func:`udf`.
Only secret names are represented. Secret values are resolved inside the
remote service and have no client request field.
"""
"""Stable remote registration envelope produced by :func:`udf`."""
name: str
artifact: FunctionArtifactRequest
signature: FunctionSignature
runtime: PythonRuntimeSpec
required_secrets: tuple[str, ...] = ()
class FunctionVersionRef(_OpenRemoteValue):
@@ -367,7 +362,7 @@ class ApplicationInput(_OpenRemoteValue):
class FunctionApplication(_OpenRemoteValue):
"""Immutable pre-declaration application of an exact Function version.
A named-struct output remains one grouped application through table
A named-struct output remains one application through table
declaration and execution.
[FunctionApplication.rename][lancedb.functions.FunctionApplication.rename]
records the result-field to table-column mapping without splitting sibling
@@ -377,7 +372,6 @@ class FunctionApplication(_OpenRemoteValue):
function: FunctionVersionRef
inputs: tuple[ApplicationInput, ...]
output: FunctionOutput
group_id: str
columns: Mapping[str, str] = Field(default_factory=dict)
def _known_dict(self) -> dict[str, Any]:
@@ -449,12 +443,10 @@ class OutputMapping(_RemoteValue):
class FunctionBinding(_RemoteValue):
"""Immutable grouped binding persisted by the Enterprise table service."""
"""Immutable Function binding persisted by the Enterprise table service."""
binding_id: str
revision: _UInt64
function: FunctionVersionRef
group_id: str
inputs: tuple[InputBinding, ...]
outputs: tuple[OutputMapping, ...]
input_schema: Optional[Mapping[str, Any]] = None
@@ -486,62 +478,60 @@ class RefreshColumnResult(_RemoteValue):
_FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")
_SECRET_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
_GRAMMAR_PRIMITIVES = (
(pa.bool_(), "bool"),
(pa.int8(), "int8"),
(pa.int16(), "int16"),
(pa.int32(), "int32"),
(pa.int64(), "int64"),
(pa.uint8(), "uint8"),
(pa.uint16(), "uint16"),
(pa.uint32(), "uint32"),
(pa.uint64(), "uint64"),
(pa.float16(), "float16"),
(pa.float32(), "float32"),
(pa.float64(), "float64"),
(pa.string(), "utf8"),
(pa.binary(), "binary"),
(pa.date32(), "date32"),
(pa.date64(), "date64"),
)
def _canonical_arrow_type(data_type: pa.DataType) -> str:
primitive_types = (
(pa.bool_(), "bool"),
(pa.int8(), "int8"),
(pa.int16(), "int16"),
(pa.int32(), "int32"),
(pa.int64(), "int64"),
(pa.uint8(), "uint8"),
(pa.uint16(), "uint16"),
(pa.uint32(), "uint32"),
(pa.uint64(), "uint64"),
(pa.float16(), "float16"),
(pa.float32(), "float32"),
(pa.float64(), "float64"),
(pa.string(), "utf8"),
(pa.large_utf8(), "large_utf8"),
(pa.binary(), "binary"),
(pa.large_binary(), "large_binary"),
(pa.date32(), "date32"),
(pa.date64(), "date64"),
)
for candidate, name in primitive_types:
"""The server's V1 Function type grammar. Anything outside it is rejected
here rather than at registration."""
for candidate, name in _GRAMMAR_PRIMITIVES:
if data_type == candidate:
return name
if pa.types.is_fixed_size_binary(data_type):
return f"fixed_size_binary[{data_type.byte_width}]"
if pa.types.is_list(data_type):
return f"list<{_canonical_arrow_type(data_type.value_type)}>"
if pa.types.is_large_list(data_type):
return f"large_list<{_canonical_arrow_type(data_type.value_type)}>"
if pa.types.is_fixed_size_list(data_type):
if pa.types.is_list(data_type) or pa.types.is_large_list(data_type):
prefix = "list" if pa.types.is_list(data_type) else "large_list"
return f"{prefix}<{_canonical_list_item(data_type)}>"
if pa.types.is_fixed_size_list(data_type) and data_type.list_size > 0:
return (
f"fixed_size_list<{_canonical_arrow_type(data_type.value_type)}>"
f"[{data_type.list_size}]"
f"fixed_size_list<{_canonical_list_item(data_type)}, {data_type.list_size}>"
)
if pa.types.is_struct(data_type):
fields = ",".join(
f"{field.name}:{_canonical_arrow_type(field.type)}" for field in data_type
)
return f"struct<{fields}>"
if pa.types.is_timestamp(data_type):
timezone = f",tz={data_type.tz}" if data_type.tz is not None else ""
return f"timestamp[{data_type.unit}{timezone}]"
if pa.types.is_time32(data_type) or pa.types.is_time64(data_type):
return f"time[{data_type.unit}]"
if pa.types.is_duration(data_type):
return f"duration[{data_type.unit}]"
if pa.types.is_decimal(data_type):
bit_width = data_type.bit_width
return f"decimal{bit_width}({data_type.precision},{data_type.scale})"
raise TypeError(f"unsupported Arrow type for Function signature: {data_type}")
def _canonical_list_item(data_type: pa.DataType) -> str:
"""The grammar names only the item type; it always means a non-nullable
child called `item`, so any other child metadata cannot be represented."""
child = data_type.value_field
if child.name != "item" or child.nullable or child.metadata:
raise TypeError(
"unsupported Arrow type for Function signature: list items must be a "
f"non-nullable field named 'item', got {child}"
)
return _canonical_arrow_type(child.type)
def _list_of(item: pa.DataType) -> pa.DataType:
return pa.list_(pa.field("item", item, nullable=False))
def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]:
nullable = False
origin = get_origin(annotation)
@@ -589,7 +579,7 @@ def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]:
value_type, value_nullable = _annotation_type(arguments[0])
if value_nullable:
raise TypeError("nullable Function list elements are not supported")
return pa.list_(value_type), nullable
return _list_of(value_type), nullable
raise TypeError(f"unsupported Function annotation: {annotation!r}")
@@ -736,6 +726,104 @@ def _literal_source(value: Any) -> str:
)
_DYNAMIC_NAMESPACE_ACCESS = frozenset(
{"globals", "locals", "vars", "eval", "exec", "compile", "__import__"}
)
# Modules that hand out namespaces (`sys.modules`, `builtins`, importers,
# introspection). The artifact's module namespace holds only the names it was
# packaged with, so reaching around it cannot be represented.
_NAMESPACE_MODULES = frozenset(
{"sys", "builtins", "importlib", "inspect", "gc", "ctypes", "types"}
)
def _namespace_acquisition(
definition: ast.FunctionDef, references: set[str]
) -> list[str]:
found = set(references & _DYNAMIC_NAMESPACE_ACCESS)
for node in ast.walk(definition):
if isinstance(node, ast.Import):
found.update(
alias.name
for alias in node.names
if alias.name.split(".")[0] in _NAMESPACE_MODULES
)
elif isinstance(node, ast.ImportFrom) and node.module:
if node.module.split(".")[0] in _NAMESPACE_MODULES:
found.add(node.module)
return sorted(found)
def _module_references(module_source: str) -> set[str]:
"""Names any scope in `module_source` binds or loads at module scope.
Python's own scope analysis on the exact text that ships: free variables
belong to an enclosing scope inside the function, and postponed
annotations are not runtime loads."""
def visit(table: symtable.SymbolTable, found: set[str]) -> None:
for symbol in table.get_symbols():
if symbol.is_global() and (
symbol.is_referenced() or symbol.is_declared_global()
):
found.add(symbol.get_name())
for child in table.get_children():
visit(child, found)
found: set[str] = set()
for table in symtable.symtable(module_source, "<udf>", "exec").get_children():
visit(table, found)
return found
def _global_source(name: str, value: Any) -> str:
"""One module-level line that rebinds `name` to `value` in the artifact:
an import for modules and importable classes/functions, a literal otherwise."""
if isinstance(value, types.ModuleType):
if value.__name__.split(".")[0] in _NAMESPACE_MODULES:
raise ValueError(
f"@udf cannot package dynamic namespace access: {value.__name__!r}"
)
try:
imported = importlib.import_module(value.__name__)
except ImportError:
imported = None
if imported is not value:
raise TypeError(
f"Function source references module {name!r} that does not import "
f"as {value.__name__!r}"
)
return f"import {value.__name__} as {name}"
module_name = getattr(value, "__module__", None)
qualname = getattr(value, "__qualname__", None)
if (
isinstance(module_name, str)
and isinstance(qualname, str)
and module_name != "__main__"
and "." not in qualname
and "<" not in qualname
):
try:
imported = getattr(importlib.import_module(module_name), qualname)
except (ImportError, AttributeError):
imported = None
if imported is value:
return f"from {module_name} import {qualname} as {name}"
return f"{name} = {_literal_source(value)}"
def _is_recursive_reference(function: Callable[..., Any], name: str) -> bool:
"""`name` inside the body means the function itself unless the module has
since bound it to something else."""
if name != function.__name__:
return False
bound = function.__globals__.get(name, function)
if bound is function:
return True
# The decorator's own result is the one wrapper known to call `function`
# unchanged; any other binding may behave differently from a self-call.
return type(bound) is UdfDefinition and bound._function is function
def _package_source(function: Callable[..., Any]) -> bytes:
if not inspect.isfunction(function) or inspect.iscoroutinefunction(function):
raise TypeError("@udf requires a synchronous Python function")
@@ -760,23 +848,46 @@ def _package_source(function: Callable[..., Any]) -> bytes:
closure = inspect.getclosurevars(function)
if closure.nonlocals:
raise ValueError("@udf cannot package functions that capture closure values")
if closure.unbound:
raise ValueError(
f"@udf source contains unresolved global names: {sorted(closure.unbound)!r}"
)
globals_source = []
for name, value in sorted(closure.globals.items()):
if isinstance(value, types.ModuleType):
globals_source.append(f"import {value.__name__} as {name}")
else:
globals_source.append(f"{name} = {_literal_source(value)}")
function_source = ast.unparse(definition)
parts = ["from __future__ import annotations"]
module_header = "from __future__ import annotations"
references = _module_references(f"{module_header}\n\n{function_source}\n")
dynamic = _namespace_acquisition(definition, references)
if dynamic:
raise ValueError(f"@udf cannot package dynamic namespace access: {dynamic!r}")
# Resolve every module-scope reference the way the interpreter would: the
# function's own globals first (a module global may shadow a builtin, and
# nested scopes are not visible to getclosurevars), then its builtins.
# The artifact runs under the standard builtins; only the exact mapping is
# provably equivalent (a subclass or copy can change lookups and hooks).
if function.__builtins__ is not vars(builtins):
raise ValueError("@udf cannot package a non-standard builtins environment")
globals_source = []
unresolved = []
for name in sorted(references):
if name == function.__name__:
if not _is_recursive_reference(function, name):
raise ValueError(
f"@udf cannot package {name!r}: the module binds that name to "
"another value, which the artifact's own definition would shadow"
)
continue
if name in function.__globals__:
globals_source.append(_global_source(name, function.__globals__[name]))
elif hasattr(builtins, name):
pass
else:
unresolved.append(name)
if unresolved:
raise ValueError(
f"@udf source contains unresolved global names: {unresolved!r}"
)
parts = [module_header]
if globals_source:
parts.extend(["", *globals_source])
parts.extend(["", function_source, ""])
return "\n".join(parts).encode("utf-8")
packaged = "\n".join(parts)
return packaged.encode("utf-8")
class UdfDefinition:
@@ -797,7 +908,6 @@ class UdfDefinition:
output_schema: Optional[pa.DataType | pa.Field | pa.Schema],
pip: tuple[str, ...],
env: Mapping[str, str],
secrets: tuple[str, ...],
python_version: Optional[str],
):
function_name = name or function.__name__
@@ -812,18 +922,6 @@ class UdfDefinition:
for key, value in environment.items()
):
raise TypeError("Function env keys and values must be strings")
required_secrets = tuple(sorted(set(secrets)))
invalid_secrets = [
secret for secret in required_secrets if not _SECRET_NAME.fullmatch(secret)
]
if invalid_secrets:
raise ValueError(f"invalid Function secret names: {invalid_secrets!r}")
overlap = set(environment) & set(required_secrets)
if overlap:
raise ValueError(
f"Function env and secret names must be disjoint: {sorted(overlap)!r}"
)
signature = _infer_signature(function, input_schema, output_schema)
source = _package_source(function)
digest = f"sha256:{hashlib.sha256(source).hexdigest()}"
@@ -852,7 +950,6 @@ class UdfDefinition:
),
signature=signature,
runtime=runtime,
required_secrets=required_secrets,
)
functools.update_wrapper(self, function)
@@ -878,7 +975,6 @@ def udf(
output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
secrets: tuple[str, ...] | list[str] = (),
python_version: Optional[str] = None,
) -> Callable[[Callable[..., Any]], UdfDefinition]: ...
@@ -891,7 +987,6 @@ def udf(
output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
secrets: tuple[str, ...] | list[str] = (),
python_version: Optional[str] = None,
):
"""Prepare a scalar Python callable for remote Function registration.
@@ -916,13 +1011,17 @@ def udf(
pip : sequence of str, optional
Pip requirements for the remote environment.
env : mapping of str to str, optional
Non-secret environment variables. Use ``secrets`` for credentials.
secrets : sequence of str, optional
Names of secrets resolved by the remote service. Secret values are not
accepted by this API or included in the registration request.
Environment variables included in the Function definition.
python_version : str, optional
Remote Python major/minor version. Defaults to the client version.
The packaged artifact is a snapshot: the function source plus exactly
the module-level names it references (modules as imports, importable
classes and functions as imports, literals inline). Code that reaches the
module namespace another way -- ``globals()``/``eval``, ``sys.modules``,
``builtins`` -- is rejected where it can be seen and otherwise
unsupported; closures and a non-standard ``__builtins__`` are rejected.
Returns
-------
UdfDefinition
@@ -934,7 +1033,7 @@ def udf(
Examples
--------
>>> from lancedb import udf
>>> @udf(pip=["numpy==2.2.0"], secrets=["MODEL_TOKEN"])
>>> @udf(pip=["numpy==2.2.0"])
... def score(value: float) -> float:
... return value * 2
>>> score(1.5)
@@ -949,7 +1048,6 @@ def udf(
output_schema=output_schema,
pip=tuple(pip),
env={} if env is None else env,
secrets=tuple(secrets),
python_version=python_version,
)
python/python/lancedb/permutation.py
+21 -5
View File
@@ -391,6 +391,15 @@ def _table_to_pickle_state(table: Table) -> dict[str, Any]:
}
def _drop_base_version(permutation_data: pa.Table) -> pa.Table:
"""Strip the recorded base version so the reader leaves the base table unpinned."""
metadata = dict(permutation_data.schema.metadata or {})
if metadata.pop(b"base_version", None) is None:
return permutation_data
metadata.pop(b"base_branch", None)
return permutation_data.replace_schema_metadata(metadata)
def _table_from_pickle_state(state: dict[str, Any]) -> Table:
from . import connect
@@ -679,11 +688,15 @@ class Permutation:
from . import connect
connection_factory = state["connection_factory"]
rebuilt_base = False
if connection_factory is not None:
base_table = connection_factory(state["base_table_name"])
elif "base_table_state" in state:
base_table = _table_from_pickle_state(state["base_table_state"])
base_state = state["base_table_state"]
rebuilt_base = base_state["kind"] == "memory"
base_table = _table_from_pickle_state(base_state)
elif "base_table_data" in state:
rebuilt_base = True
# In-memory base table inlined into the pickle; rebuild the same
# way we rebuild the in-memory permutation table.
mem_db = connect("memory://")
@@ -701,11 +714,14 @@ class Permutation:
)
permutation_table: Optional[Table] = None
if state["permutation_data"] is not None:
permutation_data = state["permutation_data"]
if permutation_data is not None:
if rebuilt_base:
# The base table was materialized from Arrow, so it is a fresh
# single-version dataset and the recorded pin cannot resolve on it.
permutation_data = _drop_base_version(permutation_data)
mem_db = connect("memory://")
permutation_table = mem_db.create_table(
"permutation", state["permutation_data"]
)
permutation_table = mem_db.create_table("permutation", permutation_data)
self.base_table = base_table
self.permutation_table = permutation_table
python/python/lancedb/remote/table.py
+4
View File
@@ -610,6 +610,7 @@ class RemoteTable(Table):
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add more data to the [Table][lancedb.table.Table].
@@ -642,6 +643,8 @@ class RemoteTable(Table):
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Not supported on LanceDB Cloud. Setting this raises.
Returns
-------
@@ -658,6 +661,7 @@ class RemoteTable(Table):
fill_value=fill_value,
progress=progress,
write_parallelism=write_parallelism,
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
)
)
finally:
python/python/lancedb/streaming.py
+565 -33
View File
@@ -19,6 +19,7 @@ above.
"""
import ctypes
import heapq
import logging
import os
import random
@@ -29,17 +30,18 @@ from collections import deque
from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
from multiprocessing import RawArray
from typing import Any, Callable, cast, Iterator, Literal, Optional, Union
from typing import Any, Callable, cast, Iterator, Literal, NamedTuple, Optional, Union
import pyarrow as pa
import pyarrow.compute as pc
import torch
from torch.utils.data import IterableDataset, get_worker_info
from torch.utils.data import DataLoader, IterableDataset, get_worker_info
from .permutation import (
Permutation,
Transforms,
permutation_builder,
_drop_base_version,
_table_from_pickle_state,
_table_to_pickle_state,
)
@@ -55,6 +57,155 @@ DEFAULT_READ_BATCH_SIZE = 64
DEFAULT_PREFETCH_BATCHES = 4
class _WorkerSample(NamedTuple):
data: Any
dataset: "StreamingDataset"
class _WorkerBatch(NamedTuple):
data: Any
state: dict
class _ConsumerIteratorLease(NamedTuple):
owner_token: int
owner_thread: int
class _CheckpointCollate:
"""Attach the worker's post-fetch state to a collated batch."""
def __init__(self, collate_fn: Callable):
self._collate_fn = collate_fn
def __call__(self, samples):
try:
if isinstance(samples, list):
if not samples:
return _WorkerBatch(self._collate_fn(samples), {})
worker_samples = samples
data = self._collate_fn([sample.data for sample in worker_samples])
dataset = worker_samples[-1].dataset
else:
data = self._collate_fn(samples.data)
dataset = samples.dataset
except StopIteration as exc:
raise RuntimeError(
"collate_fn raised StopIteration before returning a batch"
) from exc
return _WorkerBatch(data, dataset._checkpoint_snapshot())
class _StreamingDatasetAdapter(IterableDataset):
"""Yield private sample wrappers for :class:`StreamingDataLoader`."""
def __init__(self, dataset: "StreamingDataset"):
super().__init__()
self.dataset = dataset
def __iter__(self):
for sample in self.dataset._iter(consumer_checkpoint_transport=True):
yield _WorkerSample(sample, self.dataset)
def __getattr__(self, name):
dataset = self.__dict__.get("dataset")
if dataset is None:
raise AttributeError(name)
return getattr(dataset, name)
class _ConsumerCommitIterator:
def __init__(
self,
iterator,
dataset: "StreamingDataset",
*,
owner_token: int,
require_uniform: bool,
):
self._iterator = iterator
self._dataset = dataset
self._owner_token = owner_token
self._require_uniform = require_uniform
self._released = False
self._terminal = False
def __iter__(self):
return self
def __next__(self):
if self._terminal:
raise StopIteration
try:
batch = next(self._iterator)
except StopIteration:
self._terminal = True
self._release()
raise
except BaseException as exc:
self._dataset._invalidate_checkpoint(
f"a DataLoader batch failed before it was returned: {exc}"
)
raise
try:
if not isinstance(batch, _WorkerBatch):
raise RuntimeError(
"StreamingDataLoader did not receive worker checkpoint metadata"
)
self._dataset._commit_worker_state(
batch.state, require_uniform=self._require_uniform
)
return batch.data
except BaseException as exc:
self._dataset._invalidate_checkpoint(
f"a DataLoader batch failed before it was returned: {exc}"
)
raise
def _release(self) -> None:
if self.__dict__.get("_released", True):
return
self._released = True
dataset = self.__dict__.get("_dataset")
if dataset is not None:
dataset._release_consumer_iterator(self._owner_token)
def _shutdown_workers(self):
if self.__dict__.get("_released", True):
return None
self._terminal = True
iterator = self.__dict__.get("_iterator")
shutdown = getattr(iterator, "_shutdown_workers", None)
try:
if shutdown is not None:
shutdown()
else:
fetcher = getattr(iterator, "_dataset_fetcher", None)
dataset_iterator = getattr(fetcher, "dataset_iter", None)
close = getattr(dataset_iterator, "close", None)
if close is None:
raise RuntimeError(
"StreamingDataLoader could not close its inner iterator"
)
close()
except BaseException as exc:
self._dataset._invalidate_checkpoint(
f"a DataLoader iterator could not be shut down safely: {exc}"
)
raise
else:
self._release()
def __del__(self):
try:
self._shutdown_workers()
except BaseException:
pass
def __getattr__(self, name):
return getattr(self._iterator, name)
class StreamingDataset(IterableDataset):
"""An elastic, resumable PyTorch IterableDataset backed by a LanceDB table.
@@ -384,6 +535,22 @@ class StreamingDataset(IterableDataset):
# rows_skipped]
self._worker_stats: RawArray = RawArray(ctypes.c_int64, 8)
# A standard multi-process DataLoader cannot report which prefetched
# batches were actually returned to its consumer. Workers set this
# shared flag so state_dict() can reject a stale parent checkpoint
# unless StreamingDataLoader installed the consumer-commit transport.
self._untracked_worker_iteration: RawArray = RawArray(ctypes.c_int64, 1)
# Parent-side checkpoint lifecycle. A failed DataLoader task creates
# a permanent hole in that iterator's delivery stream, while a
# multi-worker checkpoint is safe to restore only after all splits
# reach the same logical step boundary.
self._checkpoint_invalid_reason: Optional[str] = None
self._consumer_checkpoint_requires_uniform = False
self._consumer_iterator_lock = threading.Lock()
self._consumer_iterator_generation = 0
self._consumer_iterator_lease: Optional[_ConsumerIteratorLease] = None
# Cumulative bytes of Arrow buffer data fetched across all iterations.
self._bytes_loaded: int = 0
# Cumulative seconds spent in LanceDB I/O and in transform functions.
@@ -396,6 +563,10 @@ class StreamingDataset(IterableDataset):
# step boundaries all splits have consumed this many samples, so a
# single scalar captures the topology-independent checkpoint state.
self._resume_offset: int = 0
# Exact yielded-sample counts for splits this process has advanced.
# Missing entries use _resume_offset, which remains the lower-bound
# checkpoint inherited from an earlier uniform/global state.
self._resume_samples: dict[int, int] = {}
# Permutation position each split has consumed through, keyed by
# global split index. Equal to _resume_offset for every split unless
# on_transform_error skipped rows, in which case skipped positions
@@ -521,11 +692,45 @@ class StreamingDataset(IterableDataset):
return self._rank_splits[start : start + splits_per_worker]
def __iter__(self) -> Iterator[dict[str, Any]]:
return self._iter()
def _iter(
self, *, consumer_checkpoint_transport: bool = False
) -> Iterator[dict[str, Any]]:
owner_token = None
previous_lease = self._consumer_iterator_lease
if consumer_checkpoint_transport:
if not self._consumer_iterator_active:
raise RuntimeError(
"StreamingDataLoader worker transport requires an active "
"parent iterator reservation"
)
else:
try:
owner_token = self._acquire_consumer_iterator()
except BaseException:
self._release_consumer_iterator_after_failed_acquire(previous_lease)
raise
try:
yield from self._iter_owned(
consumer_checkpoint_transport=consumer_checkpoint_transport
)
finally:
if owner_token is not None:
self._release_consumer_iterator(owner_token)
def _iter_owned(
self, *, consumer_checkpoint_transport: bool
) -> Iterator[dict[str, Any]]:
if self._raw_batches_ref is not None:
raise RuntimeError(
"StreamingDataset does not support concurrent iteration. "
"Only one active iterator per dataset instance is allowed."
)
real_worker = get_worker_info() is not None
if real_worker and not consumer_checkpoint_transport:
self._untracked_worker_iteration[0] = 1
my_splits = self._resolve_my_splits()
if not my_splits:
return
@@ -533,6 +738,7 @@ class StreamingDataset(IterableDataset):
# Set identity transform on each Permutation so __getitems__ returns
# the raw RecordBatch. Stage 2 applies the real transform.
permutations: list[Permutation] = []
initial_samples: list[int] = []
initial_positions: list[int] = []
for split_idx in my_splits:
perm = Permutation.from_tables(
@@ -541,21 +747,22 @@ class StreamingDataset(IterableDataset):
if self._columns is not None:
perm = perm.select_columns(self._columns)
perm = perm.with_transform(Transforms.arrow2arrow)
sample_count = self._resume_samples.get(split_idx, self._resume_offset)
# Both modes resume from absolute permutation positions. Packing
# stores them separately because it also checkpoints partial blocks.
start_pos = (
self._pack_consumed[split_idx]
if self._pack_sequences is not None
else self._resume_positions.get(split_idx, self._resume_offset)
else self._resume_positions.get(split_idx, sample_count)
)
if start_pos > 0:
perm = perm.with_skip(start_pos)
initial_samples.append(sample_count)
initial_positions.append(start_pos)
permutations.append(perm)
n = len(permutations)
split_sizes = [perm.num_rows for perm in permutations]
initial_offset = self._resume_offset
local_consumed = [0] * n
# Permutation position each split has consumed through (absolute,
# i.e. counted from the start of the unskipped split). Runs ahead of
@@ -853,6 +1060,27 @@ class StreamingDataset(IterableDataset):
for i in range(n):
_fill_io(i)
def _yield_row(i: int):
pos, row = cooked[i].popleft()
# Surface any completed prefetched failure before the
# current row becomes durable checkpoint progress.
_advance(i)
local_consumed[i] += 1
pos_consumed[i] = pos + 1
split_idx = my_splits[i]
self._resume_samples[split_idx] = (
initial_samples[i] + local_consumed[i]
)
self._resume_positions[split_idx] = pos_consumed[i]
return row
def _update_progress_stats() -> None:
if not real_worker:
self._resume_offset = min(
initial_samples[j] + local_consumed[j] for j in range(n)
)
_update_stats()
if self._pack_sequences is not None:
first_count = pack_blocks_emitted[my_splits[0]]
if any(
@@ -878,12 +1106,38 @@ class StreamingDataset(IterableDataset):
tokens.extend([pad_id] * (pack_len - len(tokens)))
block = _emit_block(i)
pack_blocks_emitted[my_splits[i]] += 1
# Checkpoint state must advance before yielding so
# StreamingDataLoader can attach the exact state to
# the batch it transports to the parent process.
_commit_pack_state()
if i == n - 1:
_commit_pack_state()
_update_stats()
yield block
return
# A checkpoint taken between round-robin split turns has
# non-uniform counts. Resume lagging splits first so the
# exact canonical sequence continues without replaying
# already-consumed rows.
if len(set(initial_samples)) > 1:
catch_up_to = max(initial_samples)
pending = [
(initial_samples[i], my_splits[i], i)
for i in range(n)
if initial_samples[i] < catch_up_to
]
heapq.heapify(pending)
while pending:
consumed, _, i = heapq.heappop(pending)
_ensure_cooked(i)
if not cooked[i]:
return
row = _yield_row(i)
if consumed + 1 < catch_up_to:
heapq.heappush(pending, (consumed + 1, my_splits[i], i))
_update_progress_stats()
yield row
while True:
# A cycle only runs if every split can still produce a
# row. Without skips all splits exhaust simultaneously
@@ -904,20 +1158,14 @@ class StreamingDataset(IterableDataset):
break
for i in range(n):
pos, row = cooked[i].popleft()
local_consumed[i] += 1
pos_consumed[i] = pos + 1
_advance(i)
row = _yield_row(i)
# After the last split in each cycle: update the
# global offset and refresh the shared-memory stats
# so the main process can observe pipeline depth
# even when __iter__ runs in a worker process.
if i == n - 1:
self._resume_offset = initial_offset + local_consumed[i]
for j, split_idx in enumerate(my_splits):
self._resume_positions[split_idx] = pos_consumed[j]
_update_stats()
_update_progress_stats()
yield row
finally:
@@ -1064,6 +1312,7 @@ class StreamingDataset(IterableDataset):
"_local_consumed_ref",
):
state[key] = None
state["_consumer_iterator_lock"] = None
return state
def __setstate__(self, state):
@@ -1074,19 +1323,31 @@ class StreamingDataset(IterableDataset):
table_state = state.pop("_table")
perm_name, perm_data = state.pop("_perm_table")
self.__dict__.update(state)
self._consumer_iterator_lock = threading.Lock()
if self._connection_factory is not None:
self._table = self._connection_factory(table_name)
else:
self._table = _table_from_pickle_state(table_state)
if table_state["kind"] == "memory":
# Rebuilt from Arrow, so the recorded pin cannot resolve on it.
perm_data = _drop_base_version(perm_data)
self._perm_table = _connect("memory://").create_table(perm_name, perm_data)
def state_dict(self) -> dict:
"""Snapshot the dataset's consumption state.
When using DataLoader workers, construct a
[StreamingDataLoader][lancedb.streaming.StreamingDataLoader]. It
commits worker state only when a prefetched batch is returned to the
trainer. A standard multi-process ``DataLoader`` cannot expose that
boundary, so calling this method after one has started raises
``RuntimeError`` instead of returning stale producer state.
In row mode, the returned dict is topology-independent at global step
boundaries. ``positions_consumed_per_split`` records how far each
split's permutation has advanced, which can differ from the sample
count when ``on_transform_error`` skips rows. Combine state dicts from
count when ``on_transform_error`` skips rows. ``StreamingDataLoader``
combines worker state in its parent process. Combine state dicts from
every rank with
[merge_state_dicts][lancedb.streaming.StreamingDataset.merge_state_dicts]
before resuming on a different topology.
@@ -1095,6 +1356,43 @@ class StreamingDataset(IterableDataset):
for every logical split. When packing is sharded, merge every rank
state with ``merge_state_dicts`` before loading it.
"""
if self._untracked_worker_iteration[0] and get_worker_info() is None:
raise RuntimeError(
"StreamingDataset cannot checkpoint a standard DataLoader with "
"num_workers > 0 because prefetched worker progress is not "
"consumer-committed. Use StreamingDataLoader instead."
)
if self._checkpoint_invalid_reason is not None:
raise RuntimeError(
"StreamingDataset checkpointing is invalid because "
f"{self._checkpoint_invalid_reason}. Load the last valid "
"checkpoint into a fresh dataset before continuing."
)
state = self._checkpoint_snapshot()
if self._pack_sequences is not None:
rank_blocks = [
state["blocks_emitted_per_split"][split] for split in self._rank_splits
]
if len(set(rank_blocks)) > 1:
raise RuntimeError(
"Packed StreamingDataset checkpointing is only safe at a "
"complete logical step boundary, when every split assigned "
"to this rank has emitted the same block count. Consume more "
"batches before calling state_dict()."
)
elif self._consumer_checkpoint_requires_uniform:
samples = state["samples_consumed_per_split"]
rank_samples = [samples[split] for split in self._rank_splits]
if len(set(rank_samples)) > 1:
raise RuntimeError(
"StreamingDataLoader checkpointing with multiple workers is "
"only safe at a complete logical step boundary, when every "
"split assigned to this rank has the same consumed-sample "
"count. Consume more batches before calling state_dict()."
)
return state
def _checkpoint_snapshot(self) -> dict:
if self._pack_sequences is not None:
return {
"shuffle_seed": self._shuffle_seed,
@@ -1108,18 +1406,141 @@ class StreamingDataset(IterableDataset):
"blocks_emitted_per_split": list(self._pack_blocks_emitted),
"pack_buffers": deepcopy(self._pack_buffers),
}
samples = [
self._resume_samples.get(split, self._resume_offset)
for split in range(self._num_splits)
]
positions = [
self._resume_positions.get(split, self._resume_offset)
self._resume_positions.get(split, samples[split])
for split in range(self._num_splits)
]
return {
"shuffle_seed": self._shuffle_seed,
"num_splits": self._num_splits,
"epoch": self._epoch,
"samples_consumed_per_split": [self._resume_offset] * self._num_splits,
"samples_consumed_per_split": samples,
"positions_consumed_per_split": positions,
}
def _invalidate_checkpoint(self, reason: str) -> None:
if self._checkpoint_invalid_reason is None:
self._checkpoint_invalid_reason = reason
@property
def _consumer_iterator_active(self) -> bool:
return self._consumer_iterator_lease is not None
@property
def _consumer_iterator_owner(self) -> Optional[int]:
lease = self._consumer_iterator_lease
return lease.owner_token if lease is not None else None
@property
def _consumer_iterator_owner_thread(self) -> Optional[int]:
lease = self._consumer_iterator_lease
return lease.owner_thread if lease is not None else None
def _acquire_consumer_iterator(self) -> int:
"""Reserve this parent dataset for one checkpoint-aware iterator."""
with self._consumer_iterator_lock:
if self._consumer_iterator_active or self._raw_batches_ref is not None:
raise RuntimeError(
"StreamingDataset does not support concurrent iteration. "
"Only one active iterator per dataset instance is allowed."
)
owner_thread = threading.get_ident()
owner_token = self._consumer_iterator_generation + 1
lease = _ConsumerIteratorLease(owner_token, owner_thread)
self._consumer_iterator_generation = owner_token
self._consumer_iterator_lease = lease
return owner_token
def _release_consumer_iterator(self, owner_token: int) -> None:
with self._consumer_iterator_lock:
lease = self._consumer_iterator_lease
if lease is not None and lease.owner_token == owner_token:
self._consumer_iterator_lease = None
def _release_consumer_iterator_after_failed_acquire(
self, previous_lease: Optional[_ConsumerIteratorLease]
) -> None:
"""Clean up when an interrupted acquire set a lease but did not return it."""
owner_thread = threading.current_thread().ident
with self._consumer_iterator_lock:
lease = self._consumer_iterator_lease
if (
lease is not None
and lease is not previous_lease
and lease.owner_thread == owner_thread
):
self._consumer_iterator_lease = None
def _commit_worker_state(self, state: dict, *, require_uniform: bool) -> None:
"""Merge one trainer-consumed worker batch into parent state."""
for key, expected in (
("shuffle_seed", self._shuffle_seed),
("num_splits", self._num_splits),
("epoch", self._epoch),
):
if state.get(key) != expected:
raise ValueError(
f"{key} mismatch in worker checkpoint: "
f"{state.get(key)} != {expected}"
)
packed = "pack_buffers" in state
if packed != (self._pack_sequences is not None):
raise ValueError("worker checkpoint mode does not match the dataset")
if packed:
for key in ("pack_sequences", "eos_id", "pad_id", "blocks_per_epoch"):
expected = getattr(self, f"_{key}")
if state.get(key) != expected:
raise ValueError(
f"{key} mismatch in worker checkpoint: "
f"{state.get(key)} != {expected}"
)
samples = state["samples_consumed_per_split"]
emitted = state["blocks_emitted_per_split"]
if len(samples) != self._num_splits or len(emitted) != self._num_splits:
raise ValueError(
"packed worker checkpoint must contain one entry per split"
)
buffers = state["pack_buffers"]
for split, (count, blocks) in enumerate(zip(samples, emitted)):
incoming = (int(blocks), int(count))
current = (
self._pack_blocks_emitted[split],
self._pack_consumed[split],
)
if incoming > current:
self._pack_blocks_emitted[split] = incoming[0]
self._pack_consumed[split] = incoming[1]
buffer = buffers.get(split, buffers.get(str(split)))
if buffer is None:
self._pack_buffers.pop(split, None)
else:
self._pack_buffers[split] = {
"tokens": list(buffer["tokens"]),
"starts": list(buffer["starts"]),
}
self._consumer_checkpoint_requires_uniform |= require_uniform
return
samples = state["samples_consumed_per_split"]
positions = state.get("positions_consumed_per_split", samples)
for split, count in enumerate(samples):
current = self._resume_samples.get(split, self._resume_offset)
self._resume_samples[split] = max(current, int(count))
for split, position in enumerate(positions):
current = self._resume_positions.get(
split, self._resume_samples.get(split, self._resume_offset)
)
self._resume_positions[split] = max(current, int(position))
self._resume_offset = min(
self._resume_samples.get(split, self._resume_offset)
for split in range(self._num_splits)
)
self._consumer_checkpoint_requires_uniform |= require_uniform
def load_state_dict(self, state: dict) -> None:
"""Resume from a previously snapshotted state.
@@ -1139,6 +1560,7 @@ class StreamingDataset(IterableDataset):
f"shuffle_seed mismatch: checkpoint has {state['shuffle_seed']}, "
f"current dataset has {self._shuffle_seed}"
)
self._consumer_checkpoint_requires_uniform = False
if "pack_buffers" in state or self._pack_sequences is not None:
for key in (
@@ -1165,14 +1587,17 @@ class StreamingDataset(IterableDataset):
return
consumed = state["samples_consumed_per_split"]
# All entries are equal at step boundaries; use the first.
if isinstance(consumed, list):
self._resume_offset = consumed[0] if consumed else 0
self._resume_offset = min(consumed) if consumed else 0
self._resume_samples = {
split: int(count) for split, count in enumerate(consumed)
}
else:
self._resume_offset = int(consumed)
self._resume_samples = {}
# Older checkpoints predate positions_consumed_per_split; without
# skipped rows positions equal sample counts, so falling back to
# _resume_offset (the .get default in __iter__) is exact.
# the per-split sample count (the .get default in __iter__) is exact.
positions = state.get("positions_consumed_per_split")
if positions is None:
self._resume_positions = {}
@@ -1185,10 +1610,11 @@ class StreamingDataset(IterableDataset):
def merge_state_dicts(states: list[dict]) -> dict:
"""Merge state dicts saved by different ranks into one exact state.
For row mode, the elementwise maximum of permutation positions recovers
splits advanced by different ranks after transform failures. For packed
mode, the state that emitted the most blocks for each logical split
supplies that split's permutation position and partial token buffer. Packed
In row mode, each rank records exact consumer-committed progress for
its own splits and lower bounds for the rest, so elementwise maxima
recover both sample counts and permutation positions. In packed mode,
the state that emitted the most blocks for each logical split supplies
that split's permutation position and partial token buffer. Packed
states must cover every rank at the same global step.
Raises ``ValueError`` if the states are empty, were not produced by
@@ -1299,17 +1725,13 @@ class StreamingDataset(IterableDataset):
merged["pack_buffers"] = merged_buffers
return merged
for state in states[1:]:
if (
state["samples_consumed_per_split"]
!= first["samples_consumed_per_split"]
):
raise ValueError(
"samples_consumed_per_split mismatch across state dicts; "
"state_dict() must be called at the same global step "
"boundary on every rank"
)
merged = dict(first)
merged["samples_consumed_per_split"] = [
max(per_split)
for per_split in zip(
*(state["samples_consumed_per_split"] for state in states)
)
]
all_positions = [
state.get(
"positions_consumed_per_split", state["samples_consumed_per_split"]
@@ -1320,3 +1742,113 @@ class StreamingDataset(IterableDataset):
max(per_split) for per_split in zip(*all_positions)
]
return merged
class StreamingDataLoader(DataLoader):
"""A PyTorch DataLoader with consumer-committed dataset checkpoints.
PyTorch workers prefetch batches ahead of the trainer, so worker-local
producer progress is not a safe checkpoint. This loader carries a state
snapshot alongside every internal batch and applies it to the parent
[StreamingDataset][lancedb.streaming.StreamingDataset] only when that batch
is returned by ``next()``.
The trainer receives the same collated batch it would receive from a
standard ``torch.utils.data.DataLoader``.
With more than one worker, row-mode ``state_dict()`` is available only at
complete logical step boundaries, when every split assigned to the rank has
the same consumed-sample count. Packed checkpoints require equal emitted-block
counts across the rank's splits for any worker count. ``persistent_workers=True``
is not supported because prefetched worker copies cannot be restored from
parent-committed state. If batch collation raises, checkpointing remains
invalid for that dataset instance; restore the last valid checkpoint into a
fresh dataset before continuing.
Only one active iterator may own a dataset at a time, including when worker
processes are used. Exhausting or explicitly shutting down the iterator
releases that ownership. ``drop_last=True`` is not supported because worker
replicas discard incomplete tails independently, which cannot produce a
topology-independent checkpoint.
Parameters are the same as ``torch.utils.data.DataLoader`` except that
``dataset`` must be a
[StreamingDataset][lancedb.streaming.StreamingDataset].
Subclasses that override ``StreamingDataset.__iter__`` are not supported
because the custom iterator cannot provide the exact per-yield checkpoint
snapshots required by this loader.
Examples
--------
>>> # dataset = StreamingDataset(table, num_splits=2)
>>> # loader = StreamingDataLoader(dataset, batch_size=8, num_workers=2)
>>> # batch = next(iter(loader))
>>> # checkpoint = dataset.state_dict()
"""
def __init__(self, dataset: StreamingDataset, *args, **kwargs):
if not isinstance(dataset, StreamingDataset):
raise TypeError("StreamingDataLoader requires a StreamingDataset")
if type(dataset).__iter__ is not StreamingDataset.__iter__:
raise TypeError(
"StreamingDataLoader does not support StreamingDataset subclasses "
"that override __iter__ because they cannot provide exact "
"per-yield checkpoint state"
)
if kwargs.get("in_order", True) is False:
raise ValueError(
"StreamingDataLoader requires in_order=True for deterministic "
"consumer checkpoints"
)
if kwargs.get("persistent_workers", False):
raise ValueError(
"StreamingDataLoader does not support persistent_workers=True "
"because worker prefetch state cannot be reset from a checkpoint"
)
self._streaming_dataset = dataset
super().__init__(_StreamingDatasetAdapter(dataset), *args, **kwargs)
if self.drop_last:
raise ValueError(
"StreamingDataLoader does not support drop_last=True because "
"discarded worker tails cannot be checkpointed "
"topology-independently"
)
self.collate_fn = _CheckpointCollate(self.collate_fn)
def __iter__(self):
dataset = self._streaming_dataset
previous_lease = dataset._consumer_iterator_lease
owner_token = None
try:
owner_token = dataset._acquire_consumer_iterator()
state = dataset._checkpoint_snapshot()
packed = dataset._pack_sequences is not None
if packed:
blocks = state["blocks_emitted_per_split"]
rank_blocks = [blocks[split] for split in dataset._rank_splits]
if len(set(rank_blocks)) > 1:
raise RuntimeError(
"StreamingDataLoader cannot start from a partial packed "
"logical step; resume from a checkpoint whose splits "
"assigned to this rank have equal emitted-block counts"
)
elif self.num_workers > 1:
samples = state["samples_consumed_per_split"]
rank_samples = [samples[split] for split in dataset._rank_splits]
if len(set(rank_samples)) > 1:
raise RuntimeError(
"StreamingDataLoader cannot start multiple workers from a "
"partial logical step; resume from a checkpoint whose "
"splits assigned to this rank have equal consumed-sample "
"counts"
)
return _ConsumerCommitIterator(
super().__iter__(),
dataset,
owner_token=owner_token,
require_uniform=self.num_workers > 1 or packed,
)
except BaseException:
if owner_token is not None:
dataset._release_consumer_iterator(owner_token)
else:
dataset._release_consumer_iterator_after_failed_acquire(previous_lease)
raise
python/python/lancedb/table.py
+18 -3
View File
@@ -1269,6 +1269,7 @@ class Table(ABC):
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add more data to the [Table][lancedb.table.Table].
@@ -1320,6 +1321,10 @@ class Table(ABC):
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Store blob URIs that sit outside registered blob bases. The row
keeps a reference, so the object has to stay readable. Local
tables only.
Returns
-------
@@ -1972,7 +1977,7 @@ class Table(ABC):
A mapping with one ``FunctionApplication`` value keeps its scalar
or named-struct result in the named table column. A bare
named-struct application expands its ordered result fields as one
atomic sibling group; aliases come from ``rename(columns=...)``.
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only on LanceDB Cloud and
Enterprise.
computed: Dict[str, str], optional
@@ -3409,6 +3414,7 @@ class LanceTable(Table):
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add data to the table.
If vector columns are missing and the table
@@ -3436,6 +3442,9 @@ class LanceTable(Table):
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Allow blob URIs outside registered bases. See :meth:`Table.add`.
Local tables only.
Returns
-------
@@ -3452,6 +3461,7 @@ class LanceTable(Table):
fill_value=fill_value,
progress=progress,
write_parallelism=write_parallelism,
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
)
)
finally:
@@ -5366,6 +5376,7 @@ class AsyncTable:
fill_value: Optional[float] = None,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add more data to the [AsyncTable][lancedb.table.AsyncTable].
@@ -5396,6 +5407,9 @@ class AsyncTable:
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Allow blob URIs outside registered bases. See :meth:`Table.add`.
Local tables only.
"""
schema = await self.schema()
@@ -5432,6 +5446,7 @@ class AsyncTable:
mode or "append",
progress=progress,
write_parallelism=write_parallelism,
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
)
except RuntimeError as e:
if "Cast error" in str(e):
@@ -6056,7 +6071,7 @@ class AsyncTable:
A mapping with one ``FunctionApplication`` value keeps its scalar
or named-struct result in the named table column. A bare
named-struct application expands its ordered result fields as one
atomic sibling group; aliases come from ``rename(columns=...)``.
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only on LanceDB Cloud and
Enterprise.
computed: Dict[str, str], optional
@@ -6093,7 +6108,7 @@ class AsyncTable:
isinstance(value, FunctionApplication) for value in transforms.values()
):
raise ValueError(
"one add_columns call declares exactly one Function sibling group"
"one add_columns call declares exactly one Function binding"
)
function_output_name, function_application = next(iter(transforms.items()))
python/python/tests/test_blob.py
+68
View File
@@ -617,3 +617,71 @@ def test_fetch_blobs_nested_path_survives_sort_after_query():
def _identifiable_payload(size: int) -> bytes:
block = 256
return b"".join(bytes([i % 256]) * block for i in range(size // block))
def _external_uri_blob_array(uris):
blob_type = lancedb.blob("image").type
storage_type = blob_type.storage_type
child_names = [field.name for field in storage_type]
assert "uri" in child_names, "blob layout no longer has a uri child"
children = [
pa.array(uris if field.name == "uri" else [None] * len(uris), type=field.type)
for field in storage_type
]
storage = pa.StructArray.from_arrays(children, fields=list(storage_type))
return pa.ExtensionArray.from_storage(blob_type, storage)
def _external_uri_table_and_rows(name, uris):
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table(name, schema=schema)
rows = pa.Table.from_arrays(
[
pa.array(range(len(uris)), type=pa.int64()),
_external_uri_blob_array(uris),
],
schema=schema,
)
return table, rows
def test_add_external_uri_struct_round_trips_with_flag(tmp_path):
payload = b"external-uri-bytes"
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(payload)
table, rows = _external_uri_table_and_rows("external_struct", [blob_path.as_uri()])
table.add(rows, allow_external_blob_outside_bases=True)
hits = table.search().to_arrow()
blobs = table.fetch_blobs("image", hits)
assert blobs[0].as_py() == payload
def test_add_external_uri_without_flag_raises(tmp_path):
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(b"unreachable")
table, rows = _external_uri_table_and_rows("external_no_flag", [blob_path.as_uri()])
with pytest.raises(ValueError, match="allow_external_blob_outside_bases"):
table.add(rows)
assert table.count_rows() == 0
def test_add_external_uri_string_round_trips_with_flag(tmp_path):
payload = b"external-uri-bytes"
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(payload)
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("external_string", schema=schema)
table.add(
[{"id": 1, "image": blob_path.as_uri()}],
allow_external_blob_outside_bases=True,
)
hits = table.search().to_arrow()
blobs = table.fetch_blobs("image", hits)
assert blobs[0].as_py() == payload
python/python/tests/test_elastic_dataloader.py
+658
View File
@@ -32,6 +32,7 @@ Parameters used throughout:
import dataclasses
import logging
import threading
from unittest.mock import patch
import lancedb
@@ -46,6 +47,7 @@ from utils import (
torch = pytest.importorskip("torch")
streaming = pytest.importorskip("lancedb.streaming")
StreamingDataset = streaming.StreamingDataset
StreamingDataLoader = streaming.StreamingDataLoader
# ---------------------------------------------------------------------------
# Dataset parameters
@@ -92,6 +94,27 @@ class FakeWorkerInfo:
num_workers: int
def _collate_with_first_batch_error(samples):
ids = [sample["id"] for sample in samples]
if ids == [0, 1]:
raise ValueError("first batch fails")
return ids
def _collate_with_first_batch_stop(samples):
ids = [sample["id"] for sample in samples]
if ids == [0, 1]:
raise StopIteration("first batch stopped")
return ids
def _collate_with_first_batch_interrupt(samples):
ids = [sample["id"] for sample in samples]
if ids == [0, 1]:
raise KeyboardInterrupt("first batch interrupted")
return ids
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@@ -1008,6 +1031,565 @@ def test_multi_worker_elastic_det_across_worker_counts(lance_table):
# ── Resumability with num_workers ─────────────────────────────────────────────
def test_streaming_dataloader_commits_only_consumed_worker_batches(tmp_path):
"""Prefetched worker state is committed only as the trainer receives it."""
db = lancedb.connect(tmp_path)
table = db.create_table(
"worker_commit", pa.table({"id": [1, 2, 3, 4, 10, 20, 30, 40]})
)
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
loader = StreamingDataLoader(
dataset,
batch_size=2,
num_workers=2,
multiprocessing_context="spawn",
prefetch_factor=4,
)
iterator = iter(loader)
try:
first = next(iterator)["id"].tolist()
assert first == [1, 2]
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2, 0]
with pytest.raises(RuntimeError, match="complete logical step boundary"):
dataset.state_dict()
second = next(iterator)["id"].tolist()
assert second == [10, 20]
checkpoint = dataset.state_dict()
assert checkpoint["samples_consumed_per_split"] == [2, 2]
uninterrupted = [batch["id"].tolist() for batch in iterator]
finally:
iterator._shutdown_workers()
resumed = StreamingDataset(table, num_splits=2, shuffle=False)
resumed.load_state_dict(checkpoint)
resumed_loader = StreamingDataLoader(
resumed,
batch_size=2,
num_workers=2,
multiprocessing_context="spawn",
prefetch_factor=4,
)
resumed_iterator = iter(resumed_loader)
try:
remaining = [batch["id"].tolist() for batch in resumed_iterator]
finally:
resumed_iterator._shutdown_workers()
assert remaining == uninterrupted == [[3, 4], [30, 40]]
def test_distributed_checkpoint_uses_rank_local_worker_boundary(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("rank_boundary", pa.table({"id": list(range(8))}))
dataset = StreamingDataset(
table,
num_splits=4,
shuffle=False,
rank=0,
world_size=2,
)
loader = StreamingDataLoader(
dataset,
batch_size=1,
num_workers=2,
multiprocessing_context="spawn",
)
iterator = iter(loader)
try:
assert next(iterator)["id"].tolist() == [0]
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [
1,
0,
0,
0,
]
with pytest.raises(RuntimeError, match="complete logical step boundary"):
dataset.state_dict()
assert next(iterator)["id"].tolist() == [2]
checkpoint = dataset.state_dict()
remaining = [batch["id"].tolist() for batch in iterator]
finally:
iterator._shutdown_workers()
assert checkpoint["samples_consumed_per_split"] == [1, 1, 0, 0]
assert remaining == [[1], [3]]
def test_standard_dataloader_rejects_stale_parent_checkpoint(tmp_path):
"""A standard DataLoader must not expose prefetched producer progress."""
db = lancedb.connect(tmp_path)
table = db.create_table("untracked_workers", pa.table({"id": [1, 2, 10, 20]}))
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
# Merely constructing the checkpoint-aware loader must not authorize a
# later plain DataLoader's worker progress.
StreamingDataLoader(dataset, batch_size=2, num_workers=0)
loader = torch.utils.data.DataLoader(
dataset,
batch_size=2,
num_workers=2,
multiprocessing_context="spawn",
)
iterator = iter(loader)
try:
assert next(iterator)["id"].tolist() == [1, 2]
with pytest.raises(RuntimeError, match="Use StreamingDataLoader"):
dataset.state_dict()
list(iterator)
finally:
iterator._shutdown_workers()
def test_streaming_dataloader_rejects_persistent_workers(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("persistent_workers", pa.table({"id": [1, 2]}))
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
with pytest.raises(ValueError, match="persistent_workers=True"):
StreamingDataLoader(
dataset,
batch_size=1,
num_workers=2,
persistent_workers=True,
)
def test_collate_failure_invalidates_consumer_checkpoint(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table(
"collate_failure", pa.table({"id": [0, 1, 2, 3, 100, 101, 102, 103]})
)
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
loader = StreamingDataLoader(
dataset,
batch_size=2,
num_workers=2,
multiprocessing_context="spawn",
collate_fn=_collate_with_first_batch_error,
prefetch_factor=2,
)
iterator = iter(loader)
try:
with pytest.raises(ValueError, match="first batch fails"):
next(iterator)
assert next(iterator) == [100, 101]
assert next(iterator) == [2, 3]
with pytest.raises(RuntimeError, match="failed before it was returned"):
dataset.state_dict()
list(iterator)
finally:
iterator._shutdown_workers()
def test_collate_stop_iteration_invalidates_consumer_checkpoint(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("collate_stop", pa.table({"id": list(range(6))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(
dataset,
batch_size=2,
num_workers=0,
collate_fn=_collate_with_first_batch_stop,
)
iterator = iter(loader)
with pytest.raises(RuntimeError, match="collate_fn raised StopIteration"):
next(iterator)
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2]
with pytest.raises(RuntimeError, match="failed before it was returned"):
dataset.state_dict()
assert list(iterator) == [[2, 3], [4, 5]]
def test_batch_base_exception_invalidates_consumer_checkpoint(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("collate_interrupt", pa.table({"id": list(range(6))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(
dataset,
batch_size=2,
num_workers=0,
collate_fn=_collate_with_first_batch_interrupt,
)
iterator = iter(loader)
with pytest.raises(KeyboardInterrupt, match="first batch interrupted"):
next(iterator)
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2]
with pytest.raises(RuntimeError, match="failed before it was returned"):
dataset.state_dict()
assert list(iterator) == [[2, 3], [4, 5]]
def test_parent_commit_base_exception_invalidates_consumer_checkpoint(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("commit_interrupt", pa.table({"id": list(range(4))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
iterator = iter(loader)
real_commit = dataset._commit_worker_state
def interrupt_after_commit(state, *, require_uniform):
real_commit(state, require_uniform=require_uniform)
raise KeyboardInterrupt("after parent commit")
with patch.object(
dataset, "_commit_worker_state", side_effect=interrupt_after_commit
):
with pytest.raises(KeyboardInterrupt, match="after parent commit"):
next(iterator)
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2]
with pytest.raises(RuntimeError, match="failed before it was returned"):
dataset.state_dict()
def test_direct_iteration_surfaces_prefetch_failure_before_committing_row(
tmp_path, monkeypatch
):
db = lancedb.connect(tmp_path)
table = db.create_table("prefetch_failure", pa.table({"id": list(range(4))}))
release = threading.Event()
failed = threading.Event()
real_getitems = streaming.Permutation.__getitems__
def controlled_getitems(permutation, indices):
if indices and indices[0] >= 2:
assert release.wait(timeout=5)
failed.set()
raise RuntimeError("later prefetched I/O failed")
return real_getitems(permutation, indices)
class SignalDict(dict):
def __setitem__(self, key, value):
super().__setitem__(key, value)
release.set()
assert failed.wait(timeout=5)
monkeypatch.setattr(streaming.Permutation, "__getitems__", controlled_getitems)
dataset = StreamingDataset(
table,
num_splits=1,
shuffle=False,
read_batch_size=2,
io_queue_depth=2,
)
dataset._resume_positions = SignalDict()
iterator = iter(dataset)
assert next(iterator)["id"] == 0
with pytest.raises(RuntimeError, match="later prefetched I/O failed"):
next(iterator)
checkpoint = dataset.state_dict()
assert checkpoint["samples_consumed_per_split"] == [1]
assert checkpoint["positions_consumed_per_split"] == [1]
@pytest.mark.parametrize("workers", [0, 1, 2])
def test_streaming_dataloader_rejects_drop_last(tmp_path, workers):
db = lancedb.connect(tmp_path)
table = db.create_table("drop_last", pa.table({"id": [0, 1, 2]}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
worker_options = {"multiprocessing_context": "spawn"} if workers else {}
with pytest.raises(ValueError, match="drop_last=True"):
StreamingDataLoader(
dataset,
batch_size=2,
num_workers=workers,
drop_last=True,
**worker_options,
)
def test_streaming_dataloader_owns_one_iterator_until_teardown(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("iterator_owner", pa.table({"id": list(range(4))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(
dataset,
batch_size=2,
num_workers=1,
multiprocessing_context="spawn",
)
first = iter(loader)
try:
assert next(first)["id"].tolist() == [0, 1]
with pytest.raises(RuntimeError, match="concurrent iteration"):
iter(loader)
finally:
first._shutdown_workers()
second = iter(loader)
try:
assert [batch["id"].tolist() for batch in second] == [[2, 3]]
except BaseException:
second._shutdown_workers()
raise
# Natural exhaustion releases ownership too.
third = iter(loader)
try:
assert list(third) == []
finally:
third._shutdown_workers()
def test_zero_worker_shutdown_closes_inner_iterator_before_release(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("zero_worker_shutdown", pa.table({"id": list(range(6))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
first = iter(loader)
assert next(first)["id"].tolist() == [0, 1]
first._shutdown_workers()
assert dataset._consumer_iterator_active is False
assert dataset._raw_batches_ref is None
second = iter(loader)
try:
with pytest.raises(StopIteration):
next(first)
assert next(second)["id"].tolist() == [2, 3]
finally:
second._shutdown_workers()
def test_direct_and_loader_admission_share_one_atomic_lease(tmp_path, monkeypatch):
db = lancedb.connect(tmp_path)
table = db.create_table("direct_loader_lease", pa.table({"id": list(range(4))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
entered = threading.Event()
release = threading.Event()
direct_result = []
direct_error = []
contender = []
real_resolve = dataset._resolve_my_splits
def controlled_resolve():
if threading.current_thread().name == "direct-start":
entered.set()
assert release.wait(timeout=5)
return real_resolve()
def advance_direct(iterator):
try:
direct_result.append(next(iterator)["id"])
except BaseException as exc:
direct_error.append(exc)
monkeypatch.setattr(dataset, "_resolve_my_splits", controlled_resolve)
direct = iter(dataset)
thread = threading.Thread(
target=advance_direct, args=(direct,), name="direct-start"
)
thread.start()
assert entered.wait(timeout=5)
try:
with pytest.raises(RuntimeError, match="concurrent iteration"):
contender.append(iter(loader))
finally:
release.set()
thread.join(timeout=5)
if contender:
contender[0]._shutdown_workers()
direct.close()
assert not thread.is_alive()
assert direct_error == []
assert direct_result == [0]
def test_loader_acquires_before_snapshot_and_cleans_interrupted_acquire(
tmp_path, monkeypatch
):
db = lancedb.connect(tmp_path)
table = db.create_table("lease_snapshot", pa.table({"id": list(range(4))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
first = iter(loader)
assert next(first)["id"].tolist() == [0, 1]
entered = threading.Event()
release = threading.Event()
pending = []
pending_errors = []
observed_snapshots = []
real_acquire = dataset._acquire_consumer_iterator
real_snapshot = dataset._checkpoint_snapshot
def controlled_acquire():
if threading.current_thread().name == "stale-start":
entered.set()
assert release.wait(timeout=5)
return real_acquire()
def recording_snapshot():
state = real_snapshot()
if threading.current_thread().name == "stale-start":
observed_snapshots.append(state["samples_consumed_per_split"])
return state
def create_pending_iterator():
try:
pending.append(iter(loader))
except BaseException as exc:
pending_errors.append(exc)
monkeypatch.setattr(dataset, "_acquire_consumer_iterator", controlled_acquire)
monkeypatch.setattr(dataset, "_checkpoint_snapshot", recording_snapshot)
thread = threading.Thread(target=create_pending_iterator, name="stale-start")
thread.start()
assert entered.wait(timeout=5)
assert next(first)["id"].tolist() == [2, 3]
with pytest.raises(StopIteration):
next(first)
release.set()
thread.join(timeout=5)
assert not thread.is_alive()
assert pending_errors == []
assert observed_snapshots == [[4]]
assert len(pending) == 1
assert list(pending[0]) == []
assert dataset.state_dict()["samples_consumed_per_split"] == [4]
def interrupted_acquire():
real_acquire()
raise KeyboardInterrupt("after acquire")
monkeypatch.setattr(dataset, "_acquire_consumer_iterator", interrupted_acquire)
with pytest.raises(KeyboardInterrupt, match="after acquire"):
iter(loader)
assert dataset._consumer_iterator_active is False
def test_consumer_iterator_lease_publication_is_atomic(tmp_path, monkeypatch):
db = lancedb.connect(tmp_path)
table = db.create_table("atomic_lease", pa.table({"id": [0, 1]}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(dataset, batch_size=1, num_workers=0)
real_get_ident = streaming.threading.get_ident
calls = 0
def interrupt_during_publication():
nonlocal calls
calls += 1
if calls == 1:
raise KeyboardInterrupt("during lease mutation")
return real_get_ident()
monkeypatch.setattr(streaming.threading, "get_ident", interrupt_during_publication)
with pytest.raises(KeyboardInterrupt, match="during lease mutation"):
iter(loader)
monkeypatch.setattr(streaming.threading, "get_ident", real_get_ident)
assert dataset._consumer_iterator_active is False
iterator = iter(loader)
try:
assert next(iterator)["id"].tolist() == [0]
finally:
iterator._shutdown_workers()
def test_streaming_dataloader_rejects_dataset_iter_override(tmp_path):
class CustomizedDataset(StreamingDataset):
def __iter__(self):
return iter([1000, 1001])
db = lancedb.connect(tmp_path)
table = db.create_table("custom_iteration", pa.table({"id": [0, 1, 2]}))
dataset = CustomizedDataset(table, num_splits=1, shuffle=False)
assert list(dataset) == [1000, 1001]
with pytest.raises(TypeError, match="override __iter__"):
StreamingDataLoader(
dataset,
batch_size=2,
num_workers=0,
collate_fn=list,
)
def test_interleaved_adapters_do_not_authorize_plain_iteration(tmp_path):
db = lancedb.connect(tmp_path)
table_a = db.create_table("adapter_a", pa.table({"id": [0, 1]}))
table_b = db.create_table("adapter_b", pa.table({"id": [10, 11]}))
dataset_a = StreamingDataset(table_a, num_splits=1, shuffle=False)
dataset_b = StreamingDataset(table_b, num_splits=1, shuffle=False)
initial_state = dataset_a.state_dict()
owner_a = dataset_a._acquire_consumer_iterator()
owner_b = dataset_b._acquire_consumer_iterator()
try:
iterator_a = iter(streaming._StreamingDatasetAdapter(dataset_a))
iterator_b = iter(streaming._StreamingDatasetAdapter(dataset_b))
assert next(iterator_a).data["id"] == 0
assert next(iterator_b).data["id"] == 10
assert [sample.data["id"] for sample in iterator_a] == [1]
assert [sample.data["id"] for sample in iterator_b] == [11]
finally:
dataset_a._release_consumer_iterator(owner_a)
dataset_b._release_consumer_iterator(owner_b)
dataset_a.load_state_dict(initial_state)
with patch(
"lancedb.streaming.get_worker_info",
return_value=FakeWorkerInfo(id=0, num_workers=1),
):
plain_iterator = iter(dataset_a)
assert next(plain_iterator)["id"] == 0
plain_iterator.close()
assert dataset_a._untracked_worker_iteration[0] == 1
with pytest.raises(RuntimeError, match="Use StreamingDataLoader"):
dataset_a.state_dict()
def test_resume_from_partial_split_cycle_preserves_remaining_order(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("partial_cycle", pa.table({"id": [1, 2, 10, 20]}))
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
iterator = iter(dataset)
assert next(iterator)["id"] == 1
checkpoint = dataset.state_dict()
iterator.close()
assert checkpoint["samples_consumed_per_split"] == [1, 0]
resumed = StreamingDataset(table, num_splits=2, shuffle=False)
resumed.load_state_dict(checkpoint)
assert [row["id"] for row in resumed] == [10, 2, 20]
def test_partial_cycle_resume_preserves_skip_truncation(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table(
"partial_skip", pa.table({"id": [0, 1, 2, 3, 100, 101, 102, 103]})
)
kwargs = dict(
num_splits=2,
shuffle=False,
transform=_failing_transform({1, 2, 3}),
on_transform_error="skip",
)
dataset = StreamingDataset(table, **kwargs)
iterator = iter(dataset)
assert next(iterator)["id"] == 0
checkpoint = dataset.state_dict()
uninterrupted = [row["id"] for row in iterator]
resumed = StreamingDataset(table, **kwargs)
resumed.load_state_dict(checkpoint)
assert [row["id"] for row in resumed] == uninterrupted == [100]
def test_multi_worker_resumability_same_topology(lance_table):
"""Checkpoint with num_workers=2, resume with num_workers=2: exact continuation."""
world_size = 1
@@ -2018,6 +2600,23 @@ def test_merge_state_dicts_validates_consistency(lance_table):
StreamingDataset.merge_state_dicts([])
def test_merge_state_dicts_combines_nonuniform_consumer_progress(lance_table):
dataset = StreamingDataset(
lance_table, num_splits=2, shuffle=False, shuffle_seed=SHUFFLE_SEED
)
rank0 = dataset.state_dict()
rank0["samples_consumed_per_split"] = [2, 0]
rank0["positions_consumed_per_split"] = [2, 0]
rank1 = dataset.state_dict()
rank1["samples_consumed_per_split"] = [0, 2]
rank1["positions_consumed_per_split"] = [0, 2]
merged = StreamingDataset.merge_state_dicts([rank0, rank1])
assert merged["samples_consumed_per_split"] == [2, 2]
assert merged["positions_consumed_per_split"] == [2, 2]
def test_load_state_dict_without_positions_key(lance_table):
"""Checkpoints from before positions_consumed_per_split existed still
resume exactly (positions equal sample counts when nothing is skipped)."""
@@ -2254,6 +2853,65 @@ def test_pack_sequences_checkpoint_resumes_on_new_topology(tmp_path):
]
def test_packed_checkpoint_requires_complete_split_cycle(tmp_path):
table = _create_token_table(tmp_path, [[1], [2], [10], [20]])
dataset = _packed_dataset(table, pack_sequences=3, blocks_per_epoch=4, num_splits=2)
iterator = iter(dataset)
next(iterator)
with pytest.raises(RuntimeError, match="complete logical step boundary"):
dataset.state_dict()
next(iterator)
assert dataset.state_dict()["blocks_emitted_per_split"] == [1, 1]
iterator.close()
def test_streaming_dataloader_commits_consumed_packed_batches(tmp_path):
table = _create_token_table(
tmp_path,
[[1], [2], [3], [4], [10], [20], [30], [40]],
)
kwargs = dict(pack_sequences=4, blocks_per_epoch=4, num_splits=2)
dataset = _packed_dataset(table, **kwargs)
loader = StreamingDataLoader(
dataset,
batch_size=1,
num_workers=2,
multiprocessing_context="spawn",
prefetch_factor=2,
)
iterator = iter(loader)
try:
next(iterator)
with pytest.raises(RuntimeError, match="complete logical step boundary"):
dataset.state_dict()
next(iterator)
checkpoint = dataset.state_dict()
uninterrupted = [batch["input_ids"].tolist() for batch in iterator]
finally:
iterator._shutdown_workers()
resumed = _packed_dataset(table, **kwargs)
resumed.load_state_dict(checkpoint)
resumed_loader = StreamingDataLoader(
resumed,
batch_size=1,
num_workers=2,
multiprocessing_context="spawn",
prefetch_factor=2,
)
resumed_iterator = iter(resumed_loader)
try:
remaining = [batch["input_ids"].tolist() for batch in resumed_iterator]
finally:
resumed_iterator._shutdown_workers()
assert checkpoint["blocks_emitted_per_split"] == [1, 1]
assert remaining == uninterrupted
def test_pack_sequences_validates_configuration_and_tokens(tmp_path):
table = _create_token_table(tmp_path, [[1, 2]])
python/python/tests/test_first_class_function_slice1.py
+5 -36
View File
@@ -37,21 +37,6 @@ def job_result(name: str) -> dict:
return json.loads(fixture(name))["result"]
def assert_no_secret_values(value):
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()
@@ -109,7 +94,6 @@ def test_function_version_identity_is_immutable_and_exact():
version = FunctionVersion.from_json(json.dumps(value))
assert version.name == "embed"
assert version.version == "fv_01K3EXACT"
assert version.required_secrets == ("HF_TOKEN",)
with pytest.raises((TypeError, ValueError)):
version.version = "fv_changed"
@@ -121,7 +105,7 @@ def test_function_version_identity_is_immutable_and_exact():
assert FunctionVersion(**changed) != version
def test_function_version_binds_named_columns_as_one_immutable_group():
def test_function_version_binds_named_columns_as_one_immutable_application():
version = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
)
@@ -131,13 +115,10 @@ def test_function_version_binds_named_columns_as_one_immutable_group():
assert application.function.name == version.name
assert application.function.version == version.version
assert application.output is version.signature.output
assert application.group_id.startswith("fg_")
assert [
(value.parameter, value.kind, value.value["path"])
for value in application.inputs
] == [("text", "column", "documents.body")]
with pytest.raises((TypeError, ValueError)):
application.group_id = "fg_changed"
def test_function_version_binding_validates_names_and_direct_columns():
@@ -156,7 +137,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_grouped"
value["version"] = "fv_multi_output"
value["signature"] = {
"inputs": [
{"name": "title", "arrow_type": "utf8", "nullable": True},
@@ -221,7 +202,6 @@ def test_function_application_uses_rename_columns_only():
assert application.columns["normalized_text"] == "search_text"
assert renamed.columns["normalized_text"] == "body_normalized"
assert renamed.function == application.function
assert renamed.group_id == application.group_id
assert not hasattr(application, "rename_outputs")
with pytest.raises(TypeError, match="immutable"):
renamed.columns["normalized_text"] = "changed"
@@ -242,7 +222,6 @@ 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.revision == 3
assert binding.function.version == "fv_01K3TEXT"
assert [output.output_ordinal for output in binding.outputs] == [0, 1]
assert binding.input_schema is not None
@@ -297,15 +276,6 @@ def test_refresh_result_rejects_non_u64_values(field):
RefreshColumnResult.from_json(json.dumps(value))
def test_canonical_client_values_contain_secret_names_only():
version = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
)
canonical = json.loads(version.to_canonical_json())
assert canonical["required_secrets"] == ["HF_TOKEN"]
assert_no_secret_values(canonical)
class _FunctionDeclarationInner:
def __init__(self):
self.calls = []
@@ -322,7 +292,7 @@ def known_application() -> FunctionApplication:
@pytest.mark.asyncio
async def test_add_columns_routes_struct_as_one_and_grouped_expansion_atomically():
async def test_add_columns_routes_struct_as_one_and_multi_output_binding_atomically():
inner = _FunctionDeclarationInner()
table = AsyncTable(inner)
application = known_application()
@@ -343,12 +313,12 @@ async def test_add_columns_routes_struct_as_one_and_grouped_expansion_atomically
@pytest.mark.asyncio
async def test_add_columns_rejects_mixed_groups_and_unknown_newer_application():
async def test_add_columns_rejects_multiple_bindings_and_unknown_newer_application():
inner = _FunctionDeclarationInner()
table = AsyncTable(inner)
application = known_application()
with pytest.raises(ValueError, match="exactly one Function sibling group"):
with pytest.raises(ValueError, match="exactly one Function binding"):
await table.add_columns({"a": application, "b": application})
future = json.loads(fixture("remote_function_application.json"))
@@ -376,7 +346,6 @@ def test_rename_requires_named_struct_and_keeps_partial_mapping_immutable():
"arrow_type": "list<float32>",
"nullable": False,
},
"group_id": "fg_scalar",
}
)
)
python/python/tests/test_first_class_function_slice2.py
+396 -30
View File
@@ -3,7 +3,12 @@
from __future__ import annotations
import base64
import contextlib
import functools
import importlib.util
import types
from datetime import date
import http.server
import json
from pathlib import Path
@@ -16,6 +21,9 @@ import pytest
import lancedb
from lancedb.functions import UdfDefinition, udf
THRESHOLD = 20
_CACHE = None
FIXTURES = (
Path(__file__).parents[3]
@@ -31,28 +39,12 @@ FIXTURES = (
@udf(
pip=["numpy>=2"],
env={"MODE": "test"},
secrets=["API_TOKEN"],
python_version="3.12",
)
def normalize_score(value: float) -> float:
return value / 100.0
def _assert_no_secret_values(value):
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_scalar_udf_matches_shared_registration_golden_and_remains_callable():
assert isinstance(normalize_score, UdfDefinition)
assert normalize_score(25.0) == 0.25
@@ -67,13 +59,397 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
"kind": "scalar_to_arrow_batch",
"version": 1,
}
assert request["required_secrets"] == ["API_TOKEN"]
_assert_no_secret_values(request)
def _run_packaged(definition, *args):
"""Execute the shipped artifact in a fresh namespace, as a worker would."""
source = base64.b64decode(definition.registration_request.artifact.content.data)
namespace: dict = {}
exec(compile(source, "<udf>", "exec"), namespace)
return namespace[definition.registration_request.artifact.entrypoint](*args)
def test_udf_packages_attribute_access_and_body_imports():
@udf
def word_norm(body: str) -> float:
import numpy as np
try:
words = body.split()
except AttributeError as error:
raise ValueError(str(error)) from error
return float(np.linalg.norm([len(w) for w in words]))
assert _run_packaged(word_norm, "aa bb") == pytest.approx(8**0.5)
def test_udf_packages_module_globals_and_global_caches():
@udf
def label(value: int) -> str:
return "big" if value >= THRESHOLD else "small"
assert _run_packaged(label, 21) == "big"
@udf
def cached(value: int) -> int:
global _CACHE
if _CACHE is None:
_CACHE = 40
return _CACHE + value
assert _run_packaged(cached, 2) == 42
def test_udf_annotations_are_not_runtime_names():
@udf
def identity(value: date) -> date:
return value
assert _run_packaged(identity, date(2026, 8, 25)) == date(2026, 8, 25)
def test_udf_nested_scopes_resolve_lexically():
@udf
def score(value: int) -> int:
offset = 2
def add_offset() -> int:
return value + offset
return add_offset() + sum(v for v in [0])
assert _run_packaged(score, 3) == 5
def test_udf_resolves_module_globals_before_builtins(tmp_path):
module_path = tmp_path / "shadowing_udfs.py"
module_path.write_text(
"max = 7\n"
"len = lambda _: 99\n"
"\n"
"def uses_literal_shadow(value: int) -> int:\n"
" def nested() -> int:\n"
" return max\n"
" return nested() + value\n"
"\n"
"def uses_callable_shadow(value: int) -> int:\n"
" def nested() -> int:\n"
" return len([1])\n"
" return nested() + value\n"
)
spec = importlib.util.spec_from_file_location("shadowing_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
# The module's `max = 7` is what the interpreter would use, so it ships.
assert _run_packaged(udf(module.uses_literal_shadow), 1) == 8
# A callable global cannot ship; it must not be silently swapped for the builtin.
with pytest.raises(TypeError, match="unsupported global value of type function"):
udf(module.uses_callable_shadow)
def test_canonical_arrow_type_is_exactly_the_grammar():
from lancedb.functions import _GRAMMAR_PRIMITIVES, _canonical_arrow_type
golden = json.loads(
(
Path(__file__).parents[3]
/ "rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json"
).read_text()
)
primitives = [
case["arrow_type"] for case in golden["valid"] if "<" not in case["arrow_type"]
]
assert [name for _, name in _GRAMMAR_PRIMITIVES] == primitives
for outside in [
pa.timestamp("us"),
pa.decimal128(10, 2),
pa.large_string(),
pa.large_binary(),
pa.binary(4),
pa.duration("s"),
pa.struct([pa.field("a", pa.int32())]),
pa.list_(pa.float32(), 0),
pa.list_(pa.timestamp("us")),
]:
with pytest.raises(TypeError, match="unsupported Arrow type"):
_canonical_arrow_type(outside)
def test_udf_nested_annotations_are_postponed_in_the_artifact():
@udf
def score(value: int) -> int:
def identity(item: date) -> date:
return item
identity(date(2026, 8, 25))
return value
assert _run_packaged(score, 3) == 3
def test_udf_ships_globals_the_body_deletes():
@udf
def clear(value: int) -> int:
global _CACHE
del _CACHE
return value
assert _run_packaged(clear, 3) == 3
def test_udf_rejects_a_module_global_that_does_not_import_as_itself(tmp_path):
module_path = tmp_path / "fake_module_udfs.py"
module_path.write_text(
"import types\n"
"np = types.ModuleType('numpy')\n"
"np.sqrt = lambda x: 0\n"
"\n"
"def score(value: int) -> int:\n"
" return int(np.sqrt(value))\n"
)
spec = importlib.util.spec_from_file_location("fake_module_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
with pytest.raises(TypeError, match="does not import as 'numpy'"):
udf(module.score)
def test_udf_rejects_a_module_level_namespace_alias(tmp_path):
module_path = tmp_path / "aliasing_udfs.py"
module_path.write_text(
"import builtins as b\n"
"THRESHOLD = 5\n"
"\n"
"def score(value: int) -> int:\n"
" return value + b.vars(b.__import__('aliasing_udfs'))['THRESHOLD']\n"
)
spec = importlib.util.spec_from_file_location("aliasing_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
with pytest.raises(ValueError, match="dynamic namespace access"):
udf(module.score)
@pytest.mark.parametrize(
"access",
[
"globals()['THRESHOLD']",
"eval('THRESHOLD')",
"(lambda g: g()['THRESHOLD'])(globals)",
"__import__('sys').modules[__name__].THRESHOLD",
"sys.modules[__name__].THRESHOLD",
],
)
def test_udf_rejects_dynamic_namespace_access(access):
namespace: dict = {}
exec(
f"def score(value: int) -> int:\n return value + {access}\n",
{"THRESHOLD": 5},
namespace,
)
with pytest.raises(ValueError, match="dynamic namespace access"):
_package_from_text(
"def score(value: int) -> int:\n"
" import sys\n"
f" return value + {access}\n"
)
def _package_from_text(source: str, module_globals: dict | None = None):
"""Load `source` as a real module file so the packager can inspect it."""
import tempfile
directory = tempfile.mkdtemp()
path = Path(directory) / "generated_udf_module.py"
path.write_text(source)
spec = importlib.util.spec_from_file_location(f"generated_udf_{id(source)}", path)
module = importlib.util.module_from_spec(spec)
if module_globals:
module.__dict__.update(module_globals)
spec.loader.exec_module(module)
functions = [
value
for value in vars(module).values()
if callable(value) and getattr(value, "__module__", None) == module.__name__
]
return udf(functions[0])
def test_udf_rejects_a_non_standard_builtins_environment():
def score(value: int) -> int:
return len([1]) + value
score.__globals__ # noqa: B018 -- real function, real globals
import builtins
patched = types.FunctionType(
score.__code__,
{"__builtins__": {**vars(builtins), "len": lambda _: 99}},
"score",
)
patched.__annotations__ = score.__annotations__
assert patched(3) == 102
with pytest.raises(ValueError, match="non-standard builtins environment"):
udf(patched)
class ReportingDict(dict): # reports standard entries, resolves differently
def __missing__(self, key):
return vars(builtins)[key]
disguised = types.FunctionType(
score.__code__, {"__builtins__": ReportingDict(len=lambda _: 99)}, "score"
)
disguised.__annotations__ = score.__annotations__
assert disguised(3) == 102
with pytest.raises(ValueError, match="non-standard builtins environment"):
udf(disguised)
hooked = types.FunctionType(
score.__code__,
{"__builtins__": {**vars(builtins), "__import__": lambda *a, **k: None}},
"score",
)
hooked.__annotations__ = score.__annotations__
with pytest.raises(ValueError, match="non-standard builtins environment"):
udf(hooked)
def test_udf_recursion_versus_a_rebound_module_name(tmp_path):
module_path = tmp_path / "rebound_udfs.py"
module_path.write_text(
"def fact(value: int) -> int:\n"
" return 1 if value <= 1 else value * fact(value - 1)\n"
"\n"
"def score(value: int) -> int:\n"
" return score + value\n"
)
spec = importlib.util.spec_from_file_location("rebound_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
assert _run_packaged(udf(module.fact), 5) == 120
raw = module.score
module.score = 10
with pytest.raises(ValueError, match="binds that name to another value"):
udf(raw)
# A wrapper that merely exposes __wrapped__ is not the function.
module.score = functools.wraps(raw)(lambda value: 41)
with pytest.raises(ValueError, match="binds that name to another value"):
udf(raw)
# The decorator's own result is; a subclass of it is not.
module.fact = udf(module.fact)
assert _run_packaged(module.fact, 4) == 24
class Twisted(UdfDefinition):
def __call__(self, *args, **kwargs):
return 41
raw_fact = module.fact._function
module.fact = Twisted(
raw_fact,
name=None,
input_schema=None,
output_schema=None,
pip=(),
env={},
python_version=None,
)
with pytest.raises(ValueError, match="binds that name to another value"):
udf(raw_fact)
def test_canonical_arrow_type_rejects_unrepresentable_list_children():
from lancedb.functions import _canonical_arrow_type
for outside in [
pa.list_(pa.float32()), # pyarrow default: nullable child
pa.list_(pa.field("custom", pa.float32(), nullable=False)),
pa.list_(pa.field("item", pa.float32(), nullable=False, metadata={"k": "v"})),
pa.list_(pa.field("item", pa.float32(), nullable=False), 0),
]:
with pytest.raises(TypeError, match="unsupported Arrow type"):
_canonical_arrow_type(outside)
assert (
_canonical_arrow_type(
pa.list_(pa.field("item", pa.float32(), nullable=False), 3)
)
== "fixed_size_list<float32, 3>"
)
def _calls_missing(value: int) -> int:
return missing(value) # noqa: F821
def _shadows_missing_in_a_comprehension(value: int) -> int:
return missing(value) + sum(missing for missing in ()) # noqa: F821
def _shadows_missing_in_a_lambda(value: int) -> int:
return (lambda missing: missing)(value) + missing # noqa: F821
@pytest.mark.parametrize(
"function",
[_calls_missing, _shadows_missing_in_a_comprehension, _shadows_missing_in_a_lambda],
)
def test_udf_rejects_a_truly_unresolved_global(function):
with pytest.raises(ValueError, match=r"unresolved global names: \['missing'\]"):
udf(function)
def _arrow_type_from_golden(spec: dict) -> pa.DataType:
kind = spec["type"]
if kind in ("list", "large_list", "fixed_size_list"):
item = _arrow_type_from_golden(spec["fields"][0]["type"])
field = pa.field("item", item, nullable=False)
if kind == "list":
return pa.list_(field)
if kind == "large_list":
return pa.large_list(field)
return pa.list_(field, spec["length"])
return {
"null": pa.null(),
"bool": pa.bool_(),
"utf8": pa.string(),
"binary": pa.binary(),
"float16": pa.float16(),
"float32": pa.float32(),
"float64": pa.float64(),
"date32": pa.date32(),
"date64": pa.date64(),
}.get(kind) or getattr(pa, kind)()
def test_arrow_type_grammar_matches_the_shared_golden():
golden = json.loads(
(
Path(__file__).parents[3]
/ "rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json"
).read_text()
)
from lancedb.functions import _canonical_arrow_type
emitted = {
case["arrow_type"]: _canonical_arrow_type(_arrow_type_from_golden(case["json"]))
for case in golden["valid"]
}
assert emitted == {
case["arrow_type"]: case["arrow_type"] for case in golden["valid"]
}
assert not set(emitted) & set(golden["invalid"])
for case in golden["server_only"]:
with pytest.raises(TypeError, match="unsupported Arrow type"):
_canonical_arrow_type(_arrow_type_from_golden(case["json"]))
def test_explicit_arrow_schema_is_deterministic():
input_schema = pa.schema([pa.field("value", pa.float32(), nullable=True)])
output_schema = pa.field("embedding", pa.list_(pa.float32(), 3), nullable=False)
output_schema = pa.field(
"embedding",
pa.list_(pa.field("item", pa.float32(), nullable=False), 3),
nullable=False,
)
@udf(input_schema=input_schema, output_schema=output_schema)
def explicit(value):
@@ -82,7 +458,7 @@ def test_explicit_arrow_schema_is_deterministic():
signature = explicit.registration_request.signature
assert signature.inputs[0].arrow_type == "float32"
assert signature.inputs[0].nullable is True
assert signature.output.arrow_type == "fixed_size_list<float32>[3]"
assert signature.output.arrow_type == "fixed_size_list<float32, 3>"
assert signature.output.nullable is False
@@ -130,14 +506,6 @@ def test_annotation_and_explicit_schema_validation_fail_closed():
return value
def test_environment_rejects_secret_value_overlap():
with pytest.raises(ValueError, match="must be disjoint"):
@udf(env={"TOKEN": "plaintext"}, secrets=["TOKEN"])
def overlapping(value: int) -> int:
return value
def test_local_function_catalog_operations_are_not_supported(tmp_path):
db = lancedb.connect(tmp_path)
message = "Function catalog operations are not supported by this database"
@@ -174,7 +542,6 @@ def _mock_remote_function_catalog():
"runtime": body["runtime"],
"runtime_digest": "sha256:runtime",
"environment_digest": "sha256:environment",
"required_secrets": body.get("required_secrets", []),
"created_at": "2026-08-21T00:00:00Z",
}
response = {"job_id": "job-register"}
@@ -233,7 +600,6 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip():
assert create_request == json.loads(
normalize_score.registration_request.to_canonical_json()
)
_assert_no_secret_values(create_request)
def test_blocking_remote_registration_returns_function_version():
python/python/tests/test_permutation.py
+25
View File
@@ -56,6 +56,31 @@ def test_execute_does_not_reenter_background_loop(tmp_path, monkeypatch):
assert permutation_tbl._conn.read_consistency_interval is None
def test_pickled_permutation_reads_pinned_version(tmp_path):
"""An unpickled copy must still read the pinned version, which also covers the
version surviving the ``to_arrow()`` round trip in ``__getstate__``."""
import pickle
db = connect(tmp_path)
tbl = db.create_table("base", pa.table({"idx": range(20)}))
permutation_tbl = permutation_builder(tbl).execute()
perm = Permutation.from_tables(tbl, permutation_tbl)
payload = pickle.dumps(perm)
# Compact so the stored row addresses no longer describe these rows at latest.
tbl.delete("true")
tbl.optimize()
assert tbl.count_rows() == 0
# Unpickle after the mutation: __setstate__ reopens at latest, so this only
# passes if the recorded version is applied on reopen.
restored = pickle.loads(payload)
assert len(restored) == 20
rows = restored.__getitems__(list(range(20)))
assert sorted(row["idx"] for row in rows) == list(range(20))
def test_split_random_counts(mem_db):
"""Test random splitting with absolute counts."""
tbl = mem_db.create_table(
python/src/table.rs
+6 -2
View File
@@ -780,15 +780,19 @@ impl Table {
})
}
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None))]
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None, allow_external_blob_outside_bases=false))]
pub fn add<'a>(
self_: PyRef<'a, Self>,
data: PyScannable,
mode: String,
progress: Option<Py<PyAny>>,
write_parallelism: Option<usize>,
allow_external_blob_outside_bases: bool,
) -> PyResult<Bound<'a, PyAny>> {
let mut op = self_.inner_ref()?.add(data);
let mut op = self_
.inner_ref()?
.add(data)
.allow_external_blob_outside_bases(allow_external_blob_outside_bases);
if mode == "append" {
op = op.mode(AddDataMode::Append);
} else if mode == "overwrite" {