From 403e4b603198824cdae5db5ab36bd81cfeda58f6 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Fri, 19 Jun 2026 18:01:54 +0800 Subject: [PATCH 01/25] feat: add lance namespace read write support --- README.md | 23 +++ daft_lance/__init__.py | 52 +++++ daft_lance/_lance.py | 15 +- daft_lance/lance_data_sink.py | 72 +++++-- daft_lance/lance_scan.py | 21 +- daft_lance/namespace.py | 217 ++++++++++++++++++++ daft_lance/utils.py | 40 +++- tests/io/lancedb/test_namespace.py | 84 ++++++++ tests/io/lancedb/test_namespace_rest_e2e.py | 94 +++++++++ 9 files changed, 597 insertions(+), 21 deletions(-) create mode 100644 daft_lance/namespace.py create mode 100644 tests/io/lancedb/test_namespace.py create mode 100644 tests/io/lancedb/test_namespace_rest_e2e.py diff --git a/README.md b/README.md index 11e7c3f..a7133b6 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,29 @@ from daft_lance import merge_columns_df merge_columns_df(df, "s3://bucket/my_dataset") ``` +### Namespace Tables + +```python +import daft +import daft_lance # installs daft.read_lance / DataFrame.write_lance namespace support + +table_id = ["my_table"] +namespace_properties = {"root": "/tmp/lance_tables"} + +df.write_lance( + namespace_impl="dir", + namespace_properties=namespace_properties, + table_id=table_id, + mode="create", +).collect() + +df = daft.read_lance( + namespace_impl="dir", + namespace_properties=namespace_properties, + table_id=table_id, +) +``` + ## Migration The migration only requires replacing `daft.io.lance` with `daft_lance`. diff --git a/daft_lance/__init__.py b/daft_lance/__init__.py index cb07862..8a1ecba 100644 --- a/daft_lance/__init__.py +++ b/daft_lance/__init__.py @@ -1,3 +1,7 @@ +from __future__ import annotations + +from typing import Any, Literal + try: import daft except ImportError: @@ -11,6 +15,54 @@ merge_columns_df, read_lance, ) +from .lance_data_sink import LanceDataSink + + +def _patch_daft_lance_api() -> None: + """Expose the daft-lance implementation through Daft's convenience APIs.""" + daft.read_lance = read_lance # type: ignore[assignment] + + from daft.dataframe import DataFrame + + original_write_lance = getattr(DataFrame, "write_lance") + if getattr(original_write_lance, "_daft_lance_namespace_patch", False): + return + + def write_lance( + self: DataFrame, + uri: Any = None, + mode: Literal["create", "append", "overwrite", "merge"] = "create", + io_config: Any | None = None, + schema: Any | None = None, + left_on: str | None = None, + right_on: str | None = None, + **kwargs: Any, + ) -> DataFrame: + if mode == "merge": + if any(k in kwargs for k in ("namespace_impl", "namespace_properties", "table_id")): + raise ValueError("write_lance(mode='merge') does not support namespace parameters yet.") + return original_write_lance( + self, + uri, + mode=mode, + io_config=io_config, + schema=schema, + left_on=left_on, + right_on=right_on, + **kwargs, + ) + + if schema is None: + schema = self.schema() + sanitized_kwargs = {k: v for k, v in kwargs.items() if k not in ("left_on", "right_on")} + sink = LanceDataSink(uri, schema, mode, io_config, **sanitized_kwargs) + return self.write_sink(sink) + + write_lance._daft_lance_namespace_patch = True # type: ignore[attr-defined] + DataFrame.write_lance = write_lance # type: ignore[assignment] + + +_patch_daft_lance_api() __all__ = [ "compact_files", diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 1bb7430..c0bc1f8 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -30,7 +30,7 @@ @PublicAPI def read_lance( - uri: str | pathlib.Path, + uri: str | pathlib.Path | None = None, io_config: IOConfig | None = None, version: str | int | None = None, asof: str | None = None, @@ -42,6 +42,10 @@ def read_lance( fragment_group_size: int | None = None, include_fragment_id: bool | None = None, checkpoint: CheckpointConfig | None = None, + *, + table_id: list[str] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, ) -> DataFrame: """Create a DataFrame from a LanceDB table. @@ -123,8 +127,8 @@ def read_lance( >>> df = daft.read_lance("s3://daft-oss-public-data/lance/words-test-dataset", io_config=io_config) >>> df.show() """ - uri_str = str(uri) - if uri_str.startswith("rest://"): + uri_str = str(uri) if uri is not None else None + if uri_str is not None and uri_str.startswith("rest://"): raise ValueError( "rest:// Lance URIs are no longer supported by daft.read_lance. " "The previous REST-namespace integration did not match the real " @@ -133,11 +137,14 @@ def read_lance( ) io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - storage_options = io_config_to_storage_options(io_config, uri_str) + storage_options = io_config_to_storage_options(io_config, uri_str) if uri_str is not None else None ds = construct_lance_dataset( uri_str, storage_options=storage_options, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, version=version, asof=asof, block_size=block_size, diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index 62c36bb..2ec69a7 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -26,6 +26,13 @@ detect_blob_v2_columns, resolve_storage_version, ) +from daft_lance.namespace import ( + get_namespace_kwargs, + get_write_fragments_kwargs, + merge_storage_options, + resolve_namespace_table, + validate_uri_or_namespace, +) if TYPE_CHECKING: from collections.abc import Iterator @@ -40,11 +47,14 @@ class LanceDataSink(DataSink[list[FragmentMetadata]]): def __init__( self, - uri: str | pathlib.Path, + uri: str | pathlib.Path | None, schema: Schema | pa.Schema, mode: Literal["create", "append", "overwrite"] = "create", io_config: IOConfig | None = None, *, + table_id: list[str] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, blob_columns: list[str] | None = None, max_rows_per_file: int = 1024 * 1024, max_rows_per_group: int = 1024, @@ -57,17 +67,26 @@ def __init__( compact_after_write: bool = True, ) -> None: self._reject_unsupported_modes(mode, use_legacy_format) - if not isinstance(uri, (str, pathlib.Path)): + validate_uri_or_namespace(uri, namespace_impl, table_id) + if uri is not None and not isinstance(uri, (str, pathlib.Path)): raise TypeError(f"Expected URI to be str or pathlib.Path, got {type(uri)}") - self._table_uri = str(uri) self._mode = mode + self._namespace_impl = namespace_impl + self._namespace_properties = namespace_properties + self._table_id = table_id self._io_config = get_context().daft_planning_config.default_io_config if io_config is None else io_config - self._storage_options = ( - storage_options - if storage_options is not None - else io_config_to_storage_options(self._io_config, self._table_uri) + base_storage_options = ( + ( + storage_options + if storage_options is not None + else io_config_to_storage_options(self._io_config, str(uri)) + ) + if uri is not None + else storage_options ) + self._table_uri, namespace_storage_options = self._resolve_table_uri(uri) + self._storage_options = merge_storage_options(base_storage_options, namespace_storage_options) self._init_lance_knobs( max_rows_per_file=max_rows_per_file, max_rows_per_group=max_rows_per_group, @@ -109,6 +128,28 @@ def __init__( ] ) + @property + def _namespace_kwargs(self) -> dict[str, object]: + return get_namespace_kwargs(self._namespace_impl, self._namespace_properties, self._table_id) + + @property + def _dataset_uri_arg(self) -> str | None: + return None if self._namespace_impl is not None and self._table_id is not None else self._table_uri + + def _resolve_table_uri(self, uri: str | pathlib.Path | None) -> tuple[str, dict[str, str] | None]: + if uri is not None: + return str(uri), None + mode = "create" if self._mode == "create" else "overwrite" if self._mode == "overwrite" else "read" + resolved_uri, namespace_storage_options = resolve_namespace_table( + namespace_impl=self._namespace_impl, + namespace_properties=self._namespace_properties, + table_id=self._table_id, + mode=mode, + ) + if resolved_uri is None: + raise ValueError("Unable to resolve Lance dataset URI from namespace.") + return resolved_uri, namespace_storage_options + @staticmethod def _reject_unsupported_modes( mode: Literal["create", "append", "overwrite"], use_legacy_format: bool | None @@ -167,7 +208,9 @@ def _absorb_existing_dataset(self) -> lance.LanceDataset | None: """ dataset: lance.LanceDataset | None try: - dataset = lance.dataset(self._table_uri, storage_options=self._storage_options) + dataset = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + ) except (ValueError, FileNotFoundError, OSError) as e: # Pinned to the Rust message format; lance has no typed exception. See test_lance_message_format_unchanged. if "was not found" in str(e): @@ -230,6 +273,7 @@ def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetada data_storage_version=self._data_storage_version, use_legacy_format=self._use_legacy_format, enable_stable_row_ids=self._enable_stable_row_ids, + **get_write_fragments_kwargs(self._namespace_impl, self._namespace_properties, self._table_id), ) # Sum on-disk sizes from fragment metadata. Lance Blob V2 sidecar .blob # files are not tracked in FragmentMetadata.files (out of scope here). @@ -240,7 +284,7 @@ def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetada def _ensure_mem_wal_dataset(self) -> lance.LanceDataset: try: - ds = lance.dataset(self._table_uri, storage_options=self._storage_options) + ds = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) except (ValueError, FileNotFoundError, OSError): ds = None @@ -250,11 +294,12 @@ def _ensure_mem_wal_dataset(self) -> lance.LanceDataset: {f.name: pa.array([], type=f.type) for f in self._effective_pyarrow_schema}, schema=self._effective_pyarrow_schema, ), - self._table_uri, + self._dataset_uri_arg, mode="create", storage_options=self._storage_options, data_storage_version=self._data_storage_version, use_legacy_format=self._use_legacy_format, + **self._namespace_kwargs, ) details = ds.mem_wal_index_details() @@ -335,6 +380,7 @@ def _finalize_cow(self, write_results: list[WriteResult[list[FragmentMetadata]]] operation, read_version=self._version, storage_options=self._storage_options, + **self._namespace_kwargs, ) stats = dataset.stats.dataset_stats() stats_dict = MicroPartition.from_pydict( @@ -348,7 +394,7 @@ def _finalize_cow(self, write_results: list[WriteResult[list[FragmentMetadata]]] return stats_dict def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadata]]]) -> MicroPartition: - dataset = lance.dataset(self._table_uri, storage_options=self._storage_options) + dataset = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) if self._compact_after_write: logger.info( @@ -359,7 +405,9 @@ def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadat from daft_lance.lance_compaction import compact_files_internal compact_files_internal(dataset) - dataset = lance.dataset(self._table_uri, storage_options=self._storage_options) + dataset = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + ) stats = dataset.stats.dataset_stats() return MicroPartition.from_pydict( diff --git a/daft_lance/lance_scan.py b/daft_lance/lance_scan.py index e0ea8fb..6acbd49 100644 --- a/daft_lance/lance_scan.py +++ b/daft_lance/lance_scan.py @@ -16,6 +16,7 @@ from daft.recordbatch import RecordBatch from ._metadata import convert_lance_schema +from .namespace import get_namespace_kwargs from .point_lookup import detect_point_lookup_columns from .utils import combine_filters_to_arrow @@ -40,7 +41,15 @@ def _lancedb_table_factory_function( "Use nearest with fragment_ids=None for index-driven global vector search." ) - ds = lance.dataset(ds_uri, **(open_kwargs or {})) + open_kwargs = dict(open_kwargs or {}) + namespace_impl = open_kwargs.pop("namespace_impl", None) + namespace_properties = open_kwargs.pop("namespace_properties", None) + table_id = open_kwargs.pop("table_id", None) + ds = lance.dataset( + None if namespace_impl is not None and table_id is not None else ds_uri, + **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), + **open_kwargs, + ) def _iter_batches() -> Iterator[PyRecordBatch]: # Iterate fragments individually; append a fragment_id column only when requested @@ -117,7 +126,15 @@ def _lancedb_count_result_function( filter: pa.compute.Expression | None = None, ) -> Iterator[PyRecordBatch]: """Use LanceDB's API to count rows and return a record batch with the count result.""" - ds = lance.dataset(ds_uri, **(open_kwargs or {})) + open_kwargs = dict(open_kwargs or {}) + namespace_impl = open_kwargs.pop("namespace_impl", None) + namespace_properties = open_kwargs.pop("namespace_properties", None) + table_id = open_kwargs.pop("table_id", None) + ds = lance.dataset( + None if namespace_impl is not None and table_id is not None else ds_uri, + **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), + **open_kwargs, + ) logger.debug("Using metadata for counting all rows") count = ds.count_rows(filter=filter) diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py new file mode 100644 index 0000000..2abe99c --- /dev/null +++ b/daft_lance/namespace.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +import os +from functools import lru_cache +from typing import Any +from urllib.parse import urlparse + +import lance + +_NAMESPACE_CACHE_SIZE = int(os.environ.get("DAFT_LANCE_NAMESPACE_CACHE_SIZE", "16")) +_PYLANCE_5 = (5, 0, 0) + + +@lru_cache(maxsize=1) +def _pylance_version() -> tuple[int, ...]: + version = getattr(lance, "__version__", "0.0.0") + parts = [] + for part in version.split(".")[:3]: + digits = "".join(ch for ch in part if ch.isdigit()) + parts.append(int(digits or "0")) + while len(parts) < 3: + parts.append(0) + return tuple(parts) + + +def has_namespace_params(namespace_impl: str | None, table_id: list[str] | None) -> bool: + return namespace_impl is not None and table_id is not None + + +def validate_uri_or_namespace( + uri: str | os.PathLike[str] | None, + namespace_impl: str | None, + table_id: list[str] | None, +) -> None: + has_uri = uri is not None + has_ns = has_namespace_params(namespace_impl, table_id) + + if namespace_impl is not None and table_id is None: + raise ValueError("'table_id' must be provided when 'namespace_impl' is provided.") + if table_id is not None and namespace_impl is None: + raise ValueError("'namespace_impl' must be provided when 'table_id' is provided.") + if has_uri and has_ns: + raise ValueError( + "Cannot provide both 'uri' and namespace parameters. Use either 'uri' OR ('namespace_impl' + 'table_id')." + ) + if not has_uri and not has_ns: + raise ValueError("Must provide either 'uri' OR ('namespace_impl' + 'table_id').") + + +def _normalize_file_uri(location: str) -> str: + parsed = urlparse(location) + if parsed.scheme == "file": + return parsed.path + return location + + +@lru_cache(maxsize=_NAMESPACE_CACHE_SIZE) +def _get_cached_namespace(namespace_impl: str, namespace_properties_tuple: tuple[tuple[str, str], ...] | None) -> Any: + import lance_namespace as ln + + namespace_properties = dict(namespace_properties_tuple) if namespace_properties_tuple else {} + return ln.connect(namespace_impl, namespace_properties) + + +def get_or_create_namespace(namespace_impl: str | None, namespace_properties: dict[str, str] | None) -> Any | None: + if namespace_impl is None: + return None + namespace_properties_tuple = tuple(sorted(namespace_properties.items())) if namespace_properties else None + return _get_cached_namespace(namespace_impl, namespace_properties_tuple) + + +def _create_storage_options_provider( + namespace_impl: str | None, + namespace_properties: dict[str, str] | None, + table_id: list[str] | None, +) -> Any | None: + if not has_namespace_params(namespace_impl, table_id): + return None + namespace = get_or_create_namespace(namespace_impl, namespace_properties) + if namespace is None or not hasattr(lance, "LanceNamespaceStorageOptionsProvider"): + return None + return lance.LanceNamespaceStorageOptionsProvider(namespace=namespace, table_id=table_id) + + +def get_namespace_kwargs( + namespace_impl: str | None, + namespace_properties: dict[str, str] | None, + table_id: list[str] | None, +) -> dict[str, Any]: + if not has_namespace_params(namespace_impl, table_id): + return {} + + namespace = get_or_create_namespace(namespace_impl, namespace_properties) + if namespace is None: + return {} + + kwargs: dict[str, Any] = {"table_id": table_id} + if _pylance_version() >= _PYLANCE_5: + kwargs["namespace_client"] = namespace + else: + kwargs["namespace"] = namespace + provider = _create_storage_options_provider(namespace_impl, namespace_properties, table_id) + if provider is not None: + kwargs["storage_options_provider"] = provider + return kwargs + + +def get_write_fragments_kwargs( + namespace_impl: str | None, + namespace_properties: dict[str, str] | None, + table_id: list[str] | None, +) -> dict[str, Any]: + if not has_namespace_params(namespace_impl, table_id): + return {} + if _pylance_version() >= _PYLANCE_5: + namespace = get_or_create_namespace(namespace_impl, namespace_properties) + if namespace is None: + return {} + return {"namespace_client": namespace, "table_id": table_id} + provider = _create_storage_options_provider(namespace_impl, namespace_properties, table_id) + return {"storage_options_provider": provider} if provider is not None else {} + + +def _storage_options(response: Any) -> dict[str, str] | None: + storage_options = getattr(response, "storage_options", None) + if storage_options is None: + return None + return dict(storage_options) + + +def _response_location(response: Any) -> str: + location = getattr(response, "location", None) or getattr(response, "table_uri", None) + if not location: + raise ValueError("Namespace response did not include a table location.") + return _normalize_file_uri(str(location)) + + +def _declare_table(namespace: Any, table_id: list[str]) -> tuple[str, dict[str, str] | None]: + try: + from lance_namespace import DeclareTableRequest + + response = namespace.declare_table(DeclareTableRequest(id=table_id, location=None)) + return _response_location(response), _storage_options(response) + except (AttributeError, NotImplementedError): + from lance_namespace import CreateEmptyTableRequest + + response = namespace.create_empty_table(CreateEmptyTableRequest(id=table_id)) + return _response_location(response), _storage_options(response) + + +def resolve_namespace_table( + *, + namespace_impl: str | None, + namespace_properties: dict[str, str] | None, + table_id: list[str] | None, + mode: str = "read", +) -> tuple[str | None, dict[str, str] | None]: + namespace = get_or_create_namespace(namespace_impl, namespace_properties) + if namespace is None or table_id is None: + return None, None + + if mode == "create": + return _declare_table(namespace, table_id) + + from lance_namespace import DescribeTableRequest + + try: + response = namespace.describe_table(DescribeTableRequest(id=table_id)) + return _response_location(response), _storage_options(response) + except Exception: + if mode == "overwrite": + return _declare_table(namespace, table_id) + raise + + +def merge_storage_options( + storage_options: dict[str, Any] | None, + namespace_storage_options: dict[str, str] | None, +) -> dict[str, Any] | None: + merged: dict[str, Any] = {} + if storage_options: + merged.update(storage_options) + if namespace_storage_options: + merged.update(namespace_storage_options) + return merged or None + + +def open_lance_dataset( + uri: str | os.PathLike[str] | None, + *, + storage_options: dict[str, Any] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, + table_id: list[str] | None = None, + mode: str = "read", + **kwargs: Any, +) -> lance.LanceDataset: + validate_uri_or_namespace(uri, namespace_impl, table_id) + resolved_uri = str(uri) if uri is not None else None + namespace_storage_options = None + if resolved_uri is None: + resolved_uri, namespace_storage_options = resolve_namespace_table( + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, + mode=mode, + ) + if resolved_uri is None: + raise ValueError("Unable to resolve Lance dataset URI.") + + merged_storage_options = merge_storage_options(storage_options, namespace_storage_options) + return lance.dataset( + None if has_namespace_params(namespace_impl, table_id) else resolved_uri, + storage_options=merged_storage_options, + **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), + **kwargs, + ) diff --git a/daft_lance/utils.py b/daft_lance/utils.py index 0eb32c7..3262f3a 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -7,6 +7,13 @@ from daft.dependencies import pa from daft.logical.schema import Schema as DaftSchema +from daft_lance.namespace import ( + get_namespace_kwargs, + has_namespace_params, + merge_storage_options, + resolve_namespace_table, + validate_uri_or_namespace, +) if TYPE_CHECKING: import pathlib @@ -82,12 +89,29 @@ def distribute_fragments_balanced(fragments: list[Any], fragment_group_size: int def construct_lance_dataset( - uri: str | pathlib.Path, + uri: str | pathlib.Path | None, version: int | str | None = None, storage_options: dict[str, Any] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, + table_id: list[str] | None = None, **kwargs: Any, ) -> lance.LanceDataset: """Construct a Lance dataset with common options.""" + validate_uri_or_namespace(uri, namespace_impl, table_id) + resolved_uri = str(uri) if uri is not None else None + namespace_storage_options = None + if resolved_uri is None: + resolved_uri, namespace_storage_options = resolve_namespace_table( + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, + mode="read", + ) + if resolved_uri is None: + raise ValueError("Unable to resolve Lance dataset URI.") + merged_storage_options = merge_storage_options(storage_options, namespace_storage_options) + original_default_scan_options = kwargs.pop("default_scan_options", None) safe_default_scan_options = None if isinstance(original_default_scan_options, dict): @@ -98,11 +122,21 @@ def construct_lance_dataset( # Non-dict defaults are forwarded as-is. kwargs["default_scan_options"] = original_default_scan_options - ds = lance.dataset(uri, storage_options=storage_options, version=version, **kwargs) + dataset_uri = None if has_namespace_params(namespace_impl, table_id) else resolved_uri + ds = lance.dataset( + dataset_uri, + storage_options=merged_storage_options, + version=version, + **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), + **kwargs, + ) effective_kwargs = { - "storage_options": storage_options, + "storage_options": merged_storage_options, "version": version, + "namespace_impl": namespace_impl, + "namespace_properties": namespace_properties, + "table_id": table_id, } effective_kwargs.update(kwargs or {}) try: diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py new file mode 100644 index 0000000..c573a56 --- /dev/null +++ b/tests/io/lancedb/test_namespace.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import pytest + +import daft +import daft_lance + + +def test_namespace_write_read_append_roundtrip(tmp_path): + namespace_properties = {"root": str(tmp_path)} + table_id = ["roundtrip"] + + df1 = daft.from_pydict({"id": [1, 2], "label": ["a", "b"]}) + df2 = daft.from_pydict({"id": [3], "label": ["c"]}) + + df1.write_lance( + namespace_impl="dir", + namespace_properties=namespace_properties, + table_id=table_id, + mode="create", + ).collect() + df2.write_lance( + namespace_impl="dir", + namespace_properties=namespace_properties, + table_id=table_id, + mode="append", + ).collect() + + result = daft.read_lance( + namespace_impl="dir", + namespace_properties=namespace_properties, + table_id=table_id, + ).to_pydict() + + assert result == {"id": [1, 2, 3], "label": ["a", "b", "c"]} + + +def test_namespace_read_supports_pushdowns(tmp_path): + namespace_properties = {"root": str(tmp_path)} + table_id = ["pushdowns"] + + daft.from_pydict( + { + "id": [1, 2, 3], + "label": ["a", "b", "c"], + "score": [10, 20, 30], + } + ).write_lance( + namespace_impl="dir", + namespace_properties=namespace_properties, + table_id=table_id, + mode="create", + ).collect() + + result = ( + daft.read_lance( + namespace_impl="dir", + namespace_properties=namespace_properties, + table_id=table_id, + ) + .where(daft.col("score") > 10) + .select("label") + .to_pydict() + ) + + assert result == {"label": ["b", "c"]} + + +def test_namespace_rejects_uri_and_namespace(tmp_path): + with pytest.raises(ValueError, match="Cannot provide both 'uri' and namespace parameters"): + daft_lance.read_lance( + str(tmp_path / "dataset"), + namespace_impl="dir", + namespace_properties={"root": str(tmp_path)}, + table_id=["tbl"], + ) + + +def test_namespace_requires_table_id(tmp_path): + with pytest.raises(ValueError, match="'table_id' must be provided"): + daft_lance.read_lance( + namespace_impl="dir", + namespace_properties={"root": str(tmp_path)}, + ) diff --git a/tests/io/lancedb/test_namespace_rest_e2e.py b/tests/io/lancedb/test_namespace_rest_e2e.py new file mode 100644 index 0000000..d9f27e9 --- /dev/null +++ b/tests/io/lancedb/test_namespace_rest_e2e.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +import os +import uuid + +import pytest + +import daft +import daft_lance # noqa: F401 - patches daft.read_lance and DataFrame.write_lance. + +pytestmark = pytest.mark.skipif( + os.environ.get("DAFT_LANCE_REST_URI") is None, + reason="Set DAFT_LANCE_REST_URI to run the Lance REST namespace integration test.", +) + + +def test_rest_namespace_write_read_append_roundtrip() -> None: + import lance + import lance_namespace as ln + from lance_namespace import CreateNamespaceRequest, DescribeTableRequest, NamespaceExistsRequest + + namespace_properties = {"uri": os.environ["DAFT_LANCE_REST_URI"]} + catalog = os.environ.get("DAFT_LANCE_REST_CATALOG", "lance_catalog") + schema = os.environ.get("DAFT_LANCE_REST_SCHEMA", "daft_ns_e2e") + table_id = [catalog, schema, f"orders_{uuid.uuid4().hex[:8]}"] + + namespace = ln.connect("rest", namespace_properties) + try: + namespace.namespace_exists(NamespaceExistsRequest(id=[catalog, schema])) + except Exception: + namespace.create_namespace(CreateNamespaceRequest(id=[catalog, schema], mode="CREATE")) + + daft.from_pydict( + { + "id": [1, 2, 3], + "label": ["a", "b", "c"], + "score": [10, 20, 30], + } + ).write_lance( + namespace_impl="rest", + namespace_properties=namespace_properties, + table_id=table_id, + mode="create", + ).collect() + + daft.from_pydict( + { + "id": [4, 5], + "label": ["d", "e"], + "score": [40, 50], + } + ).write_lance( + namespace_impl="rest", + namespace_properties=namespace_properties, + table_id=table_id, + mode="append", + ).collect() + + describe = namespace.describe_table(DescribeTableRequest(id=table_id)) + location = getattr(describe, "location", None) or getattr(describe, "table_uri", None) + assert location + + result = daft.read_lance( + namespace_impl="rest", + namespace_properties=namespace_properties, + table_id=table_id, + ).to_pydict() + assert result == { + "id": [1, 2, 3, 4, 5], + "label": ["a", "b", "c", "d", "e"], + "score": [10, 20, 30, 40, 50], + } + + filtered = ( + daft.read_lance( + namespace_impl="rest", + namespace_properties=namespace_properties, + table_id=table_id, + ) + .where(daft.col("score") >= 30) + .select("id", "label") + .to_pydict() + ) + assert filtered == {"id": [3, 4, 5], "label": ["c", "d", "e"]} + + assert ( + daft.read_lance( + namespace_impl="rest", + namespace_properties=namespace_properties, + table_id=table_id, + ).count_rows() + == 5 + ) + assert lance.dataset(None, namespace_client=namespace, table_id=table_id).count_rows() == 5 From aa34c099ef6ac9a32f75589cb70fa61c8eb44bf0 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Fri, 3 Jul 2026 18:24:14 +0900 Subject: [PATCH 02/25] feat(namespace): extend Lance Namespace support and harden implementation Build on the initial namespace read/write support: - Extend namespace params (namespace_impl/namespace_properties/table_id) to merge_columns, merge_columns_df, create_scalar_index, compact_files, threading namespace kwargs into their commit sites. - Add an explicit daft_lance.write_lance() and make the daft monkeypatch opt-in via patch_daft(); wire write_lance(mode="merge") through namespace. - Fix _declare_table fallback that always raised ImportError on the CreateEmptyTableRequest import; keep only declare_table. - Narrow overwrite resolution to only declare on TableNotFound instead of swallowing every describe_table exception. - Drop dead code (open_lance_dataset, pylance<5 compat) and dedupe the worker dataset-reopen logic into open_dataset_from_open_kwargs. - Cover all entry points with dir-namespace tests plus a Gravitino REST namespace e2e roundtrip. - Work around an upstream daft+lance native teardown SIGSEGV (unrelated to test outcomes) via a pytest_unconfigure hard-exit after reporting. --- README.md | 55 ++++- daft_lance/__init__.py | 52 ++--- daft_lance/_lance.py | 237 ++++++++++++++++++-- daft_lance/lance_merge_column.py | 5 + daft_lance/lance_scalar_index.py | 12 +- daft_lance/lance_scan.py | 22 +- daft_lance/namespace.py | 152 ++++++------- tests/io/lancedb/conftest.py | 27 +++ tests/io/lancedb/test_namespace.py | 193 +++++++++++++--- tests/io/lancedb/test_namespace_rest_e2e.py | 65 +++--- 10 files changed, 580 insertions(+), 240 deletions(-) diff --git a/README.md b/README.md index a7133b6..e88e5f1 100644 --- a/README.md +++ b/README.md @@ -40,25 +40,60 @@ merge_columns_df(df, "s3://bucket/my_dataset") ### Namespace Tables +Address Lance tables through a [Lance Namespace](https://lancedb.github.io/lance-namespace/) +(catalog) instead of a raw URI. Pass `namespace_impl` + `namespace_properties` + `table_id` +in place of `uri` — the namespace resolves the table's storage location and vends any storage +credentials. This works across `read_lance`, `write_lance`, `merge_columns_df`, +`create_scalar_index`, and `compact_files`. + ```python import daft -import daft_lance # installs daft.read_lance / DataFrame.write_lance namespace support +import daft_lance table_id = ["my_table"] -namespace_properties = {"root": "/tmp/lance_tables"} +namespace = {"namespace_impl": "dir", "namespace_properties": {"root": "/tmp/lance_tables"}} -df.write_lance( - namespace_impl="dir", - namespace_properties=namespace_properties, +daft_lance.write_lance( + daft.from_pydict({"id": [1, 2, 3]}), table_id=table_id, mode="create", + **namespace, ).collect() -df = daft.read_lance( - namespace_impl="dir", - namespace_properties=namespace_properties, - table_id=table_id, -) +df = daft_lance.read_lance(table_id=table_id, **namespace) +``` + +`uri` and the namespace parameters are mutually exclusive: provide exactly one of `uri` or +(`namespace_impl` + `table_id`). + +#### Using a REST namespace (e.g. Gravitino Lance REST server) + +```python +namespace = { + "namespace_impl": "rest", + "namespace_properties": {"uri": "http://127.0.0.1:9101/lance"}, +} +table_id = ["lance_catalog", "sales", "orders"] + +daft_lance.write_lance(df, table_id=table_id, mode="create", **namespace).collect() +daft_lance.read_lance(table_id=table_id, **namespace).show() +``` + +When the catalog holds the storage configuration (bucket, endpoint, credentials), the +`describe_table` response vends `storage_options` to the client, so you do not need to pass +object-store credentials yourself. + +#### Drop-in Daft APIs + +If you prefer Daft's own entry points, call `daft_lance.patch_daft()` once to route +`daft.read_lance` and `DataFrame.write_lance` through daft-lance (namespace parameters included): + +```python +import daft, daft_lance + +daft_lance.patch_daft() +df.write_lance(table_id=["my_table"], mode="create", **namespace).collect() +daft.read_lance(table_id=["my_table"], **namespace).show() ``` ## Migration diff --git a/daft_lance/__init__.py b/daft_lance/__init__.py index 8a1ecba..5af90b5 100644 --- a/daft_lance/__init__.py +++ b/daft_lance/__init__.py @@ -14,21 +14,25 @@ merge_columns, merge_columns_df, read_lance, + write_lance, ) from .lance_data_sink import LanceDataSink -def _patch_daft_lance_api() -> None: - """Expose the daft-lance implementation through Daft's convenience APIs.""" +def patch_daft() -> None: + """Route ``daft.read_lance`` and ``DataFrame.write_lance`` through daft-lance. + + Optional convenience for drop-in use of namespace parameters with Daft's own + APIs, e.g. ``df.write_lance(namespace_impl="dir", ...)``. Idempotent. + """ daft.read_lance = read_lance # type: ignore[assignment] from daft.dataframe import DataFrame - original_write_lance = getattr(DataFrame, "write_lance") - if getattr(original_write_lance, "_daft_lance_namespace_patch", False): + if getattr(DataFrame.write_lance, "_daft_lance_patch", False): return - def write_lance( + def _write_lance( self: DataFrame, uri: Any = None, mode: Literal["create", "append", "overwrite", "merge"] = "create", @@ -38,37 +42,29 @@ def write_lance( right_on: str | None = None, **kwargs: Any, ) -> DataFrame: - if mode == "merge": - if any(k in kwargs for k in ("namespace_impl", "namespace_properties", "table_id")): - raise ValueError("write_lance(mode='merge') does not support namespace parameters yet.") - return original_write_lance( - self, - uri, - mode=mode, - io_config=io_config, - schema=schema, - left_on=left_on, - right_on=right_on, - **kwargs, - ) - - if schema is None: - schema = self.schema() - sanitized_kwargs = {k: v for k, v in kwargs.items() if k not in ("left_on", "right_on")} - sink = LanceDataSink(uri, schema, mode, io_config, **sanitized_kwargs) - return self.write_sink(sink) - - write_lance._daft_lance_namespace_patch = True # type: ignore[attr-defined] - DataFrame.write_lance = write_lance # type: ignore[assignment] + return write_lance( + self, + uri, + mode=mode, + io_config=io_config, + schema=schema, + left_on=left_on, + right_on=right_on, + **kwargs, + ) + _write_lance._daft_lance_patch = True # type: ignore[attr-defined] + DataFrame.write_lance = _write_lance # type: ignore[method-assign] -_patch_daft_lance_api() __all__ = [ + "LanceDataSink", "compact_files", "create_scalar_index", "merge_columns", "merge_columns_df", + "patch_daft", "read_lance", "take_blobs", + "write_lance", ] diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index c0bc1f8..ab65a11 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -2,9 +2,7 @@ import pathlib from collections.abc import Callable -from typing import TYPE_CHECKING, Any - -import lance +from typing import TYPE_CHECKING, Any, Literal from daft import context from daft.api_annotations import PublicAPI @@ -15,9 +13,11 @@ from daft.logical.builder import LogicalPlanBuilder from .lance_compaction import compact_files_internal +from .lance_data_sink import LanceDataSink from .lance_merge_column import merge_columns_from_df, merge_columns_internal from .lance_scalar_index import create_scalar_index_internal from .lance_scan import LanceDBScanOperator +from .namespace import is_table_not_found, validate_uri_or_namespace from .utils import construct_lance_dataset if TYPE_CHECKING: @@ -28,6 +28,19 @@ from daft.dependencies import pa +def _effective_uri(lance_ds: LanceDataset, uri: str | pathlib.Path | None) -> str: + """The dataset location: the user-supplied URI, or the namespace-resolved one.""" + return str(uri) if uri is not None else str(lance_ds.uri) + + +def _effective_storage_options( + lance_ds: LanceDataset, storage_options: dict[str, Any] | None +) -> dict[str, Any] | None: + """Storage options actually used to open the dataset (includes namespace-vended ones).""" + open_kwargs = getattr(lance_ds, "_lance_open_kwargs", None) or {} + return open_kwargs.get("storage_options") or storage_options + + @PublicAPI def read_lance( uri: str | pathlib.Path | None = None, @@ -168,9 +181,12 @@ def read_lance( @PublicAPI def merge_columns( - uri: str | pathlib.Path, + uri: str | pathlib.Path | None = None, io_config: IOConfig | None = None, *, + table_id: list[str] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, transform: dict[str, str] | BatchUDF | Callable[[pa.lib.RecordBatch], pa.lib.RecordBatch] | None = None, read_columns: list[str] | None = None, reader_schema: pa.Schema | None = None, @@ -229,12 +245,16 @@ def merge_columns( ) io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - storage_options = storage_options or io_config_to_storage_options(io_config, uri) + if uri is not None: + storage_options = storage_options or io_config_to_storage_options(io_config, uri) # Build Lance dataset handle for committing lance_ds = construct_lance_dataset( uri, storage_options=storage_options, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, version=version, asof=asof, block_size=block_size, @@ -246,11 +266,11 @@ def merge_columns( return merge_columns_internal( lance_ds, - uri, + _effective_uri(lance_ds, uri), transform=transform, read_columns=read_columns, reader_schema=reader_schema, - storage_options=storage_options, + storage_options=_effective_storage_options(lance_ds, storage_options), daft_remote_args=daft_remote_args, concurrency=concurrency, ) @@ -259,9 +279,12 @@ def merge_columns( @PublicAPI def merge_columns_df( df: DataFrame, - uri: str | pathlib.Path, + uri: str | pathlib.Path | None = None, io_config: IOConfig | None = None, *, + table_id: list[str] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, read_columns: list[str] | None = None, reader_schema: pa.Schema | None = None, storage_options: dict[str, Any] | None = None, @@ -325,12 +348,16 @@ def merge_columns_df( >>> daft_lance.merge_columns_df(df, "s3://my-lancedb-bucket/data/") """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - storage_options = storage_options or io_config_to_storage_options(io_config, uri) + if uri is not None: + storage_options = storage_options or io_config_to_storage_options(io_config, uri) # Build Lance dataset handle for committing lance_ds = construct_lance_dataset( uri, storage_options=storage_options, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, version=version, asof=asof, block_size=block_size, @@ -349,10 +376,10 @@ def merge_columns_df( return merge_columns_from_df( df, lance_ds=lance_ds, - uri=uri, + uri=_effective_uri(lance_ds, uri), read_columns=read_columns, reader_schema=reader_schema, - storage_options=storage_options, + storage_options=_effective_storage_options(lance_ds, storage_options), daft_remote_args=daft_remote_args, concurrency=concurrency, left_on=left_on, @@ -363,9 +390,12 @@ def merge_columns_df( @PublicAPI def create_scalar_index( - uri: str | pathlib.Path, + uri: str | pathlib.Path | None = None, io_config: IOConfig | None = None, *, + table_id: list[str] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, column: str, index_type: str = "INVERTED", name: str | None = None, @@ -471,11 +501,15 @@ def create_scalar_index( >>> daft_lance.create_scalar_index("s3://my-bucket/dataset/", column="title", replace=False) """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - storage_options = storage_options or io_config_to_storage_options(io_config, str(uri)) + if uri is not None: + storage_options = storage_options or io_config_to_storage_options(io_config, str(uri)) lance_ds = construct_lance_dataset( uri, storage_options=storage_options, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, version=version, asof=asof, block_size=block_size, @@ -487,12 +521,12 @@ def create_scalar_index( create_scalar_index_internal( lance_ds=lance_ds, - uri=uri, + uri=_effective_uri(lance_ds, uri), column=column, index_type=index_type, name=name, replace=replace, - storage_options=storage_options, + storage_options=_effective_storage_options(lance_ds, storage_options), fragment_group_size=fragment_group_size, num_partitions=num_partitions, max_concurrency=max_concurrency, @@ -503,9 +537,12 @@ def create_scalar_index( @PublicAPI def compact_files( - uri: str | pathlib.Path, + uri: str | pathlib.Path | None = None, io_config: IOConfig | None = None, *, + table_id: list[str] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, storage_options: dict[str, Any] | None = None, version: int | str | None = None, asof: str | None = None, @@ -556,13 +593,17 @@ def compact_files( RuntimeError: When compaction fails or no successful results """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - storage_options = storage_options or io_config_to_storage_options( - io_config, str(uri) if isinstance(uri, pathlib.Path) else uri - ) + if uri is not None: + storage_options = storage_options or io_config_to_storage_options( + io_config, str(uri) if isinstance(uri, pathlib.Path) else uri + ) - lance_ds = lance.dataset( + lance_ds = construct_lance_dataset( uri, storage_options=storage_options, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, version=version, asof=asof, block_size=block_size, @@ -578,3 +619,159 @@ def compact_files( partition_num=partition_num, concurrency=concurrency, ) + + +def _is_missing_dataset_error(error: Exception) -> bool: + """Whether opening a dataset failed because it does not exist yet.""" + if isinstance(error, (FileNotFoundError,)): + return True + if isinstance(error, (ValueError, OSError)) and "was not found" in str(error): + return True + return is_table_not_found(error) + + +def _stats_dataframe(lance_ds: LanceDataset) -> DataFrame: + """Build the same stats DataFrame that LanceDataSink.finalize returns.""" + import pyarrow as _pa + + from daft.recordbatch import MicroPartition + + stats = lance_ds.stats.dataset_stats() + return DataFrame._from_micropartitions( + MicroPartition.from_pydict( + { + "num_fragments": _pa.array([stats["num_fragments"]], type=_pa.int64()), + "num_deleted_rows": _pa.array([stats["num_deleted_rows"]], type=_pa.int64()), + "num_small_files": _pa.array([stats["num_small_files"]], type=_pa.int64()), + "version": _pa.array([lance_ds.version], type=_pa.int64()), + } + ) + ) + + +@PublicAPI +def write_lance( + df: DataFrame, + uri: str | pathlib.Path | None = None, + mode: Literal["create", "append", "overwrite", "merge"] = "create", + io_config: IOConfig | None = None, + schema: Any | None = None, + left_on: str | None = None, + right_on: str | None = None, + *, + table_id: list[str] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, + **kwargs: Any, +) -> DataFrame: + """Write a DataFrame to a Lance table, addressed by URI or by Lance Namespace. + + Args: + df: The DataFrame to write. + uri: The URI of the Lance table. Mutually exclusive with the namespace parameters. + mode: One of "create", "append", "overwrite", or "merge" (add new columns to + existing rows; requires ``fragment_id`` and a join key column in ``df``). + io_config: A custom IOConfig to use when accessing Lance data. + schema: Desired schema to enforce during write; defaults to the DataFrame schema. + left_on/right_on: Join keys for ``mode="merge"``; default "_rowaddr". + table_id: Table identifier within the namespace, e.g. ["catalog", "schema", "table"]. + namespace_impl: Lance Namespace implementation, e.g. "dir" or "rest". + namespace_properties: Properties for connecting to the namespace, e.g. + {"root": "/data"} for "dir" or {"uri": "http://host:port"} for "rest". + **kwargs: Additional arguments forwarded to the Lance writer. + + Returns: + DataFrame: write statistics (num_fragments, num_deleted_rows, num_small_files, version). + + Examples: + >>> import daft, daft_lance + >>> df = daft.from_pydict({"id": [1, 2]}) + >>> daft_lance.write_lance( + ... df, namespace_impl="dir", namespace_properties={"root": "/tmp/tables"}, table_id=["t"] + ... ).collect() # doctest: +SKIP + """ + validate_uri_or_namespace(uri, namespace_impl, table_id) + + if schema is None: + schema = df.schema() + + if mode != "merge": + sanitized_kwargs = {k: v for k, v in kwargs.items() if k not in ("left_on", "right_on")} + sink = LanceDataSink( + uri, + schema, + mode, + io_config, + table_id=table_id, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + **sanitized_kwargs, + ) + return df.write_sink(sink) + + io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config + storage_options = kwargs.pop("storage_options", None) + if uri is not None: + storage_options = storage_options or io_config_to_storage_options(io_config, str(uri)) + + lance_ds: LanceDataset | None + try: + lance_ds = construct_lance_dataset( + uri, + storage_options=storage_options, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, + ) + except Exception as error: + if not _is_missing_dataset_error(error): + raise + lance_ds = None + + sanitized_kwargs = {k: v for k, v in kwargs.items() if k not in ("left_on", "right_on")} + + if lance_ds is None: + # Merge against a missing table degenerates to create. + sink = LanceDataSink( + uri, + schema, + "create", + io_config, + table_id=table_id, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + storage_options=storage_options, + **sanitized_kwargs, + ) + return df.write_sink(sink) + + existing_fields = {getattr(f, "name", str(f)) for f in lance_ds.schema} + meta_exclusions = {"fragment_id", "_rowaddr", "_rowid"} + new_cols = [c for c in df.column_names if c not in existing_fields and c not in meta_exclusions] + + if len(new_cols) == 0: + # No schema evolution: merge degenerates to append. + sink = LanceDataSink( + uri, + schema, + "append", + io_config, + table_id=table_id, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + storage_options=storage_options, + **sanitized_kwargs, + ) + return df.write_sink(sink) + + join_left = left_on or "_rowaddr" + join_right = right_on or join_left + merged = merge_columns_from_df( + df, + lance_ds=lance_ds, + uri=_effective_uri(lance_ds, uri), + left_on=join_left, + right_on=join_right, + storage_options=_effective_storage_options(lance_ds, storage_options), + ) + return _stats_dataframe(merged) diff --git a/daft_lance/lance_merge_column.py b/daft_lance/lance_merge_column.py index 8984120..df2d682 100644 --- a/daft_lance/lance_merge_column.py +++ b/daft_lance/lance_merge_column.py @@ -14,6 +14,8 @@ from daft.udf import cls as daft_cls from daft.udf import method +from .namespace import namespace_kwargs_for_dataset + if TYPE_CHECKING: import pathlib from collections.abc import Callable @@ -93,6 +95,7 @@ def merge_columns_internal( op, read_version=lance_ds.version, storage_options=storage_options, + **namespace_kwargs_for_dataset(lance_ds), ) @@ -473,6 +476,7 @@ def _merge_fast_path( op, read_version=lance_ds.version, storage_options=storage_options, + **namespace_kwargs_for_dataset(lance_ds), ) @@ -521,4 +525,5 @@ def _merge_slow_path( op, read_version=lance_ds.version, storage_options=storage_options, + **namespace_kwargs_for_dataset(lance_ds), ) diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index bd8e60c..b967ce6 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -14,6 +14,7 @@ import lance from daft.dependencies import pa +from daft_lance.namespace import namespace_kwargs_for_dataset from daft_lance.utils import distribute_fragments_balanced logger = logging.getLogger(__name__) @@ -386,7 +387,10 @@ def _create_segmented_index( # Reload dataset to pick up the latest version (segment files were written # by workers against the version that was current at their invocation time). - lance_ds = lance.LanceDataset(uri, storage_options=storage_options) + namespace_kwargs = namespace_kwargs_for_dataset(lance_ds) + lance_ds = lance.dataset( + None if namespace_kwargs else uri, storage_options=storage_options, **namespace_kwargs + ) index_metas = _prepare_index_segments_for_commit(lance_ds, index_type, index_metas) logger.info( @@ -460,7 +464,10 @@ def _create_partitioned_index( df.collect() logger.info("Starting index metadata merging by reloading dataset to get latest state") - lance_ds = lance.LanceDataset(uri, storage_options=storage_options) + namespace_kwargs = namespace_kwargs_for_dataset(lance_ds) + lance_ds = lance.dataset( + None if namespace_kwargs else uri, storage_options=storage_options, **namespace_kwargs + ) lance_ds.merge_index_metadata(index_id, index_type) logger.info("Starting atomic index creation and commit") @@ -506,6 +513,7 @@ def _create_partitioned_index( create_index_op, read_version=lance_ds.version, storage_options=storage_options, + **namespace_kwargs, ) logger.info("Index %s created successfully with ID %s", name, index_id) diff --git a/daft_lance/lance_scan.py b/daft_lance/lance_scan.py index 6acbd49..e350479 100644 --- a/daft_lance/lance_scan.py +++ b/daft_lance/lance_scan.py @@ -16,7 +16,7 @@ from daft.recordbatch import RecordBatch from ._metadata import convert_lance_schema -from .namespace import get_namespace_kwargs +from .namespace import open_dataset_from_open_kwargs from .point_lookup import detect_point_lookup_columns from .utils import combine_filters_to_arrow @@ -41,15 +41,7 @@ def _lancedb_table_factory_function( "Use nearest with fragment_ids=None for index-driven global vector search." ) - open_kwargs = dict(open_kwargs or {}) - namespace_impl = open_kwargs.pop("namespace_impl", None) - namespace_properties = open_kwargs.pop("namespace_properties", None) - table_id = open_kwargs.pop("table_id", None) - ds = lance.dataset( - None if namespace_impl is not None and table_id is not None else ds_uri, - **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), - **open_kwargs, - ) + ds = open_dataset_from_open_kwargs(ds_uri, open_kwargs) def _iter_batches() -> Iterator[PyRecordBatch]: # Iterate fragments individually; append a fragment_id column only when requested @@ -126,15 +118,7 @@ def _lancedb_count_result_function( filter: pa.compute.Expression | None = None, ) -> Iterator[PyRecordBatch]: """Use LanceDB's API to count rows and return a record batch with the count result.""" - open_kwargs = dict(open_kwargs or {}) - namespace_impl = open_kwargs.pop("namespace_impl", None) - namespace_properties = open_kwargs.pop("namespace_properties", None) - table_id = open_kwargs.pop("table_id", None) - ds = lance.dataset( - None if namespace_impl is not None and table_id is not None else ds_uri, - **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), - **open_kwargs, - ) + ds = open_dataset_from_open_kwargs(ds_uri, open_kwargs) logger.debug("Using metadata for counting all rows") count = ds.count_rows(filter=filter) diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index 2abe99c..3fd19c7 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -8,19 +8,6 @@ import lance _NAMESPACE_CACHE_SIZE = int(os.environ.get("DAFT_LANCE_NAMESPACE_CACHE_SIZE", "16")) -_PYLANCE_5 = (5, 0, 0) - - -@lru_cache(maxsize=1) -def _pylance_version() -> tuple[int, ...]: - version = getattr(lance, "__version__", "0.0.0") - parts = [] - for part in version.split(".")[:3]: - digits = "".join(ch for ch in part if ch.isdigit()) - parts.append(int(digits or "0")) - while len(parts) < 3: - parts.append(0) - return tuple(parts) def has_namespace_params(namespace_impl: str | None, table_id: list[str] | None) -> bool: @@ -69,56 +56,26 @@ def get_or_create_namespace(namespace_impl: str | None, namespace_properties: di return _get_cached_namespace(namespace_impl, namespace_properties_tuple) -def _create_storage_options_provider( - namespace_impl: str | None, - namespace_properties: dict[str, str] | None, - table_id: list[str] | None, -) -> Any | None: - if not has_namespace_params(namespace_impl, table_id): - return None - namespace = get_or_create_namespace(namespace_impl, namespace_properties) - if namespace is None or not hasattr(lance, "LanceNamespaceStorageOptionsProvider"): - return None - return lance.LanceNamespaceStorageOptionsProvider(namespace=namespace, table_id=table_id) - - def get_namespace_kwargs( namespace_impl: str | None, namespace_properties: dict[str, str] | None, table_id: list[str] | None, ) -> dict[str, Any]: + """Kwargs wiring a namespace client into pylance APIs (``lance.dataset``, ``commit``, ...). + + Requires pylance >= 7 which accepts ``namespace_client`` + ``table_id`` natively. + """ if not has_namespace_params(namespace_impl, table_id): return {} namespace = get_or_create_namespace(namespace_impl, namespace_properties) if namespace is None: return {} - - kwargs: dict[str, Any] = {"table_id": table_id} - if _pylance_version() >= _PYLANCE_5: - kwargs["namespace_client"] = namespace - else: - kwargs["namespace"] = namespace - provider = _create_storage_options_provider(namespace_impl, namespace_properties, table_id) - if provider is not None: - kwargs["storage_options_provider"] = provider - return kwargs + return {"namespace_client": namespace, "table_id": table_id} -def get_write_fragments_kwargs( - namespace_impl: str | None, - namespace_properties: dict[str, str] | None, - table_id: list[str] | None, -) -> dict[str, Any]: - if not has_namespace_params(namespace_impl, table_id): - return {} - if _pylance_version() >= _PYLANCE_5: - namespace = get_or_create_namespace(namespace_impl, namespace_properties) - if namespace is None: - return {} - return {"namespace_client": namespace, "table_id": table_id} - provider = _create_storage_options_provider(namespace_impl, namespace_properties, table_id) - return {"storage_options_provider": provider} if provider is not None else {} +# `lance.fragment.write_fragments` accepts the same namespace kwargs as `lance.dataset`. +get_write_fragments_kwargs = get_namespace_kwargs def _storage_options(response: Any) -> dict[str, str] | None: @@ -136,16 +93,27 @@ def _response_location(response: Any) -> str: def _declare_table(namespace: Any, table_id: list[str]) -> tuple[str, dict[str, str] | None]: - try: - from lance_namespace import DeclareTableRequest + from lance_namespace import DeclareTableRequest - response = namespace.declare_table(DeclareTableRequest(id=table_id, location=None)) - return _response_location(response), _storage_options(response) - except (AttributeError, NotImplementedError): - from lance_namespace import CreateEmptyTableRequest + response = namespace.declare_table(DeclareTableRequest(id=table_id, location=None)) + return _response_location(response), _storage_options(response) - response = namespace.create_empty_table(CreateEmptyTableRequest(id=table_id)) - return _response_location(response), _storage_options(response) + +def is_table_not_found(error: Exception) -> bool: + """Whether an exception from ``describe_table`` means the table does not exist. + + Native implementations raise the typed ``TableNotFoundError``; the Rust-backed + REST client surfaces untyped errors, so fall back to message matching. + """ + try: + from lance_namespace.errors import TableNotFoundError + + if isinstance(error, TableNotFoundError): + return True + except ImportError: + pass + message = str(error).lower() + return "not found" in message or "does not exist" in message or "no such table" in message def resolve_namespace_table( @@ -155,6 +123,11 @@ def resolve_namespace_table( table_id: list[str] | None, mode: str = "read", ) -> tuple[str | None, dict[str, str] | None]: + """Resolve a namespace table to ``(location, storage_options)``. + + ``mode="create"`` declares the table; ``mode="overwrite"`` declares it only when + it does not exist yet; other modes require the table to exist. + """ namespace = get_or_create_namespace(namespace_impl, namespace_properties) if namespace is None or table_id is None: return None, None @@ -167,8 +140,8 @@ def resolve_namespace_table( try: response = namespace.describe_table(DescribeTableRequest(id=table_id)) return _response_location(response), _storage_options(response) - except Exception: - if mode == "overwrite": + except Exception as error: + if mode == "overwrite" and is_table_not_found(error): return _declare_table(namespace, table_id) raise @@ -185,33 +158,40 @@ def merge_storage_options( return merged or None -def open_lance_dataset( - uri: str | os.PathLike[str] | None, - *, - storage_options: dict[str, Any] | None = None, - namespace_impl: str | None = None, - namespace_properties: dict[str, str] | None = None, - table_id: list[str] | None = None, - mode: str = "read", - **kwargs: Any, -) -> lance.LanceDataset: - validate_uri_or_namespace(uri, namespace_impl, table_id) - resolved_uri = str(uri) if uri is not None else None - namespace_storage_options = None - if resolved_uri is None: - resolved_uri, namespace_storage_options = resolve_namespace_table( - namespace_impl=namespace_impl, - namespace_properties=namespace_properties, - table_id=table_id, - mode=mode, - ) - if resolved_uri is None: - raise ValueError("Unable to resolve Lance dataset URI.") +def pop_namespace_params(open_kwargs: dict[str, Any]) -> tuple[str | None, dict[str, str] | None, list[str] | None]: + """Remove and return the namespace triple from a ``_lance_open_kwargs`` dict.""" + return ( + open_kwargs.pop("namespace_impl", None), + open_kwargs.pop("namespace_properties", None), + open_kwargs.pop("table_id", None), + ) + + +def namespace_kwargs_for_dataset(ds: Any) -> dict[str, Any]: + """Namespace kwargs for commit-style calls, recovered from ``ds._lance_open_kwargs``. + + ``lance.LanceDataset.commit`` is a static method, so a dataset opened through a + namespace does not carry its client into commits; re-derive it from the open + kwargs stashed by ``construct_lance_dataset``. + """ + open_kwargs = getattr(ds, "_lance_open_kwargs", None) or {} + return get_namespace_kwargs( + open_kwargs.get("namespace_impl"), + open_kwargs.get("namespace_properties"), + open_kwargs.get("table_id"), + ) + + +def open_dataset_from_open_kwargs(ds_uri: str | None, open_kwargs: dict[str, Any] | None) -> lance.LanceDataset: + """Re-open a dataset on a worker from serialized ``_lance_open_kwargs``. - merged_storage_options = merge_storage_options(storage_options, namespace_storage_options) + The namespace client is not picklable, so only the (impl, properties, table_id) + triple travels with the task; the client is re-created here via the lru cache. + """ + open_kwargs = dict(open_kwargs or {}) + namespace_impl, namespace_properties, table_id = pop_namespace_params(open_kwargs) return lance.dataset( - None if has_namespace_params(namespace_impl, table_id) else resolved_uri, - storage_options=merged_storage_options, + None if has_namespace_params(namespace_impl, table_id) else ds_uri, **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), - **kwargs, + **open_kwargs, ) diff --git a/tests/io/lancedb/conftest.py b/tests/io/lancedb/conftest.py index dda9665..464714d 100644 --- a/tests/io/lancedb/conftest.py +++ b/tests/io/lancedb/conftest.py @@ -1,6 +1,33 @@ from __future__ import annotations +import os +import sys + import pytest # Try to import lance; if it fails, all tests in this directory will be skipped. lance = pytest.importorskip("lance") + + +_exitstatus = 0 + + +def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: + global _exitstatus + _exitstatus = exitstatus + + +@pytest.hookimpl(trylast=True) +def pytest_unconfigure(config: pytest.Config) -> None: + """Hard-exit after reporting to dodge an upstream teardown crash. + + Running enough ``daft`` + ``lance`` write/read cycles in one process makes the + interpreter SIGSEGV during native finalization (in lance's runtime, inside + OpenSSL cert teardown) *after* every test has already run and been reported. + This reproduces with plain ``daft.DataFrame.write_lance`` and is unrelated to + test outcomes. Runs last, once tracebacks and the summary are already flushed, + so it never hides a failure — it only skips the crashing native finalization. + """ + sys.stdout.flush() + sys.stderr.flush() + os._exit(_exitstatus) diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index c573a56..46ad10c 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -6,64 +6,177 @@ import daft_lance +def _dir_ns(tmp_path): + return {"namespace_impl": "dir", "namespace_properties": {"root": str(tmp_path)}} + + def test_namespace_write_read_append_roundtrip(tmp_path): - namespace_properties = {"root": str(tmp_path)} + ns = _dir_ns(tmp_path) table_id = ["roundtrip"] df1 = daft.from_pydict({"id": [1, 2], "label": ["a", "b"]}) df2 = daft.from_pydict({"id": [3], "label": ["c"]}) - df1.write_lance( - namespace_impl="dir", - namespace_properties=namespace_properties, + daft_lance.write_lance(df1, table_id=table_id, mode="create", **ns).collect() + daft_lance.write_lance(df2, table_id=table_id, mode="append", **ns).collect() + + result = daft_lance.read_lance(table_id=table_id, **ns).to_pydict() + + assert result == {"id": [1, 2, 3], "label": ["a", "b", "c"]} + + +def test_namespace_overwrite(tmp_path): + ns = _dir_ns(tmp_path) + table_id = ["overwrite_tbl"] + + daft_lance.write_lance(daft.from_pydict({"id": [1, 2]}), table_id=table_id, mode="create", **ns).collect() + daft_lance.write_lance(daft.from_pydict({"id": [7, 8, 9]}), table_id=table_id, mode="overwrite", **ns).collect() + + result = daft_lance.read_lance(table_id=table_id, **ns).to_pydict() + assert result == {"id": [7, 8, 9]} + + +def test_namespace_overwrite_missing_table_declares(tmp_path): + ns = _dir_ns(tmp_path) + table_id = ["overwrite_fresh"] + + daft_lance.write_lance(daft.from_pydict({"id": [1]}), table_id=table_id, mode="overwrite", **ns).collect() + + result = daft_lance.read_lance(table_id=table_id, **ns).to_pydict() + assert result == {"id": [1]} + + +def test_namespace_read_supports_pushdowns(tmp_path): + ns = _dir_ns(tmp_path) + table_id = ["pushdowns"] + + daft_lance.write_lance( + daft.from_pydict( + { + "id": [1, 2, 3], + "label": ["a", "b", "c"], + "score": [10, 20, 30], + } + ), table_id=table_id, mode="create", + **ns, ).collect() - df2.write_lance( - namespace_impl="dir", - namespace_properties=namespace_properties, + + result = daft_lance.read_lance(table_id=table_id, **ns).where(daft.col("score") > 10).select("label").to_pydict() + + assert result == {"label": ["b", "c"]} + + +def test_namespace_count_pushdown(tmp_path): + ns = _dir_ns(tmp_path) + table_id = ["count_tbl"] + + daft_lance.write_lance(daft.from_pydict({"id": list(range(10))}), table_id=table_id, mode="create", **ns).collect() + + assert daft_lance.read_lance(table_id=table_id, **ns).count_rows() == 10 + + +def test_namespace_patch_daft(tmp_path): + ns = _dir_ns(tmp_path) + table_id = ["patched"] + + daft_lance.patch_daft() + daft_lance.patch_daft() # idempotent + + daft.from_pydict({"id": [1, 2]}).write_lance(table_id=table_id, mode="create", **ns).collect() + result = daft.read_lance(table_id=table_id, **ns).to_pydict() + assert result == {"id": [1, 2]} + + # Plain-URI writes keep working through the patched method. + uri = str(tmp_path / "plain.lance") + daft.from_pydict({"id": [5]}).write_lance(uri).collect() + assert daft.read_lance(uri).to_pydict() == {"id": [5]} + + +def test_namespace_write_lance_merge_new_columns(tmp_path): + ns = _dir_ns(tmp_path) + table_id = ["merge_tbl"] + + daft_lance.write_lance( + daft.from_pydict({"id": [1, 2, 3], "score": [10, 20, 30]}), table_id=table_id, - mode="append", + mode="create", + **ns, ).collect() - result = daft.read_lance( - namespace_impl="dir", - namespace_properties=namespace_properties, + df = daft_lance.read_lance( table_id=table_id, - ).to_pydict() + default_scan_options={"with_row_address": True}, + include_fragment_id=True, + **ns, + ) + df = df.with_column("doubled", df["score"] * 2) + daft_lance.write_lance(df, table_id=table_id, mode="merge", **ns).collect() - assert result == {"id": [1, 2, 3], "label": ["a", "b", "c"]} + result = daft_lance.read_lance(table_id=table_id, **ns).sort("id").to_pydict() + assert result["doubled"] == [20, 40, 60] -def test_namespace_read_supports_pushdowns(tmp_path): - namespace_properties = {"root": str(tmp_path)} - table_id = ["pushdowns"] +def test_namespace_merge_columns_df(tmp_path): + ns = _dir_ns(tmp_path) + table_id = ["merge_cols_df"] - daft.from_pydict( - { - "id": [1, 2, 3], - "label": ["a", "b", "c"], - "score": [10, 20, 30], - } - ).write_lance( - namespace_impl="dir", - namespace_properties=namespace_properties, + daft_lance.write_lance( + daft.from_pydict({"id": [1, 2, 3], "score": [1, 2, 3]}), table_id=table_id, mode="create", + **ns, ).collect() - result = ( - daft.read_lance( - namespace_impl="dir", - namespace_properties=namespace_properties, - table_id=table_id, - ) - .where(daft.col("score") > 10) - .select("label") - .to_pydict() + df = daft_lance.read_lance( + table_id=table_id, + default_scan_options={"with_row_address": True}, + include_fragment_id=True, + **ns, ) + df = df.with_column("tripled", df["score"] * 3) + daft_lance.merge_columns_df(df.select("fragment_id", "_rowaddr", "tripled"), table_id=table_id, **ns) - assert result == {"label": ["b", "c"]} + result = daft_lance.read_lance(table_id=table_id, **ns).sort("id").to_pydict() + assert result["tripled"] == [3, 6, 9] + + +def test_namespace_create_scalar_index(tmp_path): + ns = _dir_ns(tmp_path) + table_id = ["indexed"] + + daft_lance.write_lance( + daft.from_pydict({"id": list(range(100)), "price": [i * 2 for i in range(100)]}), + table_id=table_id, + mode="create", + **ns, + ).collect() + + daft_lance.create_scalar_index(table_id=table_id, column="price", index_type="BTREE", **ns) + + import lance + import lance_namespace as ln + from lance_namespace import DescribeTableRequest + + namespace = ln.connect("dir", {"root": str(tmp_path)}) + location = namespace.describe_table(DescribeTableRequest(id=table_id)).location + indices = lance.dataset(location).list_indices() + assert any(idx["fields"] == ["price"] for idx in indices) + + +def test_namespace_compact_files(tmp_path): + ns = _dir_ns(tmp_path) + table_id = ["compacted"] + + daft_lance.write_lance(daft.from_pydict({"id": [1]}), table_id=table_id, mode="create", **ns).collect() + for i in range(3): + daft_lance.write_lance(daft.from_pydict({"id": [i + 2]}), table_id=table_id, mode="append", **ns).collect() + + daft_lance.compact_files(table_id=table_id, compaction_options={"target_rows_per_fragment": 1024}, **ns) + + result = daft_lance.read_lance(table_id=table_id, **ns).sort("id").to_pydict() + assert result == {"id": [1, 2, 3, 4]} def test_namespace_rejects_uri_and_namespace(tmp_path): @@ -82,3 +195,13 @@ def test_namespace_requires_table_id(tmp_path): namespace_impl="dir", namespace_properties={"root": str(tmp_path)}, ) + + +def test_namespace_requires_impl(): + with pytest.raises(ValueError, match="'namespace_impl' must be provided"): + daft_lance.read_lance(table_id=["tbl"]) + + +def test_namespace_requires_uri_or_namespace(): + with pytest.raises(ValueError, match="Must provide either 'uri' OR"): + daft_lance.read_lance() diff --git a/tests/io/lancedb/test_namespace_rest_e2e.py b/tests/io/lancedb/test_namespace_rest_e2e.py index d9f27e9..a193e2e 100644 --- a/tests/io/lancedb/test_namespace_rest_e2e.py +++ b/tests/io/lancedb/test_namespace_rest_e2e.py @@ -6,7 +6,7 @@ import pytest import daft -import daft_lance # noqa: F401 - patches daft.read_lance and DataFrame.write_lance. +import daft_lance pytestmark = pytest.mark.skipif( os.environ.get("DAFT_LANCE_REST_URI") is None, @@ -23,6 +23,7 @@ def test_rest_namespace_write_read_append_roundtrip() -> None: catalog = os.environ.get("DAFT_LANCE_REST_CATALOG", "lance_catalog") schema = os.environ.get("DAFT_LANCE_REST_SCHEMA", "daft_ns_e2e") table_id = [catalog, schema, f"orders_{uuid.uuid4().hex[:8]}"] + ns = {"namespace_impl": "rest", "namespace_properties": namespace_properties} namespace = ln.connect("rest", namespace_properties) try: @@ -30,41 +31,25 @@ def test_rest_namespace_write_read_append_roundtrip() -> None: except Exception: namespace.create_namespace(CreateNamespaceRequest(id=[catalog, schema], mode="CREATE")) - daft.from_pydict( - { - "id": [1, 2, 3], - "label": ["a", "b", "c"], - "score": [10, 20, 30], - } - ).write_lance( - namespace_impl="rest", - namespace_properties=namespace_properties, + daft_lance.write_lance( + daft.from_pydict({"id": [1, 2, 3], "label": ["a", "b", "c"], "score": [10, 20, 30]}), table_id=table_id, mode="create", + **ns, ).collect() - daft.from_pydict( - { - "id": [4, 5], - "label": ["d", "e"], - "score": [40, 50], - } - ).write_lance( - namespace_impl="rest", - namespace_properties=namespace_properties, + daft_lance.write_lance( + daft.from_pydict({"id": [4, 5], "label": ["d", "e"], "score": [40, 50]}), table_id=table_id, mode="append", + **ns, ).collect() describe = namespace.describe_table(DescribeTableRequest(id=table_id)) location = getattr(describe, "location", None) or getattr(describe, "table_uri", None) assert location - result = daft.read_lance( - namespace_impl="rest", - namespace_properties=namespace_properties, - table_id=table_id, - ).to_pydict() + result = daft_lance.read_lance(table_id=table_id, **ns).to_pydict() assert result == { "id": [1, 2, 3, 4, 5], "label": ["a", "b", "c", "d", "e"], @@ -72,23 +57,23 @@ def test_rest_namespace_write_read_append_roundtrip() -> None: } filtered = ( - daft.read_lance( - namespace_impl="rest", - namespace_properties=namespace_properties, - table_id=table_id, - ) - .where(daft.col("score") >= 30) - .select("id", "label") - .to_pydict() + daft_lance.read_lance(table_id=table_id, **ns).where(daft.col("score") >= 30).select("id", "label").to_pydict() ) assert filtered == {"id": [3, 4, 5], "label": ["c", "d", "e"]} - assert ( - daft.read_lance( - namespace_impl="rest", - namespace_properties=namespace_properties, - table_id=table_id, - ).count_rows() - == 5 - ) + assert daft_lance.read_lance(table_id=table_id, **ns).count_rows() == 5 assert lance.dataset(None, namespace_client=namespace, table_id=table_id).count_rows() == 5 + + +def test_rest_namespace_patch_daft_roundtrip() -> None: + """The opt-in ``patch_daft()`` path: ``daft.read_lance`` / ``df.write_lance``.""" + namespace_properties = {"uri": os.environ["DAFT_LANCE_REST_URI"]} + catalog = os.environ.get("DAFT_LANCE_REST_CATALOG", "lance_catalog") + schema = os.environ.get("DAFT_LANCE_REST_SCHEMA", "daft_ns_e2e") + table_id = [catalog, schema, f"patched_{uuid.uuid4().hex[:8]}"] + ns = {"namespace_impl": "rest", "namespace_properties": namespace_properties} + + daft_lance.patch_daft() + + daft.from_pydict({"id": [1, 2]}).write_lance(table_id=table_id, mode="create", **ns).collect() + assert daft.read_lance(table_id=table_id, **ns).to_pydict() == {"id": [1, 2]} From 7b1a7c44f889e53611354de3ad56e024fb64aa65 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Tue, 7 Jul 2026 10:24:19 +0900 Subject: [PATCH 03/25] fix(namespace): tighten missing table handling --- daft_lance/namespace.py | 39 ++++++++++++++++++++++++++++-- tests/io/lancedb/conftest.py | 28 ++++++++++++--------- tests/io/lancedb/test_namespace.py | 37 ++++++++++++++++++++++++++++ 3 files changed, 90 insertions(+), 14 deletions(-) diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index 3fd19c7..3130008 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -103,7 +103,8 @@ def is_table_not_found(error: Exception) -> bool: """Whether an exception from ``describe_table`` means the table does not exist. Native implementations raise the typed ``TableNotFoundError``; the Rust-backed - REST client surfaces untyped errors, so fall back to message matching. + REST client may surface untyped errors, so only fall back when the error is + clearly a 404 for a table resource. """ try: from lance_namespace.errors import TableNotFoundError @@ -112,8 +113,42 @@ def is_table_not_found(error: Exception) -> bool: return True except ImportError: pass + + status_code = _error_status_code(error) + if status_code != 404: + return False + message = str(error).lower() - return "not found" in message or "does not exist" in message or "no such table" in message + return "no such table" in message or ( + "table" in message and ("not found" in message or "does not exist" in message) + ) + + +def _error_status_code(error: Exception) -> int | None: + """Best-effort HTTP status extraction for untyped namespace REST errors.""" + for attr in ("status", "status_code", "code"): + value = getattr(error, attr, None) + status_code = _coerce_status_code(value) + if status_code is not None: + return status_code + + response = getattr(error, "response", None) + if response is not None: + for attr in ("status", "status_code"): + value = getattr(response, attr, None) + status_code = _coerce_status_code(value) + if status_code is not None: + return status_code + + return None + + +def _coerce_status_code(value: Any) -> int | None: + if isinstance(value, int): + return value + if isinstance(value, str) and value.isdigit(): + return int(value) + return None def resolve_namespace_table( diff --git a/tests/io/lancedb/conftest.py b/tests/io/lancedb/conftest.py index 464714d..06e1b1d 100644 --- a/tests/io/lancedb/conftest.py +++ b/tests/io/lancedb/conftest.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import os import sys @@ -8,8 +6,21 @@ # Try to import lance; if it fails, all tests in this directory will be skipped. lance = pytest.importorskip("lance") +_NATIVE_TEARDOWN_CRASH_MARKER = "lance_native_teardown_crash_workaround" +_exitstatus: int | None = None +_hard_exit_after_session = False + + +def pytest_configure(config: pytest.Config) -> None: + config.addinivalue_line( + "markers", + f"{_NATIVE_TEARDOWN_CRASH_MARKER}: hard-exit after reported results to avoid a known Lance native teardown crash", + ) + -_exitstatus = 0 +def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: + global _hard_exit_after_session + _hard_exit_after_session = any(item.get_closest_marker(_NATIVE_TEARDOWN_CRASH_MARKER) for item in items) def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: @@ -19,15 +30,8 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: @pytest.hookimpl(trylast=True) def pytest_unconfigure(config: pytest.Config) -> None: - """Hard-exit after reporting to dodge an upstream teardown crash. - - Running enough ``daft`` + ``lance`` write/read cycles in one process makes the - interpreter SIGSEGV during native finalization (in lance's runtime, inside - OpenSSL cert teardown) *after* every test has already run and been reported. - This reproduces with plain ``daft.DataFrame.write_lance`` and is unrelated to - test outcomes. Runs last, once tracebacks and the summary are already flushed, - so it never hides a failure — it only skips the crashing native finalization. - """ + if not _hard_exit_after_session or _exitstatus is None: + return sys.stdout.flush() sys.stderr.flush() os._exit(_exitstatus) diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index 46ad10c..943bed4 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -1,9 +1,14 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest import daft import daft_lance +import daft_lance.namespace as namespace_mod + +pytestmark = pytest.mark.lance_native_teardown_crash_workaround def _dir_ns(tmp_path): @@ -46,6 +51,38 @@ def test_namespace_overwrite_missing_table_declares(tmp_path): assert result == {"id": [1]} +def test_namespace_overwrite_does_not_declare_on_ambiguous_error(monkeypatch, tmp_path): + class FakeNamespace: + declared = False + + def describe_table(self, request): + raise RuntimeError("permission denied: parent catalog does not exist") + + def declare_table(self, request): + self.declared = True + return SimpleNamespace(location=str(tmp_path / "should_not_exist.lance")) + + namespace = FakeNamespace() + monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: namespace) + + with pytest.raises(RuntimeError, match="permission denied"): + namespace_mod.resolve_namespace_table( + namespace_impl="rest", + namespace_properties={"uri": "http://namespace.example"}, + table_id=["catalog", "schema", "table"], + mode="overwrite", + ) + + assert not namespace.declared + + +def test_namespace_untyped_404_table_not_found_is_missing(): + class RestTableNotFound(RuntimeError): + status = 404 + + assert namespace_mod.is_table_not_found(RestTableNotFound("table catalog.schema.table not found")) + + def test_namespace_read_supports_pushdowns(tmp_path): ns = _dir_ns(tmp_path) table_id = ["pushdowns"] From b2e28c58b43a832539d4d8673e9f4db2c76cf8a3 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 13:04:58 +0900 Subject: [PATCH 04/25] refactor(sink): resolve namespace tables at execution time via DataSink.start() Constructing a LanceDataSink no longer talks to the namespace or opens the dataset; all resolution moves to start(), which Daft runs once on the driver before the sink is serialized to workers. This keeps declare_table from firing during plan construction and from leaving orphan declared tables behind when local parameter validation fails. Create-mode resolution is now describe-first (check_declared=True): a real existing table fails fast with a clear error suggesting overwrite/append, a declared-only stub is reused as a placeholder (write_fragments switches to overwrite when the stub materialized a dataset), and a lost declare race is recovered by re-describing and classifying. resolve_namespace_table returns a ResolvedNamespaceTable carrying uri/storage_options/placeholder state so the sink does not have to re-describe. Claude-Session: https://claude.ai/code/session_01VyD31uRyBSEmKvPQhLtM5E --- daft_lance/_lance.py | 4 +- daft_lance/lance_data_sink.py | 97 +++++++++---- daft_lance/lance_scalar_index.py | 4 +- daft_lance/namespace.py | 134 ++++++++++++++++-- daft_lance/utils.py | 5 +- tests/io/lancedb/conftest.py | 2 + tests/io/lancedb/test_lance_batching.py | 4 + .../lancedb/test_lance_data_sink_internals.py | 8 +- .../test_lance_data_sink_storage_versions.py | 4 +- tests/io/lancedb/test_mem_wal_writes.py | 5 + tests/io/lancedb/test_namespace.py | 113 +++++++++++++++ 11 files changed, 327 insertions(+), 53 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index ab65a11..12bd836 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -33,9 +33,7 @@ def _effective_uri(lance_ds: LanceDataset, uri: str | pathlib.Path | None) -> st return str(uri) if uri is not None else str(lance_ds.uri) -def _effective_storage_options( - lance_ds: LanceDataset, storage_options: dict[str, Any] | None -) -> dict[str, Any] | None: +def _effective_storage_options(lance_ds: LanceDataset, storage_options: dict[str, Any] | None) -> dict[str, Any] | None: """Storage options actually used to open the dataset (includes namespace-vended ones).""" open_kwargs = getattr(lance_ds, "_lance_open_kwargs", None) or {} return open_kwargs.get("storage_options") or storage_options diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index 2ec69a7..a995e37 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -27,6 +27,7 @@ resolve_storage_version, ) from daft_lance.namespace import ( + ResolvedNamespaceTable, get_namespace_kwargs, get_write_fragments_kwargs, merge_storage_options, @@ -72,21 +73,12 @@ def __init__( raise TypeError(f"Expected URI to be str or pathlib.Path, got {type(uri)}") self._mode = mode + self._uri = uri self._namespace_impl = namespace_impl self._namespace_properties = namespace_properties self._table_id = table_id self._io_config = get_context().daft_planning_config.default_io_config if io_config is None else io_config - base_storage_options = ( - ( - storage_options - if storage_options is not None - else io_config_to_storage_options(self._io_config, str(uri)) - ) - if uri is not None - else storage_options - ) - self._table_uri, namespace_storage_options = self._resolve_table_uri(uri) - self._storage_options = merge_storage_options(base_storage_options, namespace_storage_options) + self._user_storage_options = storage_options self._init_lance_knobs( max_rows_per_file=max_rows_per_file, max_rows_per_group=max_rows_per_group, @@ -96,18 +88,49 @@ def __init__( ) self._pyarrow_schema = self._normalize_schema(schema) self._init_blob_policy(blob_columns) + self._requested_storage_version = self._blob.apply_blob_v2_default(data_storage_version) self._use_mem_wal = use_mem_wal self._compact_after_write = compact_after_write self._mem_wal_total_rows: int = 0 self._mem_wal_total_bytes: int = 0 + # Resolved by start() on the driver, before the sink is serialized to workers. + # Construction must stay side-effect free: no namespace calls, no dataset opens. + self._table_uri: str | None = None + self._storage_options: dict[str, str] | None = None + self._data_storage_version: LanceStorageVersion | None = None + self._effective_pyarrow_schema: pa.Schema | None = None self._version: int = 0 self._table_schema: pa.Schema | None = None + self._declared_placeholder = False + self._overwrite_placeholder_dataset = False + + self._schema = Schema._from_field_name_and_types( + [ + ("num_fragments", DataType.int64()), + ("num_deleted_rows", DataType.int64()), + ("num_small_files", DataType.int64()), + ("version", DataType.int64()), + ] + ) + + def start(self) -> None: + """Resolve the target table and validate the requested mode against it. + + Runs once on the driver before the sink is serialized to workers, so this + is the single place where namespace side effects (``declare_table``) and + dataset opens happen. + """ + resolved = self._resolve_table() + self._table_uri = resolved.uri + self._declared_placeholder = resolved.is_declared_placeholder + self._storage_options = merge_storage_options(self._base_storage_options(), resolved.storage_options) + existing = self._absorb_existing_dataset() existing_version = getattr(existing, "data_storage_version", None) if existing is not None else None self._data_storage_version = resolve_storage_version( - self._blob.apply_blob_v2_default(data_storage_version), + self._requested_storage_version, existing_version, self._mode, ) @@ -119,14 +142,6 @@ def __init__( # Schema actually written to the dataset (blob columns retyped to lance.blob.v2). self._effective_pyarrow_schema = self._blob.build_effective_schema(self._pyarrow_schema) - self._schema = Schema._from_field_name_and_types( - [ - ("num_fragments", DataType.int64()), - ("num_deleted_rows", DataType.int64()), - ("num_small_files", DataType.int64()), - ("version", DataType.int64()), - ] - ) @property def _namespace_kwargs(self) -> dict[str, object]: @@ -136,19 +151,26 @@ def _namespace_kwargs(self) -> dict[str, object]: def _dataset_uri_arg(self) -> str | None: return None if self._namespace_impl is not None and self._table_id is not None else self._table_uri - def _resolve_table_uri(self, uri: str | pathlib.Path | None) -> tuple[str, dict[str, str] | None]: - if uri is not None: - return str(uri), None - mode = "create" if self._mode == "create" else "overwrite" if self._mode == "overwrite" else "read" - resolved_uri, namespace_storage_options = resolve_namespace_table( + def _resolve_table(self) -> ResolvedNamespaceTable: + if self._uri is not None: + return ResolvedNamespaceTable(uri=str(self._uri)) + mode = self._mode if self._mode in ("create", "overwrite") else "read" + resolved = resolve_namespace_table( namespace_impl=self._namespace_impl, namespace_properties=self._namespace_properties, table_id=self._table_id, mode=mode, ) - if resolved_uri is None: + if resolved is None: raise ValueError("Unable to resolve Lance dataset URI from namespace.") - return resolved_uri, namespace_storage_options + return resolved + + def _base_storage_options(self) -> dict[str, str] | None: + if self._uri is None: + return self._user_storage_options + if self._user_storage_options is not None: + return self._user_storage_options + return io_config_to_storage_options(self._io_config, str(self._uri)) @staticmethod def _reject_unsupported_modes( @@ -222,7 +244,7 @@ def _absorb_existing_dataset(self) -> lance.LanceDataset | None: if dataset is None: if self._mode == "append": raise ValueError("Cannot append to non-existent Lance dataset.") - if self._mode == "create" and self._storage_options is None: + if self._mode == "create" and self._storage_options is None and self._table_uri is not None: p = pathlib.Path(self._table_uri) if p.is_file(): raise FileExistsError("Target path points to a file, cannot create a dataset here.") @@ -232,7 +254,15 @@ def _absorb_existing_dataset(self) -> lance.LanceDataset | None: self._version = dataset.latest_version if self._mode == "create": - raise ValueError("Cannot create a Lance dataset at a location where one already exists.") + if not self._declared_placeholder: + raise ValueError( + "Cannot create a Lance dataset at a location where one already exists. " + 'Use mode="overwrite" to replace it or mode="append" to add to it.' + ) + # The namespace declared a stub table and this location already holds + # a (placeholder) dataset; overwrite it instead of failing the create. + self._overwrite_placeholder_dataset = True + self._table_schema = None if self._mode == "append" and not _pyarrow_schema_castable( blob_aware_schema_for_validation(self._pyarrow_schema, self._table_schema), @@ -261,11 +291,13 @@ def _prepare_arrow_table(self, input_table: pa.Table) -> pa.Table: return pa.Table.from_batches(input_table.to_batches(), target_schema) def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetadata]]: + assert self._table_uri is not None, "LanceDataSink.start() must run before writes" wrapped = self._blob.wrap_table(table) + write_mode = "overwrite" if self._overwrite_placeholder_dataset else self._mode fragments = lance.fragment.write_fragments( wrapped, dataset_uri=self._table_uri, - mode=self._mode, + mode=write_mode, storage_options=self._storage_options, max_rows_per_file=self._max_rows_per_file, max_rows_per_group=self._max_rows_per_group, @@ -283,6 +315,7 @@ def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetada return WriteResult(result=fragments, bytes_written=bytes_written, rows_written=wrapped.num_rows) def _ensure_mem_wal_dataset(self) -> lance.LanceDataset: + assert self._effective_pyarrow_schema is not None, "LanceDataSink.start() must run before writes" try: ds = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) except (ValueError, FileNotFoundError, OSError): @@ -369,12 +402,16 @@ def finalize(self, write_results: list[WriteResult[list[FragmentMetadata]]]) -> def _finalize_cow(self, write_results: list[WriteResult[list[FragmentMetadata]]]) -> MicroPartition: fragments = list(chain.from_iterable(write_result.result for write_result in write_results)) + assert self._effective_pyarrow_schema is not None, "LanceDataSink.start() must run before finalize" operation: lance.LanceOperation.BaseOperation if self._mode == "create" or self._mode == "overwrite": operation = lance.LanceOperation.Overwrite(self._effective_pyarrow_schema, fragments) elif self._mode == "append": operation = lance.LanceOperation.Append(fragments) + assert self._table_uri is not None, "LanceDataSink.start() must run before finalize" + # Unlike lance.dataset(), commit() requires a str base_uri even when a + # namespace client is passed; the namespace kwargs only register the commit. dataset = lance.LanceDataset.commit( self._table_uri, operation, diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index b967ce6..ef591cc 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -465,9 +465,7 @@ def _create_partitioned_index( logger.info("Starting index metadata merging by reloading dataset to get latest state") namespace_kwargs = namespace_kwargs_for_dataset(lance_ds) - lance_ds = lance.dataset( - None if namespace_kwargs else uri, storage_options=storage_options, **namespace_kwargs - ) + lance_ds = lance.dataset(None if namespace_kwargs else uri, storage_options=storage_options, **namespace_kwargs) lance_ds.merge_index_metadata(index_id, index_type) logger.info("Starting atomic index creation and commit") diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index 3130008..cb58113 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +from dataclasses import dataclass from functools import lru_cache from typing import Any from urllib.parse import urlparse @@ -92,11 +93,59 @@ def _response_location(response: Any) -> str: return _normalize_file_uri(str(location)) -def _declare_table(namespace: Any, table_id: list[str]) -> tuple[str, dict[str, str] | None]: +@dataclass(frozen=True) +class ResolvedNamespaceTable: + """A namespace table resolved to a physical location. + + ``is_declared_placeholder`` is True when the location comes from a + ``declare_table`` call (fresh or pre-existing declared-only stub): some + namespace implementations materialize a stub dataset at declare time, and a + ``mode="create"`` write must overwrite that stub instead of failing. + """ + + uri: str + storage_options: dict[str, str] | None = None + is_declared_placeholder: bool = False + + +def _resolved_from_response(response: Any, *, is_declared_placeholder: bool = False) -> ResolvedNamespaceTable: + return ResolvedNamespaceTable( + uri=_response_location(response), + storage_options=_storage_options(response), + is_declared_placeholder=is_declared_placeholder, + ) + + +def _describe_table(namespace: Any, table_id: list[str], *, check_declared: bool = False) -> Any: + from lance_namespace import DescribeTableRequest + + request = ( + DescribeTableRequest(id=table_id, check_declared=True) if check_declared else DescribeTableRequest(id=table_id) + ) + return namespace.describe_table(request) + + +def _declare_table(namespace: Any, table_id: list[str]) -> ResolvedNamespaceTable: from lance_namespace import DeclareTableRequest response = namespace.declare_table(DeclareTableRequest(id=table_id, location=None)) - return _response_location(response), _storage_options(response) + return _resolved_from_response(response, is_declared_placeholder=True) + + +def is_declared_placeholder_response(response: Any) -> bool: + """Whether a ``describe_table`` response describes a declared-only stub table. + + ``is_only_declared`` is the standard signal (only reliable when the request + was made with ``check_declared=True``); the ``lance.declared`` property is a + compatibility fallback for services that expose the state that way. + """ + if getattr(response, "is_only_declared", None) is True: + return True + for attr in ("properties", "metadata"): + mapping = getattr(response, attr, None) + if mapping and str(dict(mapping).get("lance.declared", "")).strip().lower() == "true": + return True + return False def is_table_not_found(error: Exception) -> bool: @@ -124,6 +173,28 @@ def is_table_not_found(error: Exception) -> bool: ) +def is_table_already_exists(error: Exception) -> bool: + """Whether an exception from ``declare_table`` means the table already exists. + + Mirrors :func:`is_table_not_found`: prefer the typed error, fall back to a + clearly-classified conflict from untyped REST errors. + """ + try: + from lance_namespace.errors import TableAlreadyExistsError + + if isinstance(error, TableAlreadyExistsError): + return True + except ImportError: + pass + + status_code = _error_status_code(error) + if status_code != 409: + return False + + message = str(error).lower() + return "table" in message and "already exists" in message + + def _error_status_code(error: Exception) -> int | None: """Best-effort HTTP status extraction for untyped namespace REST errors.""" for attr in ("status", "status_code", "code"): @@ -157,30 +228,67 @@ def resolve_namespace_table( namespace_properties: dict[str, str] | None, table_id: list[str] | None, mode: str = "read", -) -> tuple[str | None, dict[str, str] | None]: - """Resolve a namespace table to ``(location, storage_options)``. +) -> ResolvedNamespaceTable | None: + """Resolve a namespace table to a :class:`ResolvedNamespaceTable`. - ``mode="create"`` declares the table; ``mode="overwrite"`` declares it only when - it does not exist yet; other modes require the table to exist. + ``mode="create"`` requires the table to not exist (declared-only stubs are + reused as placeholders); ``mode="overwrite"`` declares the table when it does + not exist yet; other modes require the table to exist. Only ``mode="create"`` + and ``mode="overwrite"`` have side effects (a ``declare_table`` call). """ namespace = get_or_create_namespace(namespace_impl, namespace_properties) if namespace is None or table_id is None: - return None, None + return None if mode == "create": - return _declare_table(namespace, table_id) - - from lance_namespace import DescribeTableRequest + return _resolve_for_create(namespace, table_id) try: - response = namespace.describe_table(DescribeTableRequest(id=table_id)) - return _response_location(response), _storage_options(response) + return _resolved_from_response(_describe_table(namespace, table_id)) except Exception as error: if mode == "overwrite" and is_table_not_found(error): - return _declare_table(namespace, table_id) + return _declare_or_recover(namespace, table_id, check_declared=False) raise +def _resolve_for_create(namespace: Any, table_id: list[str]) -> ResolvedNamespaceTable: + """Describe-first create resolution. + + ``declare_table`` fails with a conflict on any existing table (declared-only + or real), so classify via ``describe_table`` before declaring, and again if + a concurrent declare wins the race in between. + """ + try: + response = _describe_table(namespace, table_id, check_declared=True) + except Exception as error: + if not is_table_not_found(error): + raise + return _declare_or_recover(namespace, table_id, check_declared=True) + return _classify_existing_for_create(response, table_id) + + +def _declare_or_recover(namespace: Any, table_id: list[str], *, check_declared: bool) -> ResolvedNamespaceTable: + try: + return _declare_table(namespace, table_id) + except Exception as error: + if not is_table_already_exists(error): + raise + # Lost a declare race; re-describe and classify what won. + response = _describe_table(namespace, table_id, check_declared=check_declared) + if not check_declared: + return _resolved_from_response(response) + return _classify_existing_for_create(response, table_id) + + +def _classify_existing_for_create(response: Any, table_id: list[str]) -> ResolvedNamespaceTable: + if is_declared_placeholder_response(response): + return _resolved_from_response(response, is_declared_placeholder=True) + raise ValueError( + f"Table {table_id!r} already exists in the namespace and cannot be created. " + 'Use mode="overwrite" to replace it or mode="append" to add to it.' + ) + + def merge_storage_options( storage_options: dict[str, Any] | None, namespace_storage_options: dict[str, str] | None, diff --git a/daft_lance/utils.py b/daft_lance/utils.py index 3262f3a..d072b23 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -102,12 +102,15 @@ def construct_lance_dataset( resolved_uri = str(uri) if uri is not None else None namespace_storage_options = None if resolved_uri is None: - resolved_uri, namespace_storage_options = resolve_namespace_table( + resolved = resolve_namespace_table( namespace_impl=namespace_impl, namespace_properties=namespace_properties, table_id=table_id, mode="read", ) + if resolved is not None: + resolved_uri = resolved.uri + namespace_storage_options = resolved.storage_options if resolved_uri is None: raise ValueError("Unable to resolve Lance dataset URI.") merged_storage_options = merge_storage_options(storage_options, namespace_storage_options) diff --git a/tests/io/lancedb/conftest.py b/tests/io/lancedb/conftest.py index 06e1b1d..69feee1 100644 --- a/tests/io/lancedb/conftest.py +++ b/tests/io/lancedb/conftest.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import os import sys diff --git a/tests/io/lancedb/test_lance_batching.py b/tests/io/lancedb/test_lance_batching.py index ea27359..c805115 100644 --- a/tests/io/lancedb/test_lance_batching.py +++ b/tests/io/lancedb/test_lance_batching.py @@ -55,6 +55,7 @@ def test_accumulate_small_micropartitions(schema, tmp_path): with patch("daft_lance.lance_data_sink.lance", fake): sink = LanceDataSink(uri=str(tmp_path / "tbl"), schema=schema, mode="create", max_rows_per_file=25) + sink.start() mps = [_make_mp(10), _make_mp(20), _make_mp(30)] results = list(sink.write(iter(mps))) @@ -69,6 +70,7 @@ def test_flush_remaining_at_end(schema, tmp_path): with patch("daft_lance.lance_data_sink.lance", fake): sink = LanceDataSink(uri=str(tmp_path / "tbl"), schema=schema, mode="create", max_rows_per_file=25) + sink.start() mps = [_make_mp(10), _make_mp(5)] results = list(sink.write(iter(mps))) @@ -83,6 +85,7 @@ def test_large_micropartition_writes_directly(schema, tmp_path): with patch("daft_lance.lance_data_sink.lance", fake): sink = LanceDataSink(uri=str(tmp_path / "tbl"), schema=schema, mode="create", max_rows_per_file=25) + sink.start() mps = [_make_mp(30)] results = list(sink.write(iter(mps))) @@ -102,6 +105,7 @@ def test_accumulation_default_batches_across_micropartitions(schema, tmp_path): with patch("daft_lance.lance_data_sink.lance", fake): sink = LanceDataSink(uri=str(tmp_path / "tbl"), schema=schema, mode="create") + sink.start() mps = [_make_mp(10), _make_mp(10), _make_mp(10)] results = list(sink.write(iter(mps))) diff --git a/tests/io/lancedb/test_lance_data_sink_internals.py b/tests/io/lancedb/test_lance_data_sink_internals.py index 47f3107..e6983c6 100644 --- a/tests/io/lancedb/test_lance_data_sink_internals.py +++ b/tests/io/lancedb/test_lance_data_sink_internals.py @@ -17,14 +17,16 @@ def test_load_existing_dataset_missing_raises_for_append(tmp_path): schema = pa.schema([("a", pa.int64())]) + sink = LanceDataSink(uri=str(tmp_path / f"missing-{uuid.uuid4()}"), schema=schema, mode="append") with pytest.raises(ValueError, match="Cannot append to non-existent Lance dataset"): - LanceDataSink(uri=str(tmp_path / f"missing-{uuid.uuid4()}"), schema=schema, mode="append") + sink.start() def test_load_existing_dataset_missing_returns_none_for_create(tmp_path): schema = pa.schema([("a", pa.int64())]) sink = LanceDataSink(uri=str(tmp_path / f"missing-{uuid.uuid4()}"), schema=schema, mode="create") - # Construction succeeds; the dataset will be created on the first write. + sink.start() + # Start succeeds; the dataset will be created on the first write. assert sink is not None @@ -143,6 +145,7 @@ def __getattr__(self, name): with patch("daft_lance.lance_data_sink.lance", fake): sink = LanceDataSink(uri=str(tmp_path / "tbl"), schema=schema, mode="create") + sink.start() write_results = list(sink.write(iter(mps))) # Default ``max_rows_per_file`` is 1_048_576. 20 partitions x 100K = 2M rows @@ -161,6 +164,7 @@ def test_prepare_arrow_table_missing_column_rejected(tmp_path): lance.write_dataset(initial, uri) # Try to append with only one column sink = LanceDataSink(uri=uri, schema=schema, mode="append") + sink.start() missing_col = pa.table({"a": pa.array([2], type=pa.int64())}) with pytest.raises(Exception): # any error is fine; the redesign currently relies on the cast to fail sink._prepare_arrow_table(missing_col) diff --git a/tests/io/lancedb/test_lance_data_sink_storage_versions.py b/tests/io/lancedb/test_lance_data_sink_storage_versions.py index 28496bc..fcf5f91 100644 --- a/tests/io/lancedb/test_lance_data_sink_storage_versions.py +++ b/tests/io/lancedb/test_lance_data_sink_storage_versions.py @@ -47,6 +47,7 @@ def test_append_inherits_storage_version(tmp_path): # inherit from the existing dataset. schema = pa.schema([("a", pa.int64())]) sink = LanceDataSink(uri=uri, schema=schema, mode="append") + sink.start() assert sink._data_storage_version == initial_version df2 = _make_df(5) @@ -64,8 +65,9 @@ def test_storage_version_mismatch_on_append_rejected(tmp_path): df1.write_lance(uri, mode="create", data_storage_version="2.1") schema = pa.schema([("a", pa.int64())]) + sink = LanceDataSink(uri=uri, schema=schema, mode="append", data_storage_version="2.2") with pytest.raises(ValueError, match="does not match existing dataset version"): - LanceDataSink(uri=uri, schema=schema, mode="append", data_storage_version="2.2") + sink.start() def test_use_legacy_format_emits_deprecation_warning(tmp_path): diff --git a/tests/io/lancedb/test_mem_wal_writes.py b/tests/io/lancedb/test_mem_wal_writes.py index dd51b80..5ac5715 100644 --- a/tests/io/lancedb/test_mem_wal_writes.py +++ b/tests/io/lancedb/test_mem_wal_writes.py @@ -205,6 +205,7 @@ def test_ensure_creates_dataset_if_missing(self, lance_dataset_path): target = os.path.join(lance_dataset_path, "new_ds") schema = pa.schema([("a", pa.int64())]) sink = LanceDataSink(uri=target, schema=schema, mode="create", use_mem_wal=True) + sink.start() ds = sink._ensure_mem_wal_dataset() assert ds is not None assert ds.mem_wal_index_details() is not None @@ -214,6 +215,7 @@ def test_ensure_reuses_existing_dataset(self, lance_dataset_path): lance.write_dataset(pa.table({"a": [1]}, schema=schema), lance_dataset_path) sink = LanceDataSink(uri=lance_dataset_path, schema=schema, mode="append", use_mem_wal=True) + sink.start() ds = sink._ensure_mem_wal_dataset() assert ds is not None assert ds.mem_wal_index_details() is not None @@ -226,6 +228,7 @@ def test_ensure_idempotent_initialization(self, lance_dataset_path): ds.initialize_mem_wal(unsharded=True) sink = LanceDataSink(uri=lance_dataset_path, schema=schema, mode="append", use_mem_wal=True) + sink.start() ds2 = sink._ensure_mem_wal_dataset() assert ds2.mem_wal_index_details() is not None @@ -236,6 +239,7 @@ def test_write_result_has_empty_fragments(self, lance_dataset_path): schema = pa.schema([("a", pa.int64())]) sink = LanceDataSink(uri=lance_dataset_path, schema=schema, mode="create", use_mem_wal=True) + sink.start() mp = MicroPartition.from_pydict({"a": [1, 2, 3]}) results = list(sink.write(iter([mp]))) assert len(results) == 1 @@ -254,6 +258,7 @@ def test_finalize_mem_wal_returns_stats(self, lance_dataset_path): use_mem_wal=True, compact_after_write=False, ) + sink.start() mp = MicroPartition.from_pydict({"a": [1, 2, 3]}) results = list(sink.write(iter([mp]))) stats_mp = sink.finalize(results) diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index 943bed4..e5a954b 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -216,6 +216,119 @@ def test_namespace_compact_files(tmp_path): assert result == {"id": [1, 2, 3, 4]} +def test_sink_construction_is_side_effect_free(tmp_path): + ns = _dir_ns(tmp_path) + schema = daft.from_pydict({"id": [1]}).schema() + + sink = daft_lance.LanceDataSink(None, schema, "create", table_id=["deferred"], **ns) + assert not (tmp_path / "deferred.lance").exists() + + sink.start() + assert (tmp_path / "deferred.lance").exists() + + +def test_sink_invalid_params_do_not_declare_table(tmp_path): + ns = _dir_ns(tmp_path) + schema = daft.from_pydict({"id": [1]}).schema() + + with pytest.raises(ValueError, match="blob_columns"): + daft_lance.LanceDataSink(None, schema, "create", table_id=["orphan"], blob_columns=["missing"], **ns) + + assert not (tmp_path / "orphan.lance").exists() + + +def test_namespace_create_on_existing_table_raises(tmp_path): + ns = _dir_ns(tmp_path) + table_id = ["exists"] + + daft_lance.write_lance(daft.from_pydict({"id": [1]}), table_id=table_id, mode="create", **ns).collect() + + with pytest.raises(ValueError, match="already exists"): + daft_lance.write_lance(daft.from_pydict({"id": [2]}), table_id=table_id, mode="create", **ns).collect() + + +def test_namespace_create_over_declared_placeholder(tmp_path): + import lance_namespace as ln + from lance_namespace import DeclareTableRequest + + ns = _dir_ns(tmp_path) + table_id = ["declared_first"] + + namespace = ln.connect("dir", {"root": str(tmp_path)}) + namespace.declare_table(DeclareTableRequest(id=table_id, location=None)) + + daft_lance.write_lance(daft.from_pydict({"id": [1, 2]}), table_id=table_id, mode="create", **ns).collect() + + assert daft_lance.read_lance(table_id=table_id, **ns).to_pydict() == {"id": [1, 2]} + + +def test_create_declare_race_recovers_placeholder(monkeypatch, tmp_path): + from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError + + class RacingNamespace: + def __init__(self): + self.describe_calls = 0 + + def describe_table(self, request): + self.describe_calls += 1 + if self.describe_calls == 1: + raise TableNotFoundError("table not found: t") + return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None, is_only_declared=True) + + def declare_table(self, request): + raise TableAlreadyExistsError("Table already exists: t") + + fake = RacingNamespace() + monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: fake) + + resolved = namespace_mod.resolve_namespace_table( + namespace_impl="rest", + namespace_properties=None, + table_id=["t"], + mode="create", + ) + assert resolved is not None + assert resolved.is_declared_placeholder + assert fake.describe_calls == 2 + + +def test_create_declare_race_lost_to_real_table_raises(monkeypatch, tmp_path): + from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError + + class RacingNamespace: + def __init__(self): + self.describe_calls = 0 + + def describe_table(self, request): + self.describe_calls += 1 + if self.describe_calls == 1: + raise TableNotFoundError("table not found: t") + return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None, is_only_declared=False) + + def declare_table(self, request): + raise TableAlreadyExistsError("Table already exists: t") + + monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: RacingNamespace()) + + with pytest.raises(ValueError, match="already exists"): + namespace_mod.resolve_namespace_table( + namespace_impl="rest", + namespace_properties=None, + table_id=["t"], + mode="create", + ) + + +def test_declared_placeholder_response_fallback_property(): + assert namespace_mod.is_declared_placeholder_response(SimpleNamespace(is_only_declared=True)) + assert namespace_mod.is_declared_placeholder_response( + SimpleNamespace(is_only_declared=None, properties={"lance.declared": "TRUE"}) + ) + assert not namespace_mod.is_declared_placeholder_response( + SimpleNamespace(is_only_declared=None, properties={}, metadata=None) + ) + + def test_namespace_rejects_uri_and_namespace(tmp_path): with pytest.raises(ValueError, match="Cannot provide both 'uri' and namespace parameters"): daft_lance.read_lance( From 37b02565faaf07712be8d8df5b0f9c1cd1c1e3b0 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 13:08:49 +0900 Subject: [PATCH 05/25] fix(namespace): thread io_config credentials into namespace-resolved locations Previously io_config was converted to storage options only when a plain uri was passed, so a namespace that resolves to s3://... without vending credentials silently ignored the user's io_config. construct_lance_dataset and LanceDataSink now derive storage options from io_config against the resolved location and layer them as: io_config-derived < user-provided storage_options < namespace-vended. Plain-uri behavior is unchanged (user-provided options still replace io_config-derived ones). Claude-Session: https://claude.ai/code/session_01VyD31uRyBSEmKvPQhLtM5E --- daft_lance/_lance.py | 21 ++---- daft_lance/lance_data_sink.py | 19 ++++-- daft_lance/namespace.py | 17 ++--- daft_lance/utils.py | 20 +++++- tests/io/lancedb/test_namespace.py | 104 +++++++++++++++++++++++++++++ 5 files changed, 149 insertions(+), 32 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 12bd836..5626c26 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -9,7 +9,6 @@ from daft.daft import IOConfig, ScanOperatorHandle from daft.dataframe import DataFrame from daft.io._checkpoint import attach_checkpoint -from daft.io.object_store_options import io_config_to_storage_options from daft.logical.builder import LogicalPlanBuilder from .lance_compaction import compact_files_internal @@ -148,11 +147,10 @@ def read_lance( ) io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - storage_options = io_config_to_storage_options(io_config, uri_str) if uri_str is not None else None ds = construct_lance_dataset( uri_str, - storage_options=storage_options, + io_config=io_config, namespace_impl=namespace_impl, namespace_properties=namespace_properties, table_id=table_id, @@ -243,13 +241,12 @@ def merge_columns( ) io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - if uri is not None: - storage_options = storage_options or io_config_to_storage_options(io_config, uri) # Build Lance dataset handle for committing lance_ds = construct_lance_dataset( uri, storage_options=storage_options, + io_config=io_config, namespace_impl=namespace_impl, namespace_properties=namespace_properties, table_id=table_id, @@ -346,13 +343,12 @@ def merge_columns_df( >>> daft_lance.merge_columns_df(df, "s3://my-lancedb-bucket/data/") """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - if uri is not None: - storage_options = storage_options or io_config_to_storage_options(io_config, uri) # Build Lance dataset handle for committing lance_ds = construct_lance_dataset( uri, storage_options=storage_options, + io_config=io_config, namespace_impl=namespace_impl, namespace_properties=namespace_properties, table_id=table_id, @@ -499,12 +495,11 @@ def create_scalar_index( >>> daft_lance.create_scalar_index("s3://my-bucket/dataset/", column="title", replace=False) """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - if uri is not None: - storage_options = storage_options or io_config_to_storage_options(io_config, str(uri)) lance_ds = construct_lance_dataset( uri, storage_options=storage_options, + io_config=io_config, namespace_impl=namespace_impl, namespace_properties=namespace_properties, table_id=table_id, @@ -591,14 +586,11 @@ def compact_files( RuntimeError: When compaction fails or no successful results """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - if uri is not None: - storage_options = storage_options or io_config_to_storage_options( - io_config, str(uri) if isinstance(uri, pathlib.Path) else uri - ) lance_ds = construct_lance_dataset( uri, storage_options=storage_options, + io_config=io_config, namespace_impl=namespace_impl, namespace_properties=namespace_properties, table_id=table_id, @@ -709,14 +701,13 @@ def write_lance( io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config storage_options = kwargs.pop("storage_options", None) - if uri is not None: - storage_options = storage_options or io_config_to_storage_options(io_config, str(uri)) lance_ds: LanceDataset | None try: lance_ds = construct_lance_dataset( uri, storage_options=storage_options, + io_config=io_config, namespace_impl=namespace_impl, namespace_properties=namespace_properties, table_id=table_id, diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index a995e37..8919772 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -125,7 +125,7 @@ def start(self) -> None: resolved = self._resolve_table() self._table_uri = resolved.uri self._declared_placeholder = resolved.is_declared_placeholder - self._storage_options = merge_storage_options(self._base_storage_options(), resolved.storage_options) + self._storage_options = self._merged_storage_options(resolved) existing = self._absorb_existing_dataset() existing_version = getattr(existing, "data_storage_version", None) if existing is not None else None @@ -165,12 +165,17 @@ def _resolve_table(self) -> ResolvedNamespaceTable: raise ValueError("Unable to resolve Lance dataset URI from namespace.") return resolved - def _base_storage_options(self) -> dict[str, str] | None: - if self._uri is None: - return self._user_storage_options - if self._user_storage_options is not None: - return self._user_storage_options - return io_config_to_storage_options(self._io_config, str(self._uri)) + def _merged_storage_options(self, resolved: ResolvedNamespaceTable) -> dict[str, str] | None: + """Layer storage options: io_config-derived < user-provided < namespace-vended. + + For a plain uri, user-provided options replace the io_config-derived ones + entirely (historical behavior). + """ + io_derived = io_config_to_storage_options(self._io_config, resolved.uri) + if self._uri is not None: + base = self._user_storage_options if self._user_storage_options is not None else io_derived + return merge_storage_options(base, resolved.storage_options) + return merge_storage_options(io_derived, self._user_storage_options, resolved.storage_options) @staticmethod def _reject_unsupported_modes( diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index cb58113..fa46768 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -289,15 +289,16 @@ def _classify_existing_for_create(response: Any, table_id: list[str]) -> Resolve ) -def merge_storage_options( - storage_options: dict[str, Any] | None, - namespace_storage_options: dict[str, str] | None, -) -> dict[str, Any] | None: +def merge_storage_options(*layers: dict[str, Any] | None) -> dict[str, Any] | None: + """Merge storage-option layers; later layers take precedence. + + Callers order layers from lowest to highest priority, conventionally: + io_config-derived < user-provided ``storage_options`` < namespace-vended. + """ merged: dict[str, Any] = {} - if storage_options: - merged.update(storage_options) - if namespace_storage_options: - merged.update(namespace_storage_options) + for layer in layers: + if layer: + merged.update(layer) return merged or None diff --git a/daft_lance/utils.py b/daft_lance/utils.py index d072b23..4b8fbe7 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -6,6 +6,7 @@ import lance from daft.dependencies import pa +from daft.io.object_store_options import io_config_to_storage_options from daft.logical.schema import Schema as DaftSchema from daft_lance.namespace import ( get_namespace_kwargs, @@ -18,6 +19,8 @@ if TYPE_CHECKING: import pathlib + from daft.daft import IOConfig + logger = logging.getLogger(__name__) @@ -92,12 +95,19 @@ def construct_lance_dataset( uri: str | pathlib.Path | None, version: int | str | None = None, storage_options: dict[str, Any] | None = None, + io_config: IOConfig | None = None, namespace_impl: str | None = None, namespace_properties: dict[str, str] | None = None, table_id: list[str] | None = None, **kwargs: Any, ) -> lance.LanceDataset: - """Construct a Lance dataset with common options.""" + """Construct a Lance dataset with common options. + + Storage options are layered from lowest to highest priority: + io_config-derived < user-provided ``storage_options`` < namespace-vended. + For a plain ``uri``, user-provided ``storage_options`` replace the + io_config-derived ones entirely (historical behavior). + """ validate_uri_or_namespace(uri, namespace_impl, table_id) resolved_uri = str(uri) if uri is not None else None namespace_storage_options = None @@ -113,7 +123,13 @@ def construct_lance_dataset( namespace_storage_options = resolved.storage_options if resolved_uri is None: raise ValueError("Unable to resolve Lance dataset URI.") - merged_storage_options = merge_storage_options(storage_options, namespace_storage_options) + + io_derived_options = io_config_to_storage_options(io_config, resolved_uri) if io_config is not None else None + if uri is not None: + base_options = storage_options if storage_options is not None else io_derived_options + merged_storage_options = merge_storage_options(base_options, namespace_storage_options) + else: + merged_storage_options = merge_storage_options(io_derived_options, storage_options, namespace_storage_options) original_default_scan_options = kwargs.pop("default_scan_options", None) safe_default_scan_options = None diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index e5a954b..1db1e76 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -329,6 +329,110 @@ def test_declared_placeholder_response_fallback_property(): ) +def test_construct_lance_dataset_storage_options_priority(monkeypatch): + import daft_lance.utils as utils_mod + from daft.io import IOConfig, S3Config + + captured = {} + + class FakeDataset: + pass + + def fake_dataset(uri, storage_options=None, version=None, **kwargs): + captured["uri"] = uri + captured["storage_options"] = storage_options + return FakeDataset() + + monkeypatch.setattr(utils_mod.lance, "dataset", fake_dataset) + monkeypatch.setattr(utils_mod, "get_namespace_kwargs", lambda *args: {}) + monkeypatch.setattr( + utils_mod, + "resolve_namespace_table", + lambda **kwargs: namespace_mod.ResolvedNamespaceTable( + uri="s3://bucket/t.lance", + storage_options={"access_key_id": "vended-key", "session_token": "vended-token"}, + ), + ) + + io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) + utils_mod.construct_lance_dataset( + None, + io_config=io_config, + storage_options={"access_key_id": "user-key", "user_option": "kept"}, + namespace_impl="rest", + namespace_properties={"uri": "http://namespace.example"}, + table_id=["t"], + ) + + merged = captured["storage_options"] + assert merged["access_key_id"] == "vended-key" # namespace-vended beats user-provided + assert merged["session_token"] == "vended-token" + assert merged["user_option"] == "kept" # user-provided keys survive + assert merged["secret_access_key"] == "io-secret" # io_config fills the gaps + assert captured["uri"] is None # namespace addressing passes uri=None + + +def test_construct_lance_dataset_io_config_reaches_namespace_location(monkeypatch): + import daft_lance.utils as utils_mod + from daft.io import IOConfig, S3Config + + captured = {} + + class FakeDataset: + pass + + def fake_dataset(uri, storage_options=None, version=None, **kwargs): + captured["storage_options"] = storage_options + return FakeDataset() + + monkeypatch.setattr(utils_mod.lance, "dataset", fake_dataset) + monkeypatch.setattr(utils_mod, "get_namespace_kwargs", lambda *args: {}) + monkeypatch.setattr( + utils_mod, + "resolve_namespace_table", + lambda **kwargs: namespace_mod.ResolvedNamespaceTable(uri="s3://bucket/t.lance", storage_options=None), + ) + + io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) + utils_mod.construct_lance_dataset( + None, + io_config=io_config, + namespace_impl="rest", + namespace_properties={"uri": "http://namespace.example"}, + table_id=["t"], + ) + + merged = captured["storage_options"] + assert merged["access_key_id"] == "io-key" + assert merged["secret_access_key"] == "io-secret" + + +def test_sink_storage_options_priority(tmp_path): + from daft.io import IOConfig, S3Config + + ns = _dir_ns(tmp_path) + schema = daft.from_pydict({"id": [1]}).schema() + io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) + + sink = daft_lance.LanceDataSink( + None, + schema, + "create", + io_config, + table_id=["t"], + storage_options={"user_option": "kept", "access_key_id": "user-key"}, + **ns, + ) + resolved = namespace_mod.ResolvedNamespaceTable( + uri="s3://bucket/t.lance", storage_options={"access_key_id": "vended-key"} + ) + + merged = sink._merged_storage_options(resolved) + assert merged["access_key_id"] == "vended-key" + assert merged["user_option"] == "kept" + assert merged["secret_access_key"] == "io-secret" + + def test_namespace_rejects_uri_and_namespace(tmp_path): with pytest.raises(ValueError, match="Cannot provide both 'uri' and namespace parameters"): daft_lance.read_lance( From 64501fbab62946cecc544ebf84cb79ca92565f01 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 13:39:37 +0900 Subject: [PATCH 06/25] chore: zero out mypy errors in daft_lance sources and namespace tests daft_lance/ sources and the namespace test files now pass mypy strict. The bulk came from LanceDataSink._namespace_kwargs being annotated dict[str, object], which poisoned every **-expansion into typed lance APIs. Also widens read_lance's default_scan_options to dict[str, Any] to match the documented usage ({"with_row_address": True}) and the other entry points. Pre-existing errors in the rest of tests/ are out of scope for this branch. Claude-Session: https://claude.ai/code/session_01VyD31uRyBSEmKvPQhLtM5E --- daft_lance/_lance.py | 2 +- daft_lance/lance_data_sink.py | 11 +-- tests/io/lancedb/test_namespace.py | 92 +++++++++++---------- tests/io/lancedb/test_namespace_rest_e2e.py | 12 +-- 4 files changed, 63 insertions(+), 54 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 5626c26..d60e044 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -47,7 +47,7 @@ def read_lance( block_size: int | None = None, commit_lock: object | None = None, index_cache_size: int | None = None, - default_scan_options: dict[str, str] | None = None, + default_scan_options: dict[str, Any] | None = None, metadata_cache_size_bytes: int | None = None, fragment_group_size: int | None = None, include_fragment_id: bool | None = None, diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index 8919772..c9b0745 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -5,7 +5,7 @@ import uuid import warnings from itertools import chain -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Any, Literal import lance from lance.fragment import FragmentMetadata @@ -144,7 +144,7 @@ def start(self) -> None: self._effective_pyarrow_schema = self._blob.build_effective_schema(self._pyarrow_schema) @property - def _namespace_kwargs(self) -> dict[str, object]: + def _namespace_kwargs(self) -> dict[str, Any]: return get_namespace_kwargs(self._namespace_impl, self._namespace_properties, self._table_id) @property @@ -255,7 +255,8 @@ def _absorb_existing_dataset(self) -> lance.LanceDataset | None: raise FileExistsError("Target path points to a file, cannot create a dataset here.") return None - self._table_schema = dataset.schema + table_schema = dataset.schema + self._table_schema = table_schema self._version = dataset.latest_version if self._mode == "create": @@ -270,8 +271,8 @@ def _absorb_existing_dataset(self) -> lance.LanceDataset | None: self._table_schema = None if self._mode == "append" and not _pyarrow_schema_castable( - blob_aware_schema_for_validation(self._pyarrow_schema, self._table_schema), - blob_aware_schema_for_validation(self._table_schema, self._table_schema), + blob_aware_schema_for_validation(self._pyarrow_schema, table_schema), + blob_aware_schema_for_validation(table_schema, table_schema), ): raise ValueError( "Schema of data does not match table schema\n" diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index 1db1e76..2463bae 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -1,6 +1,8 @@ from __future__ import annotations +from pathlib import Path from types import SimpleNamespace +from typing import Any import pytest @@ -11,11 +13,11 @@ pytestmark = pytest.mark.lance_native_teardown_crash_workaround -def _dir_ns(tmp_path): +def _dir_ns(tmp_path: Path) -> dict[str, Any]: return {"namespace_impl": "dir", "namespace_properties": {"root": str(tmp_path)}} -def test_namespace_write_read_append_roundtrip(tmp_path): +def test_namespace_write_read_append_roundtrip(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["roundtrip"] @@ -30,7 +32,7 @@ def test_namespace_write_read_append_roundtrip(tmp_path): assert result == {"id": [1, 2, 3], "label": ["a", "b", "c"]} -def test_namespace_overwrite(tmp_path): +def test_namespace_overwrite(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["overwrite_tbl"] @@ -41,7 +43,7 @@ def test_namespace_overwrite(tmp_path): assert result == {"id": [7, 8, 9]} -def test_namespace_overwrite_missing_table_declares(tmp_path): +def test_namespace_overwrite_missing_table_declares(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["overwrite_fresh"] @@ -51,14 +53,16 @@ def test_namespace_overwrite_missing_table_declares(tmp_path): assert result == {"id": [1]} -def test_namespace_overwrite_does_not_declare_on_ambiguous_error(monkeypatch, tmp_path): +def test_namespace_overwrite_does_not_declare_on_ambiguous_error( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: class FakeNamespace: declared = False - def describe_table(self, request): + def describe_table(self, request: Any) -> Any: raise RuntimeError("permission denied: parent catalog does not exist") - def declare_table(self, request): + def declare_table(self, request: Any) -> Any: self.declared = True return SimpleNamespace(location=str(tmp_path / "should_not_exist.lance")) @@ -76,14 +80,14 @@ def declare_table(self, request): assert not namespace.declared -def test_namespace_untyped_404_table_not_found_is_missing(): +def test_namespace_untyped_404_table_not_found_is_missing() -> None: class RestTableNotFound(RuntimeError): status = 404 assert namespace_mod.is_table_not_found(RestTableNotFound("table catalog.schema.table not found")) -def test_namespace_read_supports_pushdowns(tmp_path): +def test_namespace_read_supports_pushdowns(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["pushdowns"] @@ -100,12 +104,13 @@ def test_namespace_read_supports_pushdowns(tmp_path): **ns, ).collect() - result = daft_lance.read_lance(table_id=table_id, **ns).where(daft.col("score") > 10).select("label").to_pydict() + predicate = daft.col("score") > 10 # type: ignore[operator] + result = daft_lance.read_lance(table_id=table_id, **ns).where(predicate).select("label").to_pydict() assert result == {"label": ["b", "c"]} -def test_namespace_count_pushdown(tmp_path): +def test_namespace_count_pushdown(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["count_tbl"] @@ -114,7 +119,7 @@ def test_namespace_count_pushdown(tmp_path): assert daft_lance.read_lance(table_id=table_id, **ns).count_rows() == 10 -def test_namespace_patch_daft(tmp_path): +def test_namespace_patch_daft(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["patched"] @@ -122,7 +127,7 @@ def test_namespace_patch_daft(tmp_path): daft_lance.patch_daft() # idempotent daft.from_pydict({"id": [1, 2]}).write_lance(table_id=table_id, mode="create", **ns).collect() - result = daft.read_lance(table_id=table_id, **ns).to_pydict() + result = daft.read_lance(table_id=table_id, **ns).to_pydict() # type: ignore[call-arg] assert result == {"id": [1, 2]} # Plain-URI writes keep working through the patched method. @@ -131,7 +136,7 @@ def test_namespace_patch_daft(tmp_path): assert daft.read_lance(uri).to_pydict() == {"id": [5]} -def test_namespace_write_lance_merge_new_columns(tmp_path): +def test_namespace_write_lance_merge_new_columns(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["merge_tbl"] @@ -155,7 +160,7 @@ def test_namespace_write_lance_merge_new_columns(tmp_path): assert result["doubled"] == [20, 40, 60] -def test_namespace_merge_columns_df(tmp_path): +def test_namespace_merge_columns_df(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["merge_cols_df"] @@ -179,7 +184,7 @@ def test_namespace_merge_columns_df(tmp_path): assert result["tripled"] == [3, 6, 9] -def test_namespace_create_scalar_index(tmp_path): +def test_namespace_create_scalar_index(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["indexed"] @@ -199,10 +204,10 @@ def test_namespace_create_scalar_index(tmp_path): namespace = ln.connect("dir", {"root": str(tmp_path)}) location = namespace.describe_table(DescribeTableRequest(id=table_id)).location indices = lance.dataset(location).list_indices() - assert any(idx["fields"] == ["price"] for idx in indices) + assert any(idx["fields"] == ["price"] for idx in indices) # type: ignore[index] -def test_namespace_compact_files(tmp_path): +def test_namespace_compact_files(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["compacted"] @@ -216,7 +221,7 @@ def test_namespace_compact_files(tmp_path): assert result == {"id": [1, 2, 3, 4]} -def test_sink_construction_is_side_effect_free(tmp_path): +def test_sink_construction_is_side_effect_free(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) schema = daft.from_pydict({"id": [1]}).schema() @@ -227,7 +232,7 @@ def test_sink_construction_is_side_effect_free(tmp_path): assert (tmp_path / "deferred.lance").exists() -def test_sink_invalid_params_do_not_declare_table(tmp_path): +def test_sink_invalid_params_do_not_declare_table(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) schema = daft.from_pydict({"id": [1]}).schema() @@ -237,7 +242,7 @@ def test_sink_invalid_params_do_not_declare_table(tmp_path): assert not (tmp_path / "orphan.lance").exists() -def test_namespace_create_on_existing_table_raises(tmp_path): +def test_namespace_create_on_existing_table_raises(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["exists"] @@ -247,7 +252,7 @@ def test_namespace_create_on_existing_table_raises(tmp_path): daft_lance.write_lance(daft.from_pydict({"id": [2]}), table_id=table_id, mode="create", **ns).collect() -def test_namespace_create_over_declared_placeholder(tmp_path): +def test_namespace_create_over_declared_placeholder(tmp_path: Path) -> None: import lance_namespace as ln from lance_namespace import DeclareTableRequest @@ -262,20 +267,20 @@ def test_namespace_create_over_declared_placeholder(tmp_path): assert daft_lance.read_lance(table_id=table_id, **ns).to_pydict() == {"id": [1, 2]} -def test_create_declare_race_recovers_placeholder(monkeypatch, tmp_path): +def test_create_declare_race_recovers_placeholder(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError class RacingNamespace: - def __init__(self): + def __init__(self) -> None: self.describe_calls = 0 - def describe_table(self, request): + def describe_table(self, request: Any) -> Any: self.describe_calls += 1 if self.describe_calls == 1: raise TableNotFoundError("table not found: t") return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None, is_only_declared=True) - def declare_table(self, request): + def declare_table(self, request: Any) -> Any: raise TableAlreadyExistsError("Table already exists: t") fake = RacingNamespace() @@ -292,20 +297,20 @@ def declare_table(self, request): assert fake.describe_calls == 2 -def test_create_declare_race_lost_to_real_table_raises(monkeypatch, tmp_path): +def test_create_declare_race_lost_to_real_table_raises(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError class RacingNamespace: - def __init__(self): + def __init__(self) -> None: self.describe_calls = 0 - def describe_table(self, request): + def describe_table(self, request: Any) -> Any: self.describe_calls += 1 if self.describe_calls == 1: raise TableNotFoundError("table not found: t") return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None, is_only_declared=False) - def declare_table(self, request): + def declare_table(self, request: Any) -> Any: raise TableAlreadyExistsError("Table already exists: t") monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: RacingNamespace()) @@ -319,7 +324,7 @@ def declare_table(self, request): ) -def test_declared_placeholder_response_fallback_property(): +def test_declared_placeholder_response_fallback_property() -> None: assert namespace_mod.is_declared_placeholder_response(SimpleNamespace(is_only_declared=True)) assert namespace_mod.is_declared_placeholder_response( SimpleNamespace(is_only_declared=None, properties={"lance.declared": "TRUE"}) @@ -329,7 +334,7 @@ def test_declared_placeholder_response_fallback_property(): ) -def test_construct_lance_dataset_storage_options_priority(monkeypatch): +def test_construct_lance_dataset_storage_options_priority(monkeypatch: pytest.MonkeyPatch) -> None: import daft_lance.utils as utils_mod from daft.io import IOConfig, S3Config @@ -338,12 +343,12 @@ def test_construct_lance_dataset_storage_options_priority(monkeypatch): class FakeDataset: pass - def fake_dataset(uri, storage_options=None, version=None, **kwargs): + def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **kwargs: Any) -> Any: captured["uri"] = uri captured["storage_options"] = storage_options return FakeDataset() - monkeypatch.setattr(utils_mod.lance, "dataset", fake_dataset) + monkeypatch.setattr("daft_lance.utils.lance.dataset", fake_dataset) monkeypatch.setattr(utils_mod, "get_namespace_kwargs", lambda *args: {}) monkeypatch.setattr( utils_mod, @@ -365,6 +370,7 @@ def fake_dataset(uri, storage_options=None, version=None, **kwargs): ) merged = captured["storage_options"] + assert merged is not None assert merged["access_key_id"] == "vended-key" # namespace-vended beats user-provided assert merged["session_token"] == "vended-token" assert merged["user_option"] == "kept" # user-provided keys survive @@ -372,7 +378,7 @@ def fake_dataset(uri, storage_options=None, version=None, **kwargs): assert captured["uri"] is None # namespace addressing passes uri=None -def test_construct_lance_dataset_io_config_reaches_namespace_location(monkeypatch): +def test_construct_lance_dataset_io_config_reaches_namespace_location(monkeypatch: pytest.MonkeyPatch) -> None: import daft_lance.utils as utils_mod from daft.io import IOConfig, S3Config @@ -381,11 +387,11 @@ def test_construct_lance_dataset_io_config_reaches_namespace_location(monkeypatc class FakeDataset: pass - def fake_dataset(uri, storage_options=None, version=None, **kwargs): + def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **kwargs: Any) -> Any: captured["storage_options"] = storage_options return FakeDataset() - monkeypatch.setattr(utils_mod.lance, "dataset", fake_dataset) + monkeypatch.setattr("daft_lance.utils.lance.dataset", fake_dataset) monkeypatch.setattr(utils_mod, "get_namespace_kwargs", lambda *args: {}) monkeypatch.setattr( utils_mod, @@ -403,11 +409,12 @@ def fake_dataset(uri, storage_options=None, version=None, **kwargs): ) merged = captured["storage_options"] + assert merged is not None assert merged["access_key_id"] == "io-key" assert merged["secret_access_key"] == "io-secret" -def test_sink_storage_options_priority(tmp_path): +def test_sink_storage_options_priority(tmp_path: Path) -> None: from daft.io import IOConfig, S3Config ns = _dir_ns(tmp_path) @@ -428,12 +435,13 @@ def test_sink_storage_options_priority(tmp_path): ) merged = sink._merged_storage_options(resolved) + assert merged is not None assert merged["access_key_id"] == "vended-key" assert merged["user_option"] == "kept" assert merged["secret_access_key"] == "io-secret" -def test_namespace_rejects_uri_and_namespace(tmp_path): +def test_namespace_rejects_uri_and_namespace(tmp_path: Path) -> None: with pytest.raises(ValueError, match="Cannot provide both 'uri' and namespace parameters"): daft_lance.read_lance( str(tmp_path / "dataset"), @@ -443,7 +451,7 @@ def test_namespace_rejects_uri_and_namespace(tmp_path): ) -def test_namespace_requires_table_id(tmp_path): +def test_namespace_requires_table_id(tmp_path: Path) -> None: with pytest.raises(ValueError, match="'table_id' must be provided"): daft_lance.read_lance( namespace_impl="dir", @@ -451,11 +459,11 @@ def test_namespace_requires_table_id(tmp_path): ) -def test_namespace_requires_impl(): +def test_namespace_requires_impl() -> None: with pytest.raises(ValueError, match="'namespace_impl' must be provided"): daft_lance.read_lance(table_id=["tbl"]) -def test_namespace_requires_uri_or_namespace(): +def test_namespace_requires_uri_or_namespace() -> None: with pytest.raises(ValueError, match="Must provide either 'uri' OR"): daft_lance.read_lance() diff --git a/tests/io/lancedb/test_namespace_rest_e2e.py b/tests/io/lancedb/test_namespace_rest_e2e.py index a193e2e..3e388e0 100644 --- a/tests/io/lancedb/test_namespace_rest_e2e.py +++ b/tests/io/lancedb/test_namespace_rest_e2e.py @@ -2,6 +2,7 @@ import os import uuid +from typing import Any import pytest @@ -23,7 +24,7 @@ def test_rest_namespace_write_read_append_roundtrip() -> None: catalog = os.environ.get("DAFT_LANCE_REST_CATALOG", "lance_catalog") schema = os.environ.get("DAFT_LANCE_REST_SCHEMA", "daft_ns_e2e") table_id = [catalog, schema, f"orders_{uuid.uuid4().hex[:8]}"] - ns = {"namespace_impl": "rest", "namespace_properties": namespace_properties} + ns: dict[str, Any] = {"namespace_impl": "rest", "namespace_properties": namespace_properties} namespace = ln.connect("rest", namespace_properties) try: @@ -56,9 +57,8 @@ def test_rest_namespace_write_read_append_roundtrip() -> None: "score": [10, 20, 30, 40, 50], } - filtered = ( - daft_lance.read_lance(table_id=table_id, **ns).where(daft.col("score") >= 30).select("id", "label").to_pydict() - ) + predicate = daft.col("score") >= 30 # type: ignore[operator] + filtered = daft_lance.read_lance(table_id=table_id, **ns).where(predicate).select("id", "label").to_pydict() assert filtered == {"id": [3, 4, 5], "label": ["c", "d", "e"]} assert daft_lance.read_lance(table_id=table_id, **ns).count_rows() == 5 @@ -71,9 +71,9 @@ def test_rest_namespace_patch_daft_roundtrip() -> None: catalog = os.environ.get("DAFT_LANCE_REST_CATALOG", "lance_catalog") schema = os.environ.get("DAFT_LANCE_REST_SCHEMA", "daft_ns_e2e") table_id = [catalog, schema, f"patched_{uuid.uuid4().hex[:8]}"] - ns = {"namespace_impl": "rest", "namespace_properties": namespace_properties} + ns: dict[str, Any] = {"namespace_impl": "rest", "namespace_properties": namespace_properties} daft_lance.patch_daft() daft.from_pydict({"id": [1, 2]}).write_lance(table_id=table_id, mode="create", **ns).collect() - assert daft.read_lance(table_id=table_id, **ns).to_pydict() == {"id": [1, 2]} + assert daft.read_lance(table_id=table_id, **ns).to_pydict() == {"id": [1, 2]} # type: ignore[call-arg] From 0feda988381f96ff467f7ce61716dc8375f3f7e2 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 14:18:39 +0900 Subject: [PATCH 07/25] docs(namespace): document namespace params everywhere and validate orphan properties - read_lance / merge_columns(_df) / create_scalar_index / compact_files docstrings now describe table_id / namespace_impl / namespace_properties - namespace_properties without namespace_impl now fails fast instead of being silently ignored - README documents io_config fallback for namespace locations and the DAFT_LANCE_NAMESPACE_CACHE_SIZE environment variable Claude-Session: https://claude.ai/code/session_01VyD31uRyBSEmKvPQhLtM5E --- README.md | 8 +++++++- daft_lance/_lance.py | 27 ++++++++++++++++++++++++++- daft_lance/lance_data_sink.py | 2 +- daft_lance/namespace.py | 3 +++ daft_lance/utils.py | 2 +- tests/io/lancedb/test_namespace.py | 5 +++++ 6 files changed, 43 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index e88e5f1..0d9da07 100644 --- a/README.md +++ b/README.md @@ -81,7 +81,13 @@ daft_lance.read_lance(table_id=table_id, **namespace).show() When the catalog holds the storage configuration (bucket, endpoint, credentials), the `describe_table` response vends `storage_options` to the client, so you do not need to pass -object-store credentials yourself. +object-store credentials yourself. If the namespace does not vend credentials, your +`io_config` (or explicit `storage_options`) is applied to the resolved location; when both +are present, namespace-vended options take precedence. + +Namespace clients are cached per (implementation, properties) pair. The cache size defaults +to 16 and can be tuned with the `DAFT_LANCE_NAMESPACE_CACHE_SIZE` environment variable +(read once at import time). #### Drop-in Daft APIs diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index d60e044..3584b89 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -63,6 +63,11 @@ def read_lance( uri: The URI of the Lance table to read from. Accepts a local path or an object-store URI like "s3://bucket/path". io_config: A custom IOConfig to use when accessing LanceDB data. Defaults to None. + table_id: Table identifier within the namespace, e.g. ["catalog", "schema", "table"]. + Mutually exclusive with ``uri``. + namespace_impl: Lance Namespace implementation, e.g. "dir" or "rest". + namespace_properties: Properties for connecting to the namespace, e.g. + {"root": "/data"} for "dir" or {"uri": "http://host:port"} for "rest". version : optional, int | str If specified, load a specific version of the Lance dataset. Else, loads the latest version. A version number (`int`) or a tag (`str`) can be provided. @@ -205,6 +210,11 @@ def merge_columns( Args: uri: The URI of the Lance table (supports remote URLs to object stores such as `s3://` or `gs://`) io_config: A custom IOConfig to use when accessing LanceDB data. Defaults to None. + table_id: Table identifier within the namespace, e.g. ["catalog", "schema", "table"]. + Mutually exclusive with ``uri``. + namespace_impl: Lance Namespace implementation, e.g. "dir" or "rest". + namespace_properties: Properties for connecting to the namespace, e.g. + {"root": "/data"} for "dir" or {"uri": "http://host:port"} for "rest". transform: A transformation function or UDF to apply to the data. read_columns: List of column names to read for the transformation. reader_schema: Schema for the reader. @@ -305,6 +315,11 @@ def merge_columns_df( df: DataFrame containing the new columns to merge along with fragment_id and join key columns uri: URL to the LanceDB table (supports remote URLs to object stores such as `s3://` or `gs://`) io_config: A custom IOConfig to use when accessing LanceDB data. Defaults to None. + table_id: Table identifier within the namespace, e.g. ["catalog", "schema", "table"]. + Mutually exclusive with ``uri``. + namespace_impl: Lance Namespace implementation, e.g. "dir" or "rest". + namespace_properties: Properties for connecting to the namespace, e.g. + {"root": "/data"} for "dir" or {"uri": "http://host:port"} for "rest". read_columns: List of column names to read for the transformation. reader_schema: Schema for the reader. storage_options: Extra options for storage connection. @@ -417,6 +432,11 @@ def create_scalar_index( Args: uri: The URI of the Lance table (supports remote URLs to object stores such as `s3://` or `gs://`) io_config: A custom IOConfig to use when accessing LanceDB data. Defaults to None. + table_id: Table identifier within the namespace, e.g. ["catalog", "schema", "table"]. + Mutually exclusive with ``uri``. + namespace_impl: Lance Namespace implementation, e.g. "dir" or "rest". + namespace_properties: Properties for connecting to the namespace, e.g. + {"root": "/data"} for "dir" or {"uri": "http://host:port"} for "rest". column: Column name to index index_type: Type of index to build. For distributed segmented execution this supports "BITMAP", "BTREE", "INVERTED", and "FTS". @@ -556,6 +576,11 @@ def compact_files( Args: uri: The URI of the Lance table (supports remote URLs to object stores such as `s3://` or `gs://`) io_config: A custom IOConfig to use when accessing LanceDB data. Defaults to None. + table_id: Table identifier within the namespace, e.g. ["catalog", "schema", "table"]. + Mutually exclusive with ``uri``. + namespace_impl: Lance Namespace implementation, e.g. "dir" or "rest". + namespace_properties: Properties for connecting to the namespace, e.g. + {"root": "/data"} for "dir" or {"uri": "http://host:port"} for "rest". storage_options: Extra options for storage connection. version: If specified, load a specific version of the Lance dataset. asof: If specified, find the latest version created on or earlier than the given argument value. @@ -680,7 +705,7 @@ def write_lance( ... df, namespace_impl="dir", namespace_properties={"root": "/tmp/tables"}, table_id=["t"] ... ).collect() # doctest: +SKIP """ - validate_uri_or_namespace(uri, namespace_impl, table_id) + validate_uri_or_namespace(uri, namespace_impl, table_id, namespace_properties) if schema is None: schema = df.schema() diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index c9b0745..45c3dc7 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -68,7 +68,7 @@ def __init__( compact_after_write: bool = True, ) -> None: self._reject_unsupported_modes(mode, use_legacy_format) - validate_uri_or_namespace(uri, namespace_impl, table_id) + validate_uri_or_namespace(uri, namespace_impl, table_id, namespace_properties) if uri is not None and not isinstance(uri, (str, pathlib.Path)): raise TypeError(f"Expected URI to be str or pathlib.Path, got {type(uri)}") diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index fa46768..429051b 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -19,10 +19,13 @@ def validate_uri_or_namespace( uri: str | os.PathLike[str] | None, namespace_impl: str | None, table_id: list[str] | None, + namespace_properties: dict[str, str] | None = None, ) -> None: has_uri = uri is not None has_ns = has_namespace_params(namespace_impl, table_id) + if namespace_properties is not None and namespace_impl is None: + raise ValueError("'namespace_impl' must be provided when 'namespace_properties' is provided.") if namespace_impl is not None and table_id is None: raise ValueError("'table_id' must be provided when 'namespace_impl' is provided.") if table_id is not None and namespace_impl is None: diff --git a/daft_lance/utils.py b/daft_lance/utils.py index 4b8fbe7..882e1d0 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -108,7 +108,7 @@ def construct_lance_dataset( For a plain ``uri``, user-provided ``storage_options`` replace the io_config-derived ones entirely (historical behavior). """ - validate_uri_or_namespace(uri, namespace_impl, table_id) + validate_uri_or_namespace(uri, namespace_impl, table_id, namespace_properties) resolved_uri = str(uri) if uri is not None else None namespace_storage_options = None if resolved_uri is None: diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index 2463bae..b7e542a 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -459,6 +459,11 @@ def test_namespace_requires_table_id(tmp_path: Path) -> None: ) +def test_namespace_properties_require_impl(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="'namespace_impl' must be provided when 'namespace_properties'"): + daft_lance.read_lance(str(tmp_path / "dataset"), namespace_properties={"root": str(tmp_path)}) + + def test_namespace_requires_impl() -> None: with pytest.raises(ValueError, match="'namespace_impl' must be provided"): daft_lance.read_lance(table_id=["tbl"]) From 2e719564172240a442d8eda8e735102daadf6384 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 14:30:38 +0900 Subject: [PATCH 08/25] fix(namespace): address independent review findings - overwrite-mode resolution now describes with check_declared=True (initial and declare-race recovery): strict namespace impls 404 a plain describe on declared-only stubs, so the previous recovery path failed in exactly the case it existed for. A declared-only stub is a valid overwrite target. - mode=create over a namespace placeholder that materialized a dataset is rejected with a clear error when use_mem_wal=True; the MemWAL path cannot apply the overwrite-placeholder semantics. - plain-uri entry points treat storage_options={} as unset again (falsy fallthrough to io_config-derived options), matching pre-branch behavior. Claude-Session: https://claude.ai/code/session_01VyD31uRyBSEmKvPQhLtM5E --- daft_lance/lance_data_sink.py | 5 ++ daft_lance/namespace.py | 33 +++++++------ daft_lance/utils.py | 4 +- tests/io/lancedb/test_namespace.py | 74 ++++++++++++++++++++++++++++++ 4 files changed, 102 insertions(+), 14 deletions(-) diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index 45c3dc7..35f04d9 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -267,6 +267,11 @@ def _absorb_existing_dataset(self) -> lance.LanceDataset | None: ) # The namespace declared a stub table and this location already holds # a (placeholder) dataset; overwrite it instead of failing the create. + if self._use_mem_wal: + raise ValueError( + "use_mem_wal=True cannot create over a namespace-declared placeholder " + 'dataset; write with use_mem_wal=False or use mode="overwrite".' + ) self._overwrite_placeholder_dataset = True self._table_schema = None diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index 429051b..11fda3f 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -246,12 +246,18 @@ def resolve_namespace_table( if mode == "create": return _resolve_for_create(namespace, table_id) - try: - return _resolved_from_response(_describe_table(namespace, table_id)) - except Exception as error: - if mode == "overwrite" and is_table_not_found(error): - return _declare_or_recover(namespace, table_id, check_declared=False) - raise + if mode == "overwrite": + # check_declared: a plain describe may 404 on declared-only stubs, and a + # stub is a perfectly fine overwrite target. + try: + response = _describe_table(namespace, table_id, check_declared=True) + except Exception as error: + if not is_table_not_found(error): + raise + return _declare_or_recover(namespace, table_id, for_create=False) + return _resolved_from_response(response, is_declared_placeholder=is_declared_placeholder_response(response)) + + return _resolved_from_response(_describe_table(namespace, table_id)) def _resolve_for_create(namespace: Any, table_id: list[str]) -> ResolvedNamespaceTable: @@ -266,21 +272,22 @@ def _resolve_for_create(namespace: Any, table_id: list[str]) -> ResolvedNamespac except Exception as error: if not is_table_not_found(error): raise - return _declare_or_recover(namespace, table_id, check_declared=True) + return _declare_or_recover(namespace, table_id, for_create=True) return _classify_existing_for_create(response, table_id) -def _declare_or_recover(namespace: Any, table_id: list[str], *, check_declared: bool) -> ResolvedNamespaceTable: +def _declare_or_recover(namespace: Any, table_id: list[str], *, for_create: bool) -> ResolvedNamespaceTable: try: return _declare_table(namespace, table_id) except Exception as error: if not is_table_already_exists(error): raise - # Lost a declare race; re-describe and classify what won. - response = _describe_table(namespace, table_id, check_declared=check_declared) - if not check_declared: - return _resolved_from_response(response) - return _classify_existing_for_create(response, table_id) + # Lost a declare race; re-describe (check_declared so a declared-only + # winner does not 404 on strict impls) and classify what won. + response = _describe_table(namespace, table_id, check_declared=True) + if for_create: + return _classify_existing_for_create(response, table_id) + return _resolved_from_response(response, is_declared_placeholder=is_declared_placeholder_response(response)) def _classify_existing_for_create(response: Any, table_id: list[str]) -> ResolvedNamespaceTable: diff --git a/daft_lance/utils.py b/daft_lance/utils.py index 882e1d0..fb4cebc 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -126,7 +126,9 @@ def construct_lance_dataset( io_derived_options = io_config_to_storage_options(io_config, resolved_uri) if io_config is not None else None if uri is not None: - base_options = storage_options if storage_options is not None else io_derived_options + # Falsy check on purpose: entry points historically treated an empty dict + # like None and fell through to the io_config-derived options. + base_options = storage_options or io_derived_options merged_storage_options = merge_storage_options(base_options, namespace_storage_options) else: merged_storage_options = merge_storage_options(io_derived_options, storage_options, namespace_storage_options) diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index b7e542a..b72e172 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -324,6 +324,80 @@ def declare_table(self, request: Any) -> Any: ) +def test_overwrite_resolves_declared_only_stub_on_strict_impl(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Strict impls 404 a plain describe on declared-only stubs; overwrite must still find them.""" + from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError + + class StrictNamespace: + def describe_table(self, request: Any) -> Any: + if getattr(request, "check_declared", None) is True: + return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None, is_only_declared=True) + raise TableNotFoundError("table not found: t") + + def declare_table(self, request: Any) -> Any: + raise TableAlreadyExistsError("Table already exists: t") + + monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: StrictNamespace()) + + resolved = namespace_mod.resolve_namespace_table( + namespace_impl="rest", + namespace_properties=None, + table_id=["t"], + mode="overwrite", + ) + assert resolved is not None + assert resolved.uri.endswith("t.lance") + assert resolved.is_declared_placeholder + + +def test_mem_wal_create_over_materialized_placeholder_rejected(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + import daft_lance.lance_data_sink as sink_mod + + ns = _dir_ns(tmp_path) + schema = daft.from_pydict({"id": [1]}).schema() + sink = daft_lance.LanceDataSink(None, schema, "create", table_id=["t"], use_mem_wal=True, **ns) + + monkeypatch.setattr( + sink_mod, + "resolve_namespace_table", + lambda **kwargs: namespace_mod.ResolvedNamespaceTable( + uri=str(tmp_path / "t.lance"), is_declared_placeholder=True + ), + ) + import pyarrow as pa + + fake_dataset = SimpleNamespace(schema=pa.schema([("id", pa.int64())]), latest_version=1) + monkeypatch.setattr("daft_lance.lance_data_sink.lance.dataset", lambda *args, **kwargs: fake_dataset) + + with pytest.raises(ValueError, match="use_mem_wal"): + sink.start() + + +def test_construct_lance_dataset_empty_storage_options_falls_back_to_io_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import daft_lance.utils as utils_mod + from daft.io import IOConfig, S3Config + + captured = {} + + class FakeDataset: + pass + + def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **kwargs: Any) -> Any: + captured["storage_options"] = storage_options + return FakeDataset() + + monkeypatch.setattr("daft_lance.utils.lance.dataset", fake_dataset) + + io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) + utils_mod.construct_lance_dataset("s3://bucket/t.lance", storage_options={}, io_config=io_config) + + merged = captured["storage_options"] + assert merged is not None + assert merged["access_key_id"] == "io-key" + + def test_declared_placeholder_response_fallback_property() -> None: assert namespace_mod.is_declared_placeholder_response(SimpleNamespace(is_only_declared=True)) assert namespace_mod.is_declared_placeholder_response( From 6fb5437401dcf0969bcfc9f48495275bed858df3 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 15:07:40 +0900 Subject: [PATCH 09/25] test(namespace): lock in sink pickle round-trip and compaction fragment count - pickle a started sink and run write()/finalize() on the copy, codifying the driver-resolves/worker-writes contract instead of relying on manual verification - compact_files test now asserts the fragment count actually shrinks, not just data correctness Claude-Session: https://claude.ai/code/session_01VyD31uRyBSEmKvPQhLtM5E --- tests/io/lancedb/test_namespace.py | 32 ++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index b72e172..82784d5 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -215,8 +215,17 @@ def test_namespace_compact_files(tmp_path: Path) -> None: for i in range(3): daft_lance.write_lance(daft.from_pydict({"id": [i + 2]}), table_id=table_id, mode="append", **ns).collect() + import lance + import lance_namespace as ln + from lance_namespace import DescribeTableRequest + + location = ln.connect("dir", {"root": str(tmp_path)}).describe_table(DescribeTableRequest(id=table_id)).location + fragments_before = len(lance.dataset(location).get_fragments()) + daft_lance.compact_files(table_id=table_id, compaction_options={"target_rows_per_fragment": 1024}, **ns) + assert len(lance.dataset(location).get_fragments()) < fragments_before + result = daft_lance.read_lance(table_id=table_id, **ns).sort("id").to_pydict() assert result == {"id": [1, 2, 3, 4]} @@ -398,6 +407,29 @@ def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **k assert merged["access_key_id"] == "io-key" +def test_sink_survives_pickle_after_start(tmp_path: Path) -> None: + """start() state must round-trip through pickle: workers run write() on a copy.""" + import pickle + + ns = _dir_ns(tmp_path) + schema = daft.from_pydict({"id": [1]}).schema() + + sink = daft_lance.LanceDataSink(None, schema, "create", table_id=["pickled"], **ns) + sink.start() + + worker_sink = pickle.loads(pickle.dumps(sink)) + assert worker_sink._table_uri == sink._table_uri + assert worker_sink._storage_options == sink._storage_options + assert worker_sink._effective_pyarrow_schema == sink._effective_pyarrow_schema + + from daft.recordbatch import MicroPartition + + results = list(worker_sink.write(iter([MicroPartition.from_pydict({"id": [1, 2]})]))) + stats = sink.finalize(results).to_pydict() + assert stats["version"] == [1] + assert daft_lance.read_lance(table_id=["pickled"], **ns).to_pydict() == {"id": [1, 2]} + + def test_declared_placeholder_response_fallback_property() -> None: assert namespace_mod.is_declared_placeholder_response(SimpleNamespace(is_only_declared=True)) assert namespace_mod.is_declared_placeholder_response( From 7833d582d205c1b5f0c90fd7c3ffef9d82adef30 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 15:32:43 +0900 Subject: [PATCH 10/25] fix(namespace): request vended credentials explicitly and unify empty storage_options handling - describe_table/declare_table now pass vend_credentials=True: per the lance-namespace spec, whether credentials are returned is implementation-defined when the flag is unset, so the documented "namespace vends storage_options" behavior was not guaranteed - LanceDataSink now treats storage_options={} the same as the read entry points (falsy fallthrough to io_config-derived options); previously write_lance dropped io_config credentials for that input while read_lance kept them Claude-Session: https://claude.ai/code/session_01VyD31uRyBSEmKvPQhLtM5E --- daft_lance/lance_data_sink.py | 7 ++--- daft_lance/namespace.py | 12 +++++---- tests/io/lancedb/test_namespace.py | 43 ++++++++++++++++++++++++++++++ 3 files changed, 54 insertions(+), 8 deletions(-) diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index 35f04d9..b0ea082 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -168,12 +168,13 @@ def _resolve_table(self) -> ResolvedNamespaceTable: def _merged_storage_options(self, resolved: ResolvedNamespaceTable) -> dict[str, str] | None: """Layer storage options: io_config-derived < user-provided < namespace-vended. - For a plain uri, user-provided options replace the io_config-derived ones - entirely (historical behavior). + For a plain uri, non-empty user-provided options replace the + io_config-derived ones entirely (falsy fallthrough, matching the read + entry points via construct_lance_dataset). """ io_derived = io_config_to_storage_options(self._io_config, resolved.uri) if self._uri is not None: - base = self._user_storage_options if self._user_storage_options is not None else io_derived + base = self._user_storage_options or io_derived return merge_storage_options(base, resolved.storage_options) return merge_storage_options(io_derived, self._user_storage_options, resolved.storage_options) diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index 11fda3f..08e755d 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -122,16 +122,18 @@ def _resolved_from_response(response: Any, *, is_declared_placeholder: bool = Fa def _describe_table(namespace: Any, table_id: list[str], *, check_declared: bool = False) -> Any: from lance_namespace import DescribeTableRequest - request = ( - DescribeTableRequest(id=table_id, check_declared=True) if check_declared else DescribeTableRequest(id=table_id) - ) - return namespace.describe_table(request) + # vend_credentials is explicit: when unset, whether the namespace returns + # storage credentials is implementation-defined. + kwargs: dict[str, Any] = {"id": table_id, "vend_credentials": True} + if check_declared: + kwargs["check_declared"] = True + return namespace.describe_table(DescribeTableRequest(**kwargs)) def _declare_table(namespace: Any, table_id: list[str]) -> ResolvedNamespaceTable: from lance_namespace import DeclareTableRequest - response = namespace.declare_table(DeclareTableRequest(id=table_id, location=None)) + response = namespace.declare_table(DeclareTableRequest(id=table_id, location=None, vend_credentials=True)) return _resolved_from_response(response, is_declared_placeholder=True) diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index 82784d5..cd65dc8 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -430,6 +430,49 @@ def test_sink_survives_pickle_after_start(tmp_path: Path) -> None: assert daft_lance.read_lance(table_id=["pickled"], **ns).to_pydict() == {"id": [1, 2]} +def test_namespace_requests_explicitly_vend_credentials(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Whether a namespace vends credentials is implementation-defined unless requested.""" + from lance_namespace.errors import TableNotFoundError + + requests: list[Any] = [] + + class RecordingNamespace: + def describe_table(self, request: Any) -> Any: + requests.append(request) + raise TableNotFoundError("table not found: t") + + def declare_table(self, request: Any) -> Any: + requests.append(request) + return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None) + + monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: RecordingNamespace()) + + for mode in ("create", "overwrite"): + namespace_mod.resolve_namespace_table( + namespace_impl="rest", namespace_properties=None, table_id=["t"], mode=mode + ) + with pytest.raises(TableNotFoundError): + namespace_mod.resolve_namespace_table( + namespace_impl="rest", namespace_properties=None, table_id=["t"], mode="read" + ) + + assert requests, "expected describe/declare requests to be issued" + assert all(request.vend_credentials is True for request in requests) + + +def test_sink_empty_storage_options_falls_back_to_io_config(tmp_path: Path) -> None: + """storage_options={} must behave like the read entry points: fall through to io_config.""" + from daft.io import IOConfig, S3Config + + io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) + schema = daft.from_pydict({"id": [1]}).schema() + + sink = daft_lance.LanceDataSink("s3://bucket/t.lance", schema, "create", io_config, storage_options={}) + merged = sink._merged_storage_options(namespace_mod.ResolvedNamespaceTable(uri="s3://bucket/t.lance")) + assert merged is not None + assert merged["access_key_id"] == "io-key" + + def test_declared_placeholder_response_fallback_property() -> None: assert namespace_mod.is_declared_placeholder_response(SimpleNamespace(is_only_declared=True)) assert namespace_mod.is_declared_placeholder_response( From 72cb1c79ad8222bcddcf2fd04cc7c3e158957ea6 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 15:52:25 +0900 Subject: [PATCH 11/25] docs(namespace): point rest:// uri error at the namespace parameters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The old message predated namespace support and directed users to "the daft-lance package" — which is where they already are. Show the actual read_lance(namespace_impl="rest", ...) usage instead. Claude-Session: https://claude.ai/code/session_01VyD31uRyBSEmKvPQhLtM5E --- daft_lance/_lance.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 3584b89..f23cdfb 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -145,10 +145,12 @@ def read_lance( uri_str = str(uri) if uri is not None else None if uri_str is not None and uri_str.startswith("rest://"): raise ValueError( - "rest:// Lance URIs are no longer supported by daft.read_lance. " - "The previous REST-namespace integration did not match the real " - "lance-namespace API and has been removed. Use a local path or " - "object-store URI, or the daft-lance package." + "rest:// Lance URIs are not supported. To read a table through a REST " + "Lance Namespace, use the namespace parameters instead of a uri: " + 'read_lance(namespace_impl="rest", ' + 'namespace_properties={"uri": "http://host:port"}, ' + 'table_id=["catalog", "schema", "table"]). ' + "Otherwise pass a local path or object-store URI." ) io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config From 8627210ae4ea25f34f2a9f6d209043d4f6623c5d Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 16:26:06 +0900 Subject: [PATCH 12/25] docs(namespace): explain the integration layer's design for reviewers Module docstring covers the three load-bearing decisions (serialize the triple + per-process client cache, describe-first resolution, defensive error classification), and resolve_namespace_table documents the full create/overwrite/read state machine including the declare-race legs. Claude-Session: https://claude.ai/code/session_01VyD31uRyBSEmKvPQhLtM5E --- daft_lance/namespace.py | 62 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index 08e755d..5516768 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -1,3 +1,30 @@ +"""Lance Namespace (catalog) integration layer. + +Tables can be addressed by a ``(namespace_impl, namespace_properties, table_id)`` +triple instead of a raw uri; this module owns everything between that triple and +a pylance call. Three design decisions shape the code: + +1. **Only the triple is serialized, never the client.** Namespace clients are + not picklable, so distributed tasks carry the triple and every process + rebuilds its client through a per-process ``lru_cache`` + (:func:`get_or_create_namespace`). This is why entry points thread the three + raw parameters around instead of a client object. + +2. **Table resolution is describe-first** (:func:`resolve_namespace_table`). + The protocol only offers ``declare_table`` (conflicts on ANY existing table, + even a declared-only stub) and ``describe_table`` (404s on declared-only + stubs unless ``check_declared=True``), so the create/overwrite semantics an + engine needs — "reuse a declared-only stub, reject a real table, survive a + concurrent declare" — are assembled here from those two primitives. + +3. **Error classification is defensive** (:func:`is_table_not_found` / + :func:`is_table_already_exists`): native namespace impls raise typed errors, + but the Rust-backed REST client can surface untyped ones, so both fall back + to strict status-code + message matching. Never branch on a bare + ``except Exception`` around namespace calls — an auth or network failure + must not be mistaken for "table does not exist". +""" + from __future__ import annotations import os @@ -12,6 +39,7 @@ def has_namespace_params(namespace_impl: str | None, table_id: list[str] | None) -> bool: + """Whether namespace addressing is in effect (``namespace_properties`` stays optional).""" return namespace_impl is not None and table_id is not None @@ -21,6 +49,7 @@ def validate_uri_or_namespace( table_id: list[str] | None, namespace_properties: dict[str, str] | None = None, ) -> None: + """Enforce that exactly one addressing style is used: ``uri`` XOR the namespace triple.""" has_uri = uri is not None has_ns = has_namespace_params(namespace_impl, table_id) @@ -39,6 +68,13 @@ def validate_uri_or_namespace( def _normalize_file_uri(location: str) -> str: + """Strip a ``file://`` scheme to a plain filesystem path. + + Namespace impls return locations as URIs (dir namespace vends + ``file:///...``), but the location is later used where a plain path is + required: ``write_fragments(dataset_uri=...)``, ``LanceDataset.commit``, + and ``pathlib`` checks in the sink. Object-store URIs pass through as-is. + """ parsed = urlparse(location) if parsed.scheme == "file": return parsed.path @@ -54,6 +90,14 @@ def _get_cached_namespace(namespace_impl: str, namespace_properties_tuple: tuple def get_or_create_namespace(namespace_impl: str | None, namespace_properties: dict[str, str] | None) -> Any | None: + """Per-process namespace client pool. + + ``ln.connect`` may build HTTP clients / perform auth, and workers re-derive + the client for every scan task and fragment write, so connections are cached + per (impl, properties). Properties are canonicalized to a sorted tuple + because ``lru_cache`` keys must be hashable and dict ordering must not + create duplicate connections. + """ if namespace_impl is None: return None namespace_properties_tuple = tuple(sorted(namespace_properties.items())) if namespace_properties else None @@ -240,6 +284,19 @@ def resolve_namespace_table( reused as placeholders); ``mode="overwrite"`` declares the table when it does not exist yet; other modes require the table to exist. Only ``mode="create"`` and ``mode="overwrite"`` have side effects (a ``declare_table`` call). + + Resolution state machine (describe always with ``check_declared=True`` except + in read mode):: + + create: describe ─ found ──► real table? error : reuse stub (placeholder) + └ 404 ───► declare ─ conflict? re-describe & classify + overwrite: describe ─ found ──► use it (stub or real) + └ 404 ───► declare ─ conflict? re-describe & use winner + read: describe ─ found ──► use it + └ error ─► raise + + The "declare ─ conflict?" legs exist because a concurrent writer can win the + declare race between our describe and our declare. """ namespace = get_or_create_namespace(namespace_impl, namespace_properties) if namespace is None or table_id is None: @@ -279,6 +336,11 @@ def _resolve_for_create(namespace: Any, table_id: list[str]) -> ResolvedNamespac def _declare_or_recover(namespace: Any, table_id: list[str], *, for_create: bool) -> ResolvedNamespaceTable: + """Declare the table; if a concurrent declare won the race, classify the winner. + + ``for_create=True`` keeps create semantics (a real-table winner is an error); + ``for_create=False`` (overwrite) accepts whatever won. + """ try: return _declare_table(namespace, table_id) except Exception as error: From 3bf594381e3a1052b4fd91f172869c787ed2d755 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 18:51:37 +0900 Subject: [PATCH 13/25] refactor(namespace): drop patch_daft in favor of daft_lance entry points MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit patch_daft() was introduced on this branch and never released, so it can be removed without a deprecation cycle. Until Eventual-Inc/Daft#7282 delegates DataFrame.write_lance / read_lance natively, namespace users call daft_lance.read_lance / daft_lance.write_lance directly — which the README documents as the primary usage anyway. This also removes the only monkeypatch in the package and avoids having to keep the patch compatible with the upstream signature change coming in #7282. Claude-Session: https://claude.ai/code/session_01VyD31uRyBSEmKvPQhLtM5E --- README.md | 15 ++------ daft_lance/__init__.py | 42 --------------------- tests/io/lancedb/test_namespace.py | 17 --------- tests/io/lancedb/test_namespace_rest_e2e.py | 14 ------- 4 files changed, 4 insertions(+), 84 deletions(-) diff --git a/README.md b/README.md index 0d9da07..106d6bd 100644 --- a/README.md +++ b/README.md @@ -89,18 +89,11 @@ Namespace clients are cached per (implementation, properties) pair. The cache si to 16 and can be tuned with the `DAFT_LANCE_NAMESPACE_CACHE_SIZE` environment variable (read once at import time). -#### Drop-in Daft APIs +#### Daft's own entry points -If you prefer Daft's own entry points, call `daft_lance.patch_daft()` once to route -`daft.read_lance` and `DataFrame.write_lance` through daft-lance (namespace parameters included): - -```python -import daft, daft_lance - -daft_lance.patch_daft() -df.write_lance(table_id=["my_table"], mode="create", **namespace).collect() -daft.read_lance(table_id=["my_table"], **namespace).show() -``` +Native namespace support in `daft.read_lance` / `DataFrame.write_lance` is tracked in +[Eventual-Inc/Daft#7282](https://github.com/Eventual-Inc/Daft/issues/7282); until that lands, +use the `daft_lance` entry points shown above for namespace-addressed tables. ## Migration diff --git a/daft_lance/__init__.py b/daft_lance/__init__.py index 5af90b5..322dbc4 100644 --- a/daft_lance/__init__.py +++ b/daft_lance/__init__.py @@ -1,7 +1,5 @@ from __future__ import annotations -from typing import Any, Literal - try: import daft except ImportError: @@ -18,52 +16,12 @@ ) from .lance_data_sink import LanceDataSink - -def patch_daft() -> None: - """Route ``daft.read_lance`` and ``DataFrame.write_lance`` through daft-lance. - - Optional convenience for drop-in use of namespace parameters with Daft's own - APIs, e.g. ``df.write_lance(namespace_impl="dir", ...)``. Idempotent. - """ - daft.read_lance = read_lance # type: ignore[assignment] - - from daft.dataframe import DataFrame - - if getattr(DataFrame.write_lance, "_daft_lance_patch", False): - return - - def _write_lance( - self: DataFrame, - uri: Any = None, - mode: Literal["create", "append", "overwrite", "merge"] = "create", - io_config: Any | None = None, - schema: Any | None = None, - left_on: str | None = None, - right_on: str | None = None, - **kwargs: Any, - ) -> DataFrame: - return write_lance( - self, - uri, - mode=mode, - io_config=io_config, - schema=schema, - left_on=left_on, - right_on=right_on, - **kwargs, - ) - - _write_lance._daft_lance_patch = True # type: ignore[attr-defined] - DataFrame.write_lance = _write_lance # type: ignore[method-assign] - - __all__ = [ "LanceDataSink", "compact_files", "create_scalar_index", "merge_columns", "merge_columns_df", - "patch_daft", "read_lance", "take_blobs", "write_lance", diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index cd65dc8..709f3d9 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -119,23 +119,6 @@ def test_namespace_count_pushdown(tmp_path: Path) -> None: assert daft_lance.read_lance(table_id=table_id, **ns).count_rows() == 10 -def test_namespace_patch_daft(tmp_path: Path) -> None: - ns = _dir_ns(tmp_path) - table_id = ["patched"] - - daft_lance.patch_daft() - daft_lance.patch_daft() # idempotent - - daft.from_pydict({"id": [1, 2]}).write_lance(table_id=table_id, mode="create", **ns).collect() - result = daft.read_lance(table_id=table_id, **ns).to_pydict() # type: ignore[call-arg] - assert result == {"id": [1, 2]} - - # Plain-URI writes keep working through the patched method. - uri = str(tmp_path / "plain.lance") - daft.from_pydict({"id": [5]}).write_lance(uri).collect() - assert daft.read_lance(uri).to_pydict() == {"id": [5]} - - def test_namespace_write_lance_merge_new_columns(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["merge_tbl"] diff --git a/tests/io/lancedb/test_namespace_rest_e2e.py b/tests/io/lancedb/test_namespace_rest_e2e.py index 3e388e0..ce3b089 100644 --- a/tests/io/lancedb/test_namespace_rest_e2e.py +++ b/tests/io/lancedb/test_namespace_rest_e2e.py @@ -63,17 +63,3 @@ def test_rest_namespace_write_read_append_roundtrip() -> None: assert daft_lance.read_lance(table_id=table_id, **ns).count_rows() == 5 assert lance.dataset(None, namespace_client=namespace, table_id=table_id).count_rows() == 5 - - -def test_rest_namespace_patch_daft_roundtrip() -> None: - """The opt-in ``patch_daft()`` path: ``daft.read_lance`` / ``df.write_lance``.""" - namespace_properties = {"uri": os.environ["DAFT_LANCE_REST_URI"]} - catalog = os.environ.get("DAFT_LANCE_REST_CATALOG", "lance_catalog") - schema = os.environ.get("DAFT_LANCE_REST_SCHEMA", "daft_ns_e2e") - table_id = [catalog, schema, f"patched_{uuid.uuid4().hex[:8]}"] - ns: dict[str, Any] = {"namespace_impl": "rest", "namespace_properties": namespace_properties} - - daft_lance.patch_daft() - - daft.from_pydict({"id": [1, 2]}).write_lance(table_id=table_id, mode="create", **ns).collect() - assert daft.read_lance(table_id=table_id, **ns).to_pydict() == {"id": [1, 2]} # type: ignore[call-arg] From fc7cda3b2d466df3156ed59f0913adc472b5493b Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 19:13:33 +0900 Subject: [PATCH 14/25] style: format scalar index namespace reload --- daft_lance/lance_scalar_index.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index ef591cc..2cd2fa2 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -388,9 +388,7 @@ def _create_segmented_index( # Reload dataset to pick up the latest version (segment files were written # by workers against the version that was current at their invocation time). namespace_kwargs = namespace_kwargs_for_dataset(lance_ds) - lance_ds = lance.dataset( - None if namespace_kwargs else uri, storage_options=storage_options, **namespace_kwargs - ) + lance_ds = lance.dataset(None if namespace_kwargs else uri, storage_options=storage_options, **namespace_kwargs) index_metas = _prepare_index_segments_for_commit(lance_ds, index_type, index_metas) logger.info( From 6f3bd945182d8ab36887764ac05725488c43911a Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Thu, 16 Jul 2026 22:44:03 +0900 Subject: [PATCH 15/25] chore: drop unused __future__ import from package init Claude-Session: https://claude.ai/code/session_01VyD31uRyBSEmKvPQhLtM5E --- daft_lance/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/daft_lance/__init__.py b/daft_lance/__init__.py index 322dbc4..051934d 100644 --- a/daft_lance/__init__.py +++ b/daft_lance/__init__.py @@ -1,5 +1,3 @@ -from __future__ import annotations - try: import daft except ImportError: From ea3293d4f742212084f583e7a4af35f91bd7692c Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Mon, 20 Jul 2026 18:31:59 +0900 Subject: [PATCH 16/25] refactor(namespace): remove legacy and unrelated compatibility paths --- daft_lance/__init__.py | 2 - daft_lance/_lance.py | 126 ++---------- daft_lance/lance_data_sink.py | 24 +-- daft_lance/namespace.py | 210 ++------------------ tests/io/lancedb/conftest.py | 33 --- tests/io/lancedb/test_namespace.py | 176 ++++------------ tests/io/lancedb/test_namespace_rest_e2e.py | 35 ++++ 7 files changed, 110 insertions(+), 496 deletions(-) diff --git a/daft_lance/__init__.py b/daft_lance/__init__.py index 051934d..1a0323d 100644 --- a/daft_lance/__init__.py +++ b/daft_lance/__init__.py @@ -12,10 +12,8 @@ read_lance, write_lance, ) -from .lance_data_sink import LanceDataSink __all__ = [ - "LanceDataSink", "compact_files", "create_scalar_index", "merge_columns", diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index f23cdfb..4e2d321 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -16,7 +16,7 @@ from .lance_merge_column import merge_columns_from_df, merge_columns_internal from .lance_scalar_index import create_scalar_index_internal from .lance_scan import LanceDBScanOperator -from .namespace import is_table_not_found, validate_uri_or_namespace +from .namespace import validate_uri_or_namespace from .utils import construct_lance_dataset if TYPE_CHECKING: @@ -638,43 +638,13 @@ def compact_files( ) -def _is_missing_dataset_error(error: Exception) -> bool: - """Whether opening a dataset failed because it does not exist yet.""" - if isinstance(error, (FileNotFoundError,)): - return True - if isinstance(error, (ValueError, OSError)) and "was not found" in str(error): - return True - return is_table_not_found(error) - - -def _stats_dataframe(lance_ds: LanceDataset) -> DataFrame: - """Build the same stats DataFrame that LanceDataSink.finalize returns.""" - import pyarrow as _pa - - from daft.recordbatch import MicroPartition - - stats = lance_ds.stats.dataset_stats() - return DataFrame._from_micropartitions( - MicroPartition.from_pydict( - { - "num_fragments": _pa.array([stats["num_fragments"]], type=_pa.int64()), - "num_deleted_rows": _pa.array([stats["num_deleted_rows"]], type=_pa.int64()), - "num_small_files": _pa.array([stats["num_small_files"]], type=_pa.int64()), - "version": _pa.array([lance_ds.version], type=_pa.int64()), - } - ) - ) - - @PublicAPI def write_lance( df: DataFrame, uri: str | pathlib.Path | None = None, - mode: Literal["create", "append", "overwrite", "merge"] = "create", + mode: Literal["create", "append", "overwrite"] = "create", io_config: IOConfig | None = None, schema: Any | None = None, - left_on: str | None = None, - right_on: str | None = None, *, table_id: list[str] | None = None, namespace_impl: str | None = None, @@ -686,11 +656,9 @@ def write_lance( Args: df: The DataFrame to write. uri: The URI of the Lance table. Mutually exclusive with the namespace parameters. - mode: One of "create", "append", "overwrite", or "merge" (add new columns to - existing rows; requires ``fragment_id`` and a join key column in ``df``). + mode: One of "create", "append", or "overwrite". io_config: A custom IOConfig to use when accessing Lance data. schema: Desired schema to enforce during write; defaults to the DataFrame schema. - left_on/right_on: Join keys for ``mode="merge"``; default "_rowaddr". table_id: Table identifier within the namespace, e.g. ["catalog", "schema", "table"]. namespace_impl: Lance Namespace implementation, e.g. "dir" or "rest". namespace_properties: Properties for connecting to the namespace, e.g. @@ -712,82 +680,14 @@ def write_lance( if schema is None: schema = df.schema() - if mode != "merge": - sanitized_kwargs = {k: v for k, v in kwargs.items() if k not in ("left_on", "right_on")} - sink = LanceDataSink( - uri, - schema, - mode, - io_config, - table_id=table_id, - namespace_impl=namespace_impl, - namespace_properties=namespace_properties, - **sanitized_kwargs, - ) - return df.write_sink(sink) - - io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - storage_options = kwargs.pop("storage_options", None) - - lance_ds: LanceDataset | None - try: - lance_ds = construct_lance_dataset( - uri, - storage_options=storage_options, - io_config=io_config, - namespace_impl=namespace_impl, - namespace_properties=namespace_properties, - table_id=table_id, - ) - except Exception as error: - if not _is_missing_dataset_error(error): - raise - lance_ds = None - - sanitized_kwargs = {k: v for k, v in kwargs.items() if k not in ("left_on", "right_on")} - - if lance_ds is None: - # Merge against a missing table degenerates to create. - sink = LanceDataSink( - uri, - schema, - "create", - io_config, - table_id=table_id, - namespace_impl=namespace_impl, - namespace_properties=namespace_properties, - storage_options=storage_options, - **sanitized_kwargs, - ) - return df.write_sink(sink) - - existing_fields = {getattr(f, "name", str(f)) for f in lance_ds.schema} - meta_exclusions = {"fragment_id", "_rowaddr", "_rowid"} - new_cols = [c for c in df.column_names if c not in existing_fields and c not in meta_exclusions] - - if len(new_cols) == 0: - # No schema evolution: merge degenerates to append. - sink = LanceDataSink( - uri, - schema, - "append", - io_config, - table_id=table_id, - namespace_impl=namespace_impl, - namespace_properties=namespace_properties, - storage_options=storage_options, - **sanitized_kwargs, - ) - return df.write_sink(sink) - - join_left = left_on or "_rowaddr" - join_right = right_on or join_left - merged = merge_columns_from_df( - df, - lance_ds=lance_ds, - uri=_effective_uri(lance_ds, uri), - left_on=join_left, - right_on=join_right, - storage_options=_effective_storage_options(lance_ds, storage_options), + sink = LanceDataSink( + uri, + schema, + mode, + io_config, + table_id=table_id, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + **kwargs, ) - return _stats_dataframe(merged) + return df.write_sink(sink) diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index b0ea082..bda33d1 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -103,8 +103,6 @@ def __init__( self._effective_pyarrow_schema: pa.Schema | None = None self._version: int = 0 self._table_schema: pa.Schema | None = None - self._declared_placeholder = False - self._overwrite_placeholder_dataset = False self._schema = Schema._from_field_name_and_types( [ @@ -124,7 +122,6 @@ def start(self) -> None: """ resolved = self._resolve_table() self._table_uri = resolved.uri - self._declared_placeholder = resolved.is_declared_placeholder self._storage_options = self._merged_storage_options(resolved) existing = self._absorb_existing_dataset() @@ -261,20 +258,10 @@ def _absorb_existing_dataset(self) -> lance.LanceDataset | None: self._version = dataset.latest_version if self._mode == "create": - if not self._declared_placeholder: - raise ValueError( - "Cannot create a Lance dataset at a location where one already exists. " - 'Use mode="overwrite" to replace it or mode="append" to add to it.' - ) - # The namespace declared a stub table and this location already holds - # a (placeholder) dataset; overwrite it instead of failing the create. - if self._use_mem_wal: - raise ValueError( - "use_mem_wal=True cannot create over a namespace-declared placeholder " - 'dataset; write with use_mem_wal=False or use mode="overwrite".' - ) - self._overwrite_placeholder_dataset = True - self._table_schema = None + raise ValueError( + "Cannot create a Lance dataset at a location where one already exists. " + 'Use mode="overwrite" to replace it or mode="append" to add to it.' + ) if self._mode == "append" and not _pyarrow_schema_castable( blob_aware_schema_for_validation(self._pyarrow_schema, table_schema), @@ -305,11 +292,10 @@ def _prepare_arrow_table(self, input_table: pa.Table) -> pa.Table: def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetadata]]: assert self._table_uri is not None, "LanceDataSink.start() must run before writes" wrapped = self._blob.wrap_table(table) - write_mode = "overwrite" if self._overwrite_placeholder_dataset else self._mode fragments = lance.fragment.write_fragments( wrapped, dataset_uri=self._table_uri, - mode=write_mode, + mode=self._mode, storage_options=self._storage_options, max_rows_per_file=self._max_rows_per_file, max_rows_per_group=self._max_rows_per_group, diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index 5516768..fd6b2f4 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -10,19 +10,10 @@ (:func:`get_or_create_namespace`). This is why entry points thread the three raw parameters around instead of a client object. -2. **Table resolution is describe-first** (:func:`resolve_namespace_table`). - The protocol only offers ``declare_table`` (conflicts on ANY existing table, - even a declared-only stub) and ``describe_table`` (404s on declared-only - stubs unless ``check_declared=True``), so the create/overwrite semantics an - engine needs — "reuse a declared-only stub, reject a real table, survive a - concurrent declare" — are assembled here from those two primitives. - -3. **Error classification is defensive** (:func:`is_table_not_found` / - :func:`is_table_already_exists`): native namespace impls raise typed errors, - but the Rust-backed REST client can surface untyped ones, so both fall back - to strict status-code + message matching. Never branch on a bare - ``except Exception`` around namespace calls — an auth or network failure - must not be mistaken for "table does not exist". +2. **Table creation delegates to the namespace.** ``mode="create"`` maps to + the atomic ``declare_table`` operation, while overwrite describes first and + declares only when the namespace raises its typed ``TableNotFoundError``. + Other namespace failures propagate unchanged. """ from __future__ import annotations @@ -142,133 +133,32 @@ def _response_location(response: Any) -> str: @dataclass(frozen=True) class ResolvedNamespaceTable: - """A namespace table resolved to a physical location. - - ``is_declared_placeholder`` is True when the location comes from a - ``declare_table`` call (fresh or pre-existing declared-only stub): some - namespace implementations materialize a stub dataset at declare time, and a - ``mode="create"`` write must overwrite that stub instead of failing. - """ + """A namespace table resolved to a physical location and storage options.""" uri: str storage_options: dict[str, str] | None = None - is_declared_placeholder: bool = False -def _resolved_from_response(response: Any, *, is_declared_placeholder: bool = False) -> ResolvedNamespaceTable: +def _resolved_from_response(response: Any) -> ResolvedNamespaceTable: return ResolvedNamespaceTable( uri=_response_location(response), storage_options=_storage_options(response), - is_declared_placeholder=is_declared_placeholder, ) -def _describe_table(namespace: Any, table_id: list[str], *, check_declared: bool = False) -> Any: +def _describe_table(namespace: Any, table_id: list[str]) -> Any: from lance_namespace import DescribeTableRequest # vend_credentials is explicit: when unset, whether the namespace returns # storage credentials is implementation-defined. - kwargs: dict[str, Any] = {"id": table_id, "vend_credentials": True} - if check_declared: - kwargs["check_declared"] = True - return namespace.describe_table(DescribeTableRequest(**kwargs)) + return namespace.describe_table(DescribeTableRequest(id=table_id, vend_credentials=True)) def _declare_table(namespace: Any, table_id: list[str]) -> ResolvedNamespaceTable: from lance_namespace import DeclareTableRequest response = namespace.declare_table(DeclareTableRequest(id=table_id, location=None, vend_credentials=True)) - return _resolved_from_response(response, is_declared_placeholder=True) - - -def is_declared_placeholder_response(response: Any) -> bool: - """Whether a ``describe_table`` response describes a declared-only stub table. - - ``is_only_declared`` is the standard signal (only reliable when the request - was made with ``check_declared=True``); the ``lance.declared`` property is a - compatibility fallback for services that expose the state that way. - """ - if getattr(response, "is_only_declared", None) is True: - return True - for attr in ("properties", "metadata"): - mapping = getattr(response, attr, None) - if mapping and str(dict(mapping).get("lance.declared", "")).strip().lower() == "true": - return True - return False - - -def is_table_not_found(error: Exception) -> bool: - """Whether an exception from ``describe_table`` means the table does not exist. - - Native implementations raise the typed ``TableNotFoundError``; the Rust-backed - REST client may surface untyped errors, so only fall back when the error is - clearly a 404 for a table resource. - """ - try: - from lance_namespace.errors import TableNotFoundError - - if isinstance(error, TableNotFoundError): - return True - except ImportError: - pass - - status_code = _error_status_code(error) - if status_code != 404: - return False - - message = str(error).lower() - return "no such table" in message or ( - "table" in message and ("not found" in message or "does not exist" in message) - ) - - -def is_table_already_exists(error: Exception) -> bool: - """Whether an exception from ``declare_table`` means the table already exists. - - Mirrors :func:`is_table_not_found`: prefer the typed error, fall back to a - clearly-classified conflict from untyped REST errors. - """ - try: - from lance_namespace.errors import TableAlreadyExistsError - - if isinstance(error, TableAlreadyExistsError): - return True - except ImportError: - pass - - status_code = _error_status_code(error) - if status_code != 409: - return False - - message = str(error).lower() - return "table" in message and "already exists" in message - - -def _error_status_code(error: Exception) -> int | None: - """Best-effort HTTP status extraction for untyped namespace REST errors.""" - for attr in ("status", "status_code", "code"): - value = getattr(error, attr, None) - status_code = _coerce_status_code(value) - if status_code is not None: - return status_code - - response = getattr(error, "response", None) - if response is not None: - for attr in ("status", "status_code"): - value = getattr(response, attr, None) - status_code = _coerce_status_code(value) - if status_code is not None: - return status_code - - return None - - -def _coerce_status_code(value: Any) -> int | None: - if isinstance(value, int): - return value - if isinstance(value, str) and value.isdigit(): - return int(value) - return None + return _resolved_from_response(response) def resolve_namespace_table( @@ -280,87 +170,25 @@ def resolve_namespace_table( ) -> ResolvedNamespaceTable | None: """Resolve a namespace table to a :class:`ResolvedNamespaceTable`. - ``mode="create"`` requires the table to not exist (declared-only stubs are - reused as placeholders); ``mode="overwrite"`` declares the table when it does - not exist yet; other modes require the table to exist. Only ``mode="create"`` - and ``mode="overwrite"`` have side effects (a ``declare_table`` call). - - Resolution state machine (describe always with ``check_declared=True`` except - in read mode):: - - create: describe ─ found ──► real table? error : reuse stub (placeholder) - └ 404 ───► declare ─ conflict? re-describe & classify - overwrite: describe ─ found ──► use it (stub or real) - └ 404 ───► declare ─ conflict? re-describe & use winner - read: describe ─ found ──► use it - └ error ─► raise - - The "declare ─ conflict?" legs exist because a concurrent writer can win the - declare race between our describe and our declare. + ``mode="create"`` atomically declares a new table. ``mode="overwrite"`` + declares the table only when a typed ``TableNotFoundError`` is raised; + other modes require the table to exist. """ namespace = get_or_create_namespace(namespace_impl, namespace_properties) if namespace is None or table_id is None: return None if mode == "create": - return _resolve_for_create(namespace, table_id) - - if mode == "overwrite": - # check_declared: a plain describe may 404 on declared-only stubs, and a - # stub is a perfectly fine overwrite target. - try: - response = _describe_table(namespace, table_id, check_declared=True) - except Exception as error: - if not is_table_not_found(error): - raise - return _declare_or_recover(namespace, table_id, for_create=False) - return _resolved_from_response(response, is_declared_placeholder=is_declared_placeholder_response(response)) - - return _resolved_from_response(_describe_table(namespace, table_id)) - - -def _resolve_for_create(namespace: Any, table_id: list[str]) -> ResolvedNamespaceTable: - """Describe-first create resolution. - - ``declare_table`` fails with a conflict on any existing table (declared-only - or real), so classify via ``describe_table`` before declaring, and again if - a concurrent declare wins the race in between. - """ - try: - response = _describe_table(namespace, table_id, check_declared=True) - except Exception as error: - if not is_table_not_found(error): - raise - return _declare_or_recover(namespace, table_id, for_create=True) - return _classify_existing_for_create(response, table_id) - + return _declare_table(namespace, table_id) -def _declare_or_recover(namespace: Any, table_id: list[str], *, for_create: bool) -> ResolvedNamespaceTable: - """Declare the table; if a concurrent declare won the race, classify the winner. + from lance_namespace.errors import TableNotFoundError - ``for_create=True`` keeps create semantics (a real-table winner is an error); - ``for_create=False`` (overwrite) accepts whatever won. - """ try: - return _declare_table(namespace, table_id) - except Exception as error: - if not is_table_already_exists(error): - raise - # Lost a declare race; re-describe (check_declared so a declared-only - # winner does not 404 on strict impls) and classify what won. - response = _describe_table(namespace, table_id, check_declared=True) - if for_create: - return _classify_existing_for_create(response, table_id) - return _resolved_from_response(response, is_declared_placeholder=is_declared_placeholder_response(response)) - - -def _classify_existing_for_create(response: Any, table_id: list[str]) -> ResolvedNamespaceTable: - if is_declared_placeholder_response(response): - return _resolved_from_response(response, is_declared_placeholder=True) - raise ValueError( - f"Table {table_id!r} already exists in the namespace and cannot be created. " - 'Use mode="overwrite" to replace it or mode="append" to add to it.' - ) + return _resolved_from_response(_describe_table(namespace, table_id)) + except TableNotFoundError: + if mode == "overwrite": + return _declare_table(namespace, table_id) + raise def merge_storage_options(*layers: dict[str, Any] | None) -> dict[str, Any] | None: diff --git a/tests/io/lancedb/conftest.py b/tests/io/lancedb/conftest.py index 69feee1..dda9665 100644 --- a/tests/io/lancedb/conftest.py +++ b/tests/io/lancedb/conftest.py @@ -1,39 +1,6 @@ from __future__ import annotations -import os -import sys - import pytest # Try to import lance; if it fails, all tests in this directory will be skipped. lance = pytest.importorskip("lance") - -_NATIVE_TEARDOWN_CRASH_MARKER = "lance_native_teardown_crash_workaround" -_exitstatus: int | None = None -_hard_exit_after_session = False - - -def pytest_configure(config: pytest.Config) -> None: - config.addinivalue_line( - "markers", - f"{_NATIVE_TEARDOWN_CRASH_MARKER}: hard-exit after reported results to avoid a known Lance native teardown crash", - ) - - -def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: - global _hard_exit_after_session - _hard_exit_after_session = any(item.get_closest_marker(_NATIVE_TEARDOWN_CRASH_MARKER) for item in items) - - -def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: - global _exitstatus - _exitstatus = exitstatus - - -@pytest.hookimpl(trylast=True) -def pytest_unconfigure(config: pytest.Config) -> None: - if not _hard_exit_after_session or _exitstatus is None: - return - sys.stdout.flush() - sys.stderr.flush() - os._exit(_exitstatus) diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index 709f3d9..11c0109 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -9,8 +9,7 @@ import daft import daft_lance import daft_lance.namespace as namespace_mod - -pytestmark = pytest.mark.lance_native_teardown_crash_workaround +from daft_lance.lance_data_sink import LanceDataSink def _dir_ns(tmp_path: Path) -> dict[str, Any]: @@ -80,13 +79,6 @@ def declare_table(self, request: Any) -> Any: assert not namespace.declared -def test_namespace_untyped_404_table_not_found_is_missing() -> None: - class RestTableNotFound(RuntimeError): - status = 404 - - assert namespace_mod.is_table_not_found(RestTableNotFound("table catalog.schema.table not found")) - - def test_namespace_read_supports_pushdowns(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["pushdowns"] @@ -119,30 +111,6 @@ def test_namespace_count_pushdown(tmp_path: Path) -> None: assert daft_lance.read_lance(table_id=table_id, **ns).count_rows() == 10 -def test_namespace_write_lance_merge_new_columns(tmp_path: Path) -> None: - ns = _dir_ns(tmp_path) - table_id = ["merge_tbl"] - - daft_lance.write_lance( - daft.from_pydict({"id": [1, 2, 3], "score": [10, 20, 30]}), - table_id=table_id, - mode="create", - **ns, - ).collect() - - df = daft_lance.read_lance( - table_id=table_id, - default_scan_options={"with_row_address": True}, - include_fragment_id=True, - **ns, - ) - df = df.with_column("doubled", df["score"] * 2) - daft_lance.write_lance(df, table_id=table_id, mode="merge", **ns).collect() - - result = daft_lance.read_lance(table_id=table_id, **ns).sort("id").to_pydict() - assert result["doubled"] == [20, 40, 60] - - def test_namespace_merge_columns_df(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["merge_cols_df"] @@ -187,7 +155,7 @@ def test_namespace_create_scalar_index(tmp_path: Path) -> None: namespace = ln.connect("dir", {"root": str(tmp_path)}) location = namespace.describe_table(DescribeTableRequest(id=table_id)).location indices = lance.dataset(location).list_indices() - assert any(idx["fields"] == ["price"] for idx in indices) # type: ignore[index] + assert any(idx["fields"] == ["price"] for idx in indices) def test_namespace_compact_files(tmp_path: Path) -> None: @@ -217,7 +185,7 @@ def test_sink_construction_is_side_effect_free(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) schema = daft.from_pydict({"id": [1]}).schema() - sink = daft_lance.LanceDataSink(None, schema, "create", table_id=["deferred"], **ns) + sink = LanceDataSink(None, schema, "create", table_id=["deferred"], **ns) assert not (tmp_path / "deferred.lance").exists() sink.start() @@ -229,24 +197,27 @@ def test_sink_invalid_params_do_not_declare_table(tmp_path: Path) -> None: schema = daft.from_pydict({"id": [1]}).schema() with pytest.raises(ValueError, match="blob_columns"): - daft_lance.LanceDataSink(None, schema, "create", table_id=["orphan"], blob_columns=["missing"], **ns) + LanceDataSink(None, schema, "create", table_id=["orphan"], blob_columns=["missing"], **ns) assert not (tmp_path / "orphan.lance").exists() def test_namespace_create_on_existing_table_raises(tmp_path: Path) -> None: + from lance_namespace.errors import TableAlreadyExistsError + ns = _dir_ns(tmp_path) table_id = ["exists"] daft_lance.write_lance(daft.from_pydict({"id": [1]}), table_id=table_id, mode="create", **ns).collect() - with pytest.raises(ValueError, match="already exists"): + with pytest.raises(TableAlreadyExistsError, match="already exists"): daft_lance.write_lance(daft.from_pydict({"id": [2]}), table_id=table_id, mode="create", **ns).collect() -def test_namespace_create_over_declared_placeholder(tmp_path: Path) -> None: +def test_namespace_create_on_declared_table_raises(tmp_path: Path) -> None: import lance_namespace as ln from lance_namespace import DeclareTableRequest + from lance_namespace.errors import TableAlreadyExistsError ns = _dir_ns(tmp_path) table_id = ["declared_first"] @@ -254,115 +225,54 @@ def test_namespace_create_over_declared_placeholder(tmp_path: Path) -> None: namespace = ln.connect("dir", {"root": str(tmp_path)}) namespace.declare_table(DeclareTableRequest(id=table_id, location=None)) - daft_lance.write_lance(daft.from_pydict({"id": [1, 2]}), table_id=table_id, mode="create", **ns).collect() - - assert daft_lance.read_lance(table_id=table_id, **ns).to_pydict() == {"id": [1, 2]} - + with pytest.raises(TableAlreadyExistsError, match="already exists"): + daft_lance.write_lance(daft.from_pydict({"id": [1, 2]}), table_id=table_id, mode="create", **ns).collect() -def test_create_declare_race_recovers_placeholder(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError - class RacingNamespace: - def __init__(self) -> None: - self.describe_calls = 0 +def test_namespace_create_declares_without_describe(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + requests: list[Any] = [] + class RecordingNamespace: def describe_table(self, request: Any) -> Any: - self.describe_calls += 1 - if self.describe_calls == 1: - raise TableNotFoundError("table not found: t") - return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None, is_only_declared=True) + raise AssertionError("create must not describe before declaring") def declare_table(self, request: Any) -> Any: - raise TableAlreadyExistsError("Table already exists: t") + requests.append(request) + return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None) - fake = RacingNamespace() - monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: fake) + monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: RecordingNamespace()) resolved = namespace_mod.resolve_namespace_table( - namespace_impl="rest", - namespace_properties=None, - table_id=["t"], - mode="create", + namespace_impl="rest", namespace_properties=None, table_id=["t"], mode="create" ) - assert resolved is not None - assert resolved.is_declared_placeholder - assert fake.describe_calls == 2 - - -def test_create_declare_race_lost_to_real_table_raises(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError - - class RacingNamespace: - def __init__(self) -> None: - self.describe_calls = 0 - def describe_table(self, request: Any) -> Any: - self.describe_calls += 1 - if self.describe_calls == 1: - raise TableNotFoundError("table not found: t") - return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None, is_only_declared=False) - - def declare_table(self, request: Any) -> Any: - raise TableAlreadyExistsError("Table already exists: t") - - monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: RacingNamespace()) - - with pytest.raises(ValueError, match="already exists"): - namespace_mod.resolve_namespace_table( - namespace_impl="rest", - namespace_properties=None, - table_id=["t"], - mode="create", - ) + assert resolved is not None + assert resolved.uri.endswith("t.lance") + assert len(requests) == 1 + assert requests[0].vend_credentials is True -def test_overwrite_resolves_declared_only_stub_on_strict_impl(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - """Strict impls 404 a plain describe on declared-only stubs; overwrite must still find them.""" - from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError +def test_namespace_overwrite_uses_plain_describe(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + requests: list[Any] = [] - class StrictNamespace: + class RecordingNamespace: def describe_table(self, request: Any) -> Any: - if getattr(request, "check_declared", None) is True: - return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None, is_only_declared=True) - raise TableNotFoundError("table not found: t") + requests.append(request) + return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None) def declare_table(self, request: Any) -> Any: - raise TableAlreadyExistsError("Table already exists: t") + raise AssertionError("an existing overwrite target must not be declared") - monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: StrictNamespace()) + monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: RecordingNamespace()) resolved = namespace_mod.resolve_namespace_table( - namespace_impl="rest", - namespace_properties=None, - table_id=["t"], - mode="overwrite", + namespace_impl="rest", namespace_properties=None, table_id=["t"], mode="overwrite" ) - assert resolved is not None - assert resolved.uri.endswith("t.lance") - assert resolved.is_declared_placeholder - - -def test_mem_wal_create_over_materialized_placeholder_rejected(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - import daft_lance.lance_data_sink as sink_mod - - ns = _dir_ns(tmp_path) - schema = daft.from_pydict({"id": [1]}).schema() - sink = daft_lance.LanceDataSink(None, schema, "create", table_id=["t"], use_mem_wal=True, **ns) - - monkeypatch.setattr( - sink_mod, - "resolve_namespace_table", - lambda **kwargs: namespace_mod.ResolvedNamespaceTable( - uri=str(tmp_path / "t.lance"), is_declared_placeholder=True - ), - ) - import pyarrow as pa - fake_dataset = SimpleNamespace(schema=pa.schema([("id", pa.int64())]), latest_version=1) - monkeypatch.setattr("daft_lance.lance_data_sink.lance.dataset", lambda *args, **kwargs: fake_dataset) - - with pytest.raises(ValueError, match="use_mem_wal"): - sink.start() + assert resolved is not None + assert len(requests) == 1 + assert requests[0].model_fields_set == {"id", "vend_credentials"} + assert requests[0].vend_credentials is True def test_construct_lance_dataset_empty_storage_options_falls_back_to_io_config( @@ -397,7 +307,7 @@ def test_sink_survives_pickle_after_start(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) schema = daft.from_pydict({"id": [1]}).schema() - sink = daft_lance.LanceDataSink(None, schema, "create", table_id=["pickled"], **ns) + sink = LanceDataSink(None, schema, "create", table_id=["pickled"], **ns) sink.start() worker_sink = pickle.loads(pickle.dumps(sink)) @@ -450,22 +360,12 @@ def test_sink_empty_storage_options_falls_back_to_io_config(tmp_path: Path) -> N io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) schema = daft.from_pydict({"id": [1]}).schema() - sink = daft_lance.LanceDataSink("s3://bucket/t.lance", schema, "create", io_config, storage_options={}) + sink = LanceDataSink("s3://bucket/t.lance", schema, "create", io_config, storage_options={}) merged = sink._merged_storage_options(namespace_mod.ResolvedNamespaceTable(uri="s3://bucket/t.lance")) assert merged is not None assert merged["access_key_id"] == "io-key" -def test_declared_placeholder_response_fallback_property() -> None: - assert namespace_mod.is_declared_placeholder_response(SimpleNamespace(is_only_declared=True)) - assert namespace_mod.is_declared_placeholder_response( - SimpleNamespace(is_only_declared=None, properties={"lance.declared": "TRUE"}) - ) - assert not namespace_mod.is_declared_placeholder_response( - SimpleNamespace(is_only_declared=None, properties={}, metadata=None) - ) - - def test_construct_lance_dataset_storage_options_priority(monkeypatch: pytest.MonkeyPatch) -> None: import daft_lance.utils as utils_mod from daft.io import IOConfig, S3Config @@ -553,7 +453,7 @@ def test_sink_storage_options_priority(tmp_path: Path) -> None: schema = daft.from_pydict({"id": [1]}).schema() io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) - sink = daft_lance.LanceDataSink( + sink = LanceDataSink( None, schema, "create", diff --git a/tests/io/lancedb/test_namespace_rest_e2e.py b/tests/io/lancedb/test_namespace_rest_e2e.py index ce3b089..2445cc3 100644 --- a/tests/io/lancedb/test_namespace_rest_e2e.py +++ b/tests/io/lancedb/test_namespace_rest_e2e.py @@ -19,6 +19,7 @@ def test_rest_namespace_write_read_append_roundtrip() -> None: import lance import lance_namespace as ln from lance_namespace import CreateNamespaceRequest, DescribeTableRequest, NamespaceExistsRequest + from lance_namespace.errors import TableAlreadyExistsError namespace_properties = {"uri": os.environ["DAFT_LANCE_REST_URI"]} catalog = os.environ.get("DAFT_LANCE_REST_CATALOG", "lance_catalog") @@ -63,3 +64,37 @@ def test_rest_namespace_write_read_append_roundtrip() -> None: assert daft_lance.read_lance(table_id=table_id, **ns).count_rows() == 5 assert lance.dataset(None, namespace_client=namespace, table_id=table_id).count_rows() == 5 + + with pytest.raises(TableAlreadyExistsError): + daft_lance.write_lance( + daft.from_pydict({"id": [99], "label": ["duplicate"], "score": [990]}), + table_id=table_id, + mode="create", + **ns, + ).collect() + assert daft_lance.read_lance(table_id=table_id, **ns).count_rows() == 5 + + daft_lance.write_lance( + daft.from_pydict({"id": [6], "label": ["overwritten"], "score": [60]}), + table_id=table_id, + mode="overwrite", + **ns, + ).collect() + assert daft_lance.read_lance(table_id=table_id, **ns).to_pydict() == { + "id": [6], + "label": ["overwritten"], + "score": [60], + } + + missing_table_id = [catalog, schema, f"overwrite_missing_{uuid.uuid4().hex[:8]}"] + daft_lance.write_lance( + daft.from_pydict({"id": [7], "label": ["created-by-overwrite"], "score": [70]}), + table_id=missing_table_id, + mode="overwrite", + **ns, + ).collect() + assert daft_lance.read_lance(table_id=missing_table_id, **ns).to_pydict() == { + "id": [7], + "label": ["created-by-overwrite"], + "score": [70], + } From fe11b9a99394349ce1e0e9892d81c6501d391784 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Tue, 21 Jul 2026 11:36:02 +0900 Subject: [PATCH 17/25] refactor(namespace): make Lance open context explicit --- daft_lance/_lance.py | 50 +++++------ daft_lance/lance_merge_column.py | 22 +++-- daft_lance/lance_scalar_index.py | 22 +++-- daft_lance/lance_scan.py | 25 +++--- daft_lance/namespace.py | 19 +--- daft_lance/utils.py | 86 ++++++++++++++---- tests/io/lancedb/test_lancedb_point_lookup.py | 2 +- .../io/lancedb/test_lancedb_vector_search.py | 11 +++ tests/io/lancedb/test_namespace.py | 90 ++++++++++++++++++- 9 files changed, 238 insertions(+), 89 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 4e2d321..83adc33 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -17,7 +17,7 @@ from .lance_scalar_index import create_scalar_index_internal from .lance_scan import LanceDBScanOperator from .namespace import validate_uri_or_namespace -from .utils import construct_lance_dataset +from .utils import construct_lance_dataset_handle if TYPE_CHECKING: from lance.dataset import LanceDataset @@ -27,17 +27,6 @@ from daft.dependencies import pa -def _effective_uri(lance_ds: LanceDataset, uri: str | pathlib.Path | None) -> str: - """The dataset location: the user-supplied URI, or the namespace-resolved one.""" - return str(uri) if uri is not None else str(lance_ds.uri) - - -def _effective_storage_options(lance_ds: LanceDataset, storage_options: dict[str, Any] | None) -> dict[str, Any] | None: - """Storage options actually used to open the dataset (includes namespace-vended ones).""" - open_kwargs = getattr(lance_ds, "_lance_open_kwargs", None) or {} - return open_kwargs.get("storage_options") or storage_options - - @PublicAPI def read_lance( uri: str | pathlib.Path | None = None, @@ -155,7 +144,7 @@ def read_lance( io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - ds = construct_lance_dataset( + dataset_handle = construct_lance_dataset_handle( uri_str, io_config=io_config, namespace_impl=namespace_impl, @@ -171,9 +160,11 @@ def read_lance( ) lance_operator = LanceDBScanOperator( - ds, + dataset_handle.dataset, fragment_group_size=fragment_group_size, include_fragment_id=include_fragment_id, + open_kwargs=dataset_handle.open_kwargs, + default_scan_options=dataset_handle.default_scan_options, ) handle = ScanOperatorHandle.from_python_scan_operator(lance_operator) @@ -255,7 +246,7 @@ def merge_columns( io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config # Build Lance dataset handle for committing - lance_ds = construct_lance_dataset( + dataset_handle = construct_lance_dataset_handle( uri, storage_options=storage_options, io_config=io_config, @@ -272,12 +263,13 @@ def merge_columns( ) return merge_columns_internal( - lance_ds, - _effective_uri(lance_ds, uri), + dataset_handle.dataset, + dataset_handle.uri, transform=transform, read_columns=read_columns, reader_schema=reader_schema, - storage_options=_effective_storage_options(lance_ds, storage_options), + storage_options=dataset_handle.storage_options, + namespace_kwargs=dataset_handle.namespace_kwargs, daft_remote_args=daft_remote_args, concurrency=concurrency, ) @@ -362,7 +354,7 @@ def merge_columns_df( io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config # Build Lance dataset handle for committing - lance_ds = construct_lance_dataset( + dataset_handle = construct_lance_dataset_handle( uri, storage_options=storage_options, io_config=io_config, @@ -386,11 +378,12 @@ def merge_columns_df( return merge_columns_from_df( df, - lance_ds=lance_ds, - uri=_effective_uri(lance_ds, uri), + lance_ds=dataset_handle.dataset, + uri=dataset_handle.uri, read_columns=read_columns, reader_schema=reader_schema, - storage_options=_effective_storage_options(lance_ds, storage_options), + storage_options=dataset_handle.storage_options, + namespace_kwargs=dataset_handle.namespace_kwargs, daft_remote_args=daft_remote_args, concurrency=concurrency, left_on=left_on, @@ -518,7 +511,7 @@ def create_scalar_index( """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - lance_ds = construct_lance_dataset( + dataset_handle = construct_lance_dataset_handle( uri, storage_options=storage_options, io_config=io_config, @@ -535,13 +528,14 @@ def create_scalar_index( ) create_scalar_index_internal( - lance_ds=lance_ds, - uri=_effective_uri(lance_ds, uri), + lance_ds=dataset_handle.dataset, + uri=dataset_handle.uri, column=column, index_type=index_type, name=name, replace=replace, - storage_options=_effective_storage_options(lance_ds, storage_options), + storage_options=dataset_handle.storage_options, + namespace_kwargs=dataset_handle.namespace_kwargs, fragment_group_size=fragment_group_size, num_partitions=num_partitions, max_concurrency=max_concurrency, @@ -614,7 +608,7 @@ def compact_files( """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config - lance_ds = construct_lance_dataset( + dataset_handle = construct_lance_dataset_handle( uri, storage_options=storage_options, io_config=io_config, @@ -631,7 +625,7 @@ def compact_files( ) return compact_files_internal( - lance_ds=lance_ds, + lance_ds=dataset_handle.dataset, compaction_options=compaction_options, partition_num=partition_num, concurrency=concurrency, diff --git a/daft_lance/lance_merge_column.py b/daft_lance/lance_merge_column.py index df2d682..44f44ea 100644 --- a/daft_lance/lance_merge_column.py +++ b/daft_lance/lance_merge_column.py @@ -14,8 +14,6 @@ from daft.udf import cls as daft_cls from daft.udf import method -from .namespace import namespace_kwargs_for_dataset - if TYPE_CHECKING: import pathlib from collections.abc import Callable @@ -60,6 +58,7 @@ def merge_columns_internal( read_columns: list[str] | None = None, reader_schema: pa.Schema | None = None, storage_options: dict[str, Any] | None = None, + namespace_kwargs: dict[str, Any] | None = None, daft_remote_args: dict[str, Any] | None = None, concurrency: int | None = None, ) -> lance.LanceDataset: @@ -95,7 +94,7 @@ def merge_columns_internal( op, read_version=lance_ds.version, storage_options=storage_options, - **namespace_kwargs_for_dataset(lance_ds), + **(namespace_kwargs or {}), ) @@ -372,6 +371,7 @@ def merge_columns_from_df( read_columns: list[str] | None = None, reader_schema: pa.Schema | None = None, storage_options: dict[str, Any] | None = None, + namespace_kwargs: dict[str, Any] | None = None, daft_remote_args: dict[str, Any] | None = None, concurrency: int | None = None, left_on: str | None = "_rowaddr", @@ -416,7 +416,14 @@ def merge_columns_from_df( use_fast_path = _can_use_fast_path(df, lance_ds, join_key) if use_fast_path: - return _merge_fast_path(df, lance_ds, uri, new_cols, storage_options=storage_options) + return _merge_fast_path( + df, + lance_ds, + uri, + new_cols, + storage_options=storage_options, + namespace_kwargs=namespace_kwargs, + ) else: return _merge_slow_path( df, @@ -428,6 +435,7 @@ def merge_columns_from_df( reader_schema, batch_size, storage_options=storage_options, + namespace_kwargs=namespace_kwargs, ) @@ -437,6 +445,7 @@ def _merge_fast_path( uri: str | pathlib.Path, new_column_names: list[str], storage_options: dict[str, Any] | None = None, + namespace_kwargs: dict[str, Any] | None = None, ) -> lance.LanceDataset: """Metadata-only add_columns: write raw .lance files and stitch into fragment metadata.""" handler = FastPathFragmentWriter(lance_ds, str(uri), new_column_names, storage_options=storage_options) @@ -476,7 +485,7 @@ def _merge_fast_path( op, read_version=lance_ds.version, storage_options=storage_options, - **namespace_kwargs_for_dataset(lance_ds), + **(namespace_kwargs or {}), ) @@ -490,6 +499,7 @@ def _merge_slow_path( reader_schema: pa.Schema | None, batch_size: int | None, storage_options: dict[str, Any] | None = None, + namespace_kwargs: dict[str, Any] | None = None, ) -> lance.LanceDataset: """Original keyed-join merge path: rewrites fragment data.""" handler_udf = GroupFragmentMergeUDF( @@ -525,5 +535,5 @@ def _merge_slow_path( op, read_version=lance_ds.version, storage_options=storage_options, - **namespace_kwargs_for_dataset(lance_ds), + **(namespace_kwargs or {}), ) diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index 2cd2fa2..6d12969 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -14,7 +14,6 @@ import lance from daft.dependencies import pa -from daft_lance.namespace import namespace_kwargs_for_dataset from daft_lance.utils import distribute_fragments_balanced logger = logging.getLogger(__name__) @@ -140,6 +139,7 @@ def create_scalar_index_internal( name: str | None = None, replace: bool = False, storage_options: dict[str, Any] | None = None, + namespace_kwargs: dict[str, Any] | None = None, fragment_group_size: int | None = None, num_partitions: int | None = None, max_concurrency: int | None = None, @@ -297,6 +297,7 @@ def create_scalar_index_internal( index_type=index_type, name=name, storage_options=storage_options, + namespace_kwargs=namespace_kwargs, fragment_data=fragment_data, fragment_ids_to_use=fragment_ids_to_use, num_partitions=num_partitions, @@ -312,6 +313,7 @@ def create_scalar_index_internal( name=name, replace=replace, storage_options=storage_options, + namespace_kwargs=namespace_kwargs, fragment_data=fragment_data, fragment_ids_to_use=fragment_ids_to_use, num_partitions=num_partitions, @@ -328,6 +330,7 @@ def _create_segmented_index( index_type: str, name: str, storage_options: dict[str, Any] | None, + namespace_kwargs: dict[str, Any] | None, fragment_data: list[dict[str, list[int]]], fragment_ids_to_use: list[int], num_partitions: int | None, @@ -387,8 +390,11 @@ def _create_segmented_index( # Reload dataset to pick up the latest version (segment files were written # by workers against the version that was current at their invocation time). - namespace_kwargs = namespace_kwargs_for_dataset(lance_ds) - lance_ds = lance.dataset(None if namespace_kwargs else uri, storage_options=storage_options, **namespace_kwargs) + lance_ds = lance.dataset( + None if namespace_kwargs else uri, + storage_options=storage_options, + **(namespace_kwargs or {}), + ) index_metas = _prepare_index_segments_for_commit(lance_ds, index_type, index_metas) logger.info( @@ -423,6 +429,7 @@ def _create_partitioned_index( name: str, replace: bool, storage_options: dict[str, Any] | None, + namespace_kwargs: dict[str, Any] | None, fragment_data: list[dict[str, list[int]]], fragment_ids_to_use: list[int], num_partitions: int | None, @@ -462,8 +469,11 @@ def _create_partitioned_index( df.collect() logger.info("Starting index metadata merging by reloading dataset to get latest state") - namespace_kwargs = namespace_kwargs_for_dataset(lance_ds) - lance_ds = lance.dataset(None if namespace_kwargs else uri, storage_options=storage_options, **namespace_kwargs) + lance_ds = lance.dataset( + None if namespace_kwargs else uri, + storage_options=storage_options, + **(namespace_kwargs or {}), + ) lance_ds.merge_index_metadata(index_id, index_type) logger.info("Starting atomic index creation and commit") @@ -509,7 +519,7 @@ def _create_partitioned_index( create_index_op, read_version=lance_ds.version, storage_options=storage_options, - **namespace_kwargs, + **(namespace_kwargs or {}), ) logger.info("Index %s created successfully with ID %s", name, index_id) diff --git a/daft_lance/lance_scan.py b/daft_lance/lance_scan.py index e350479..18cb957 100644 --- a/daft_lance/lance_scan.py +++ b/daft_lance/lance_scan.py @@ -135,8 +135,12 @@ def __init__( ds: lance.LanceDataset, fragment_group_size: int | None = None, include_fragment_id: bool | None = False, + open_kwargs: dict[str, Any] | None = None, + default_scan_options: dict[str, Any] | None = None, ): self._ds = ds + self._open_kwargs = dict(open_kwargs or {}) + self._default_scan_options = default_scan_options self._pushed_filters: list[PyExpr] | None = None self._remaining_filters: list[PyExpr] | None = None self._fragment_group_size = fragment_group_size @@ -256,11 +260,10 @@ def _create_count_rows_scan_task(self, pushdowns: PyPushdowns) -> Iterator[ScanT """Create scan task for counting rows.""" fields = pushdowns.aggregation_required_column_names() new_schema = Schema.from_pyarrow_schema(pa.schema([pa.field(fields[0], pa.uint64())])) - open_kwargs = getattr(self._ds, "_lance_open_kwargs", None) yield ScanTask.python_factory_func_scan_task( module=_lancedb_count_result_function.__module__, func_name=_lancedb_count_result_function.__name__, - func_args=(self._ds.uri, open_kwargs, fields[0], self._combine_filters_to_arrow()), + func_args=(self._ds.uri, self._open_kwargs, fields[0], self._combine_filters_to_arrow()), schema=new_schema._schema, num_rows=1, size_bytes=None, @@ -276,7 +279,6 @@ def _create_scan_tasks_with_limit_and_no_filters( assert self._pushed_filters is None, "Expected no filters when creating scan tasks with limit and no filters" assert pushdowns.limit is not None, "Expected a limit when creating scan tasks with limit and no filters" - open_kwargs = getattr(self._ds, "_lance_open_kwargs", None) fragments = self._ds.get_fragments() remaining_limit = pushdowns.limit @@ -301,7 +303,7 @@ def _create_scan_tasks_with_limit_and_no_filters( func_name=_lancedb_table_factory_function.__name__, func_args=( self._ds.uri, - open_kwargs, + self._open_kwargs, [fragment.fragment_id], required_columns, None, @@ -321,7 +323,6 @@ def _create_regular_scan_tasks( self, pushdowns: PyPushdowns, required_columns: list[str] | None, nearest_option: dict[str, Any] | None = None ) -> Iterator[ScanTask]: """Create regular scan tasks without count pushdown.""" - open_kwargs = getattr(self._ds, "_lance_open_kwargs", None) fragments = self._ds.get_fragments() pushed_expr = self._combine_filters_to_arrow() @@ -336,7 +337,7 @@ def _python_factory_func_scan_task( func_name=_lancedb_table_factory_function.__name__, func_args=( self._ds.uri, - open_kwargs, + self._open_kwargs, fragment_ids, required_columns, pushed_expr, @@ -460,15 +461,15 @@ def _should_use_index_for_point_lookup(self) -> bool: def _nearest_default_option(self) -> dict[str, Any] | None: """Return the default nearest option configured on the Lance dataset, if any. - Prefer Daft-specific `_daft_default_scan_options` to preserve options stripped before `lance.dataset` (e.g., `nearest`). + Daft keeps the original options explicitly because unsupported planning + keys such as ``nearest`` are stripped before calling ``lance.dataset``. """ - default_opts = getattr(self._ds, "_daft_default_scan_options", None) + default_opts = self._default_scan_options if not isinstance(default_opts, dict): + # Preserve direct ``LanceDBScanOperator(ds)`` compatibility. This + # is Lance's own stored constructor option, not Daft metadata + # attached to the third-party object. default_opts = getattr(self._ds, "_default_scan_options", None) - if not isinstance(default_opts, dict): - open_kwargs = getattr(self._ds, "_lance_open_kwargs", None) - if isinstance(open_kwargs, dict): - default_opts = open_kwargs.get("default_scan_options") if not isinstance(default_opts, dict): return None diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index fd6b2f4..a88a4cd 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -205,7 +205,7 @@ def merge_storage_options(*layers: dict[str, Any] | None) -> dict[str, Any] | No def pop_namespace_params(open_kwargs: dict[str, Any]) -> tuple[str | None, dict[str, str] | None, list[str] | None]: - """Remove and return the namespace triple from a ``_lance_open_kwargs`` dict.""" + """Remove and return the namespace triple from serialized open arguments.""" return ( open_kwargs.pop("namespace_impl", None), open_kwargs.pop("namespace_properties", None), @@ -213,23 +213,8 @@ def pop_namespace_params(open_kwargs: dict[str, Any]) -> tuple[str | None, dict[ ) -def namespace_kwargs_for_dataset(ds: Any) -> dict[str, Any]: - """Namespace kwargs for commit-style calls, recovered from ``ds._lance_open_kwargs``. - - ``lance.LanceDataset.commit`` is a static method, so a dataset opened through a - namespace does not carry its client into commits; re-derive it from the open - kwargs stashed by ``construct_lance_dataset``. - """ - open_kwargs = getattr(ds, "_lance_open_kwargs", None) or {} - return get_namespace_kwargs( - open_kwargs.get("namespace_impl"), - open_kwargs.get("namespace_properties"), - open_kwargs.get("table_id"), - ) - - def open_dataset_from_open_kwargs(ds_uri: str | None, open_kwargs: dict[str, Any] | None) -> lance.LanceDataset: - """Re-open a dataset on a worker from serialized ``_lance_open_kwargs``. + """Re-open a dataset on a worker from serialized open arguments. The namespace client is not picklable, so only the (impl, properties, table_id) triple travels with the task; the client is re-created here via the lru cache. diff --git a/daft_lance/utils.py b/daft_lance/utils.py index fb4cebc..7447fd1 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any import lance @@ -24,6 +25,34 @@ logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class LanceDatasetHandle: + """An opened Lance dataset together with the context needed to reuse it. + + ``lance.LanceDataset`` does not expose all arguments used to open it. Keep + those arguments in this Daft-owned value object instead of attaching + private attributes to the third-party dataset instance. + """ + + dataset: lance.LanceDataset + uri: str + open_kwargs: dict[str, Any] = field(repr=False) + default_scan_options: dict[str, Any] | None = field(default=None, repr=False) + + @property + def storage_options(self) -> dict[str, Any] | None: + options = self.open_kwargs.get("storage_options") + return options if isinstance(options, dict) else None + + @property + def namespace_kwargs(self) -> dict[str, Any]: + return get_namespace_kwargs( + self.open_kwargs.get("namespace_impl"), + self.open_kwargs.get("namespace_properties"), + self.open_kwargs.get("table_id"), + ) + + def distribute_fragments_balanced(fragments: list[Any], fragment_group_size: int) -> list[dict[str, list[int]]]: """Distribute fragments across workers using a balanced algorithm considering fragment sizes.""" if fragment_group_size <= 0: @@ -91,7 +120,7 @@ def distribute_fragments_balanced(fragments: list[Any], fragment_group_size: int return non_empty_batches -def construct_lance_dataset( +def construct_lance_dataset_handle( uri: str | pathlib.Path | None, version: int | str | None = None, storage_options: dict[str, Any] | None = None, @@ -100,8 +129,8 @@ def construct_lance_dataset( namespace_properties: dict[str, str] | None = None, table_id: list[str] | None = None, **kwargs: Any, -) -> lance.LanceDataset: - """Construct a Lance dataset with common options. +) -> LanceDatasetHandle: + """Construct a Lance dataset and retain its reusable open context. Storage options are layered from lowest to highest priority: io_config-derived < user-provided ``storage_options`` < namespace-vended. @@ -144,7 +173,7 @@ def construct_lance_dataset( kwargs["default_scan_options"] = original_default_scan_options dataset_uri = None if has_namespace_params(namespace_impl, table_id) else resolved_uri - ds = lance.dataset( + dataset = lance.dataset( dataset_uri, storage_options=merged_storage_options, version=version, @@ -160,19 +189,42 @@ def construct_lance_dataset( "table_id": table_id, } effective_kwargs.update(kwargs or {}) - try: - ds._lance_open_kwargs = effective_kwargs # type: ignore[attr-defined] - except Exception: - pass - - # Preserve the full user-provided defaults (including nearest) for Daft's planning - # even if we stripped keys out before calling `lance.dataset`. - try: - ds._daft_default_scan_options = original_default_scan_options # type: ignore[attr-defined] - except Exception: - pass - - return ds + return LanceDatasetHandle( + dataset=dataset, + uri=resolved_uri, + open_kwargs=effective_kwargs, + # Preserve the full user-provided defaults (including nearest) for + # Daft's planning even if keys were stripped before calling Lance. + default_scan_options=original_default_scan_options if isinstance(original_default_scan_options, dict) else None, + ) + + +def construct_lance_dataset( + uri: str | pathlib.Path | None, + version: int | str | None = None, + storage_options: dict[str, Any] | None = None, + io_config: IOConfig | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, + table_id: list[str] | None = None, + **kwargs: Any, +) -> lance.LanceDataset: + """Construct a Lance dataset with common options. + + This compatibility wrapper preserves the historical return type. Internal + callers that also need the resolved URI or reusable open arguments should + use :func:`construct_lance_dataset_handle`. + """ + return construct_lance_dataset_handle( + uri, + version=version, + storage_options=storage_options, + io_config=io_config, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, + **kwargs, + ).dataset def combine_filters_to_arrow(predicates: list[Any] | None) -> pa.compute.Expression | None: diff --git a/tests/io/lancedb/test_lancedb_point_lookup.py b/tests/io/lancedb/test_lancedb_point_lookup.py index 92b07ff..63d2dd7 100644 --- a/tests/io/lancedb/test_lancedb_point_lookup.py +++ b/tests/io/lancedb/test_lancedb_point_lookup.py @@ -104,7 +104,7 @@ def test_scanner_without_fragments(lance_dataset, idx_type): # Invoke factory with fragment_ids=None to exercise index-driven fragment selection gen = lance_scan._lancedb_table_factory_function( ds.uri, - getattr(ds, "_lance_open_kwargs", None), + {}, None, ["id", "value"], arrow_filter, diff --git a/tests/io/lancedb/test_lancedb_vector_search.py b/tests/io/lancedb/test_lancedb_vector_search.py index 54d4077..babd899 100755 --- a/tests/io/lancedb/test_lancedb_vector_search.py +++ b/tests/io/lancedb/test_lancedb_vector_search.py @@ -9,6 +9,7 @@ import daft from daft import col +from daft_lance.lance_scan import LanceDBScanOperator def build_single_fragment_dataset(tmp_path_factory) -> str: @@ -115,6 +116,16 @@ def test_nearest_global_single_scan_task(tmp_path_factory) -> None: assert "Num Scan Tasks = 1" in explain_output +def test_direct_scan_operator_uses_lance_default_nearest(tmp_path_factory) -> None: + dataset_path = build_multi_fragment_dataset(tmp_path_factory) + nearest = {"column": "vector", "q": pa.array([0.0, 0.0], type=pa.float32()), "k": 1} + dataset = lance.dataset(dataset_path, default_scan_options={"nearest": nearest}) + + scan = LanceDBScanOperator(dataset) + + assert scan._nearest_default_option() == nearest + + def test_nearest_metric_cosine_k1(tmp_path_factory) -> None: dataset_path = build_metric_dataset(tmp_path_factory) diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index 11c0109..83b02f7 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -16,6 +16,13 @@ def _dir_ns(tmp_path: Path) -> dict[str, Any]: return {"namespace_impl": "dir", "namespace_properties": {"root": str(tmp_path)}} +def _double_score(batch: Any) -> Any: + import pyarrow as pa + import pyarrow.compute as pc + + return pa.RecordBatch.from_arrays([pc.multiply(batch["score"], 2)], ["doubled"]) + + def test_namespace_write_read_append_roundtrip(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["roundtrip"] @@ -135,6 +142,51 @@ def test_namespace_merge_columns_df(tmp_path: Path) -> None: assert result["tripled"] == [3, 6, 9] +def test_namespace_merge_columns_transform(tmp_path: Path) -> None: + ns = _dir_ns(tmp_path) + table_id = ["merge_cols_transform"] + + daft_lance.write_lance( + daft.from_pydict({"id": [1, 2, 3], "score": [10, 20, 30]}), + table_id=table_id, + mode="create", + **ns, + ).collect() + + daft_lance.merge_columns(table_id=table_id, transform=_double_score, read_columns=["score"], **ns) + + result = daft_lance.read_lance(table_id=table_id, **ns).sort("id").to_pydict() + assert result["doubled"] == [20, 40, 60] + + +def test_namespace_merge_columns_df_slow_path(tmp_path: Path) -> None: + ns = _dir_ns(tmp_path) + table_id = ["merge_cols_slow"] + + daft_lance.write_lance( + daft.from_pydict({"id": [1, 2, 3], "score": [10, 20, 30]}), + table_id=table_id, + mode="create", + **ns, + ).collect() + + source = daft_lance.read_lance( + table_id=table_id, + default_scan_options={"with_row_address": True}, + include_fragment_id=True, + **ns, + ).limit(2) + source = source.with_column("partial_score", source["score"] * 10) + daft_lance.merge_columns_df( + source.select("fragment_id", "_rowaddr", "partial_score"), + table_id=table_id, + **ns, + ) + + result = daft_lance.read_lance(table_id=table_id, **ns).sort("id").to_pydict() + assert result["partial_score"] == [100, 200, None] + + def test_namespace_create_scalar_index(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["indexed"] @@ -158,6 +210,36 @@ def test_namespace_create_scalar_index(tmp_path: Path) -> None: assert any(idx["fields"] == ["price"] for idx in indices) +@pytest.mark.parametrize("segmented", [False, True]) +def test_namespace_create_distributed_inverted_index(tmp_path: Path, segmented: bool) -> None: + ns = _dir_ns(tmp_path) + table_id = [f"inverted_{segmented}"] + + daft_lance.write_lance( + daft.from_pydict({"id": list(range(20)), "text": [f"document {i}" for i in range(20)]}), + table_id=table_id, + mode="create", + **ns, + ).collect() + + daft_lance.create_scalar_index( + table_id=table_id, + column="text", + index_type="INVERTED", + name="text_idx", + segmented=segmented, + **ns, + ) + + import lance + import lance_namespace as ln + from lance_namespace import DescribeTableRequest + + location = ln.connect("dir", {"root": str(tmp_path)}).describe_table(DescribeTableRequest(id=table_id)).location + indices = lance.dataset(location).list_indices() + assert any(idx["name"] == "text_idx" and idx["fields"] == ["text"] for idx in indices) + + def test_namespace_compact_files(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["compacted"] @@ -293,9 +375,10 @@ def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **k monkeypatch.setattr("daft_lance.utils.lance.dataset", fake_dataset) io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) - utils_mod.construct_lance_dataset("s3://bucket/t.lance", storage_options={}, io_config=io_config) + dataset = utils_mod.construct_lance_dataset("s3://bucket/t.lance", storage_options={}, io_config=io_config) merged = captured["storage_options"] + assert isinstance(dataset, FakeDataset) assert merged is not None assert merged["access_key_id"] == "io-key" @@ -392,7 +475,7 @@ def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **k ) io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) - utils_mod.construct_lance_dataset( + handle = utils_mod.construct_lance_dataset_handle( None, io_config=io_config, storage_options={"access_key_id": "user-key", "user_option": "kept"}, @@ -408,6 +491,9 @@ def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **k assert merged["user_option"] == "kept" # user-provided keys survive assert merged["secret_access_key"] == "io-secret" # io_config fills the gaps assert captured["uri"] is None # namespace addressing passes uri=None + assert handle.storage_options == merged + assert handle.uri == "s3://bucket/t.lance" + assert not hasattr(handle.dataset, "_lance_open_kwargs") def test_construct_lance_dataset_io_config_reaches_namespace_location(monkeypatch: pytest.MonkeyPatch) -> None: From f7f665c6d65ac750c835362e661e885fdee9cb59 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Tue, 21 Jul 2026 12:17:07 +0900 Subject: [PATCH 18/25] fix(namespace): propagate managed versioning through commits --- daft_lance/_lance.py | 13 +++++-- daft_lance/lance_data_sink.py | 22 ++++++++--- daft_lance/lance_scalar_index.py | 5 ++- daft_lance/namespace.py | 17 ++++++++- daft_lance/utils.py | 14 +++++++ .../lancedb/test_lance_data_sink_internals.py | 24 +++++++++++- tests/io/lancedb/test_namespace.py | 38 +++++++++++++++++-- tests/io/lancedb/test_namespace_rest_e2e.py | 24 +++++++++++- 8 files changed, 138 insertions(+), 19 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 83adc33..a58cc5d 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -101,7 +101,7 @@ def read_lance( each fragment will be processed individually (default behavior). include_fragment_id : Optional, bool Whether to display fragment_id. - if you have the behavior of 'merge_columns_df' or 'write_lance(mode = 'merge')', the `include_fragment_id` must be set to True + Set this to True when preparing input for ``merge_columns_df``. checkpoint: Optional :class:`daft.CheckpointConfig` for progress tracking across runs. Bundles the checkpoint store, the source key column (``on=``), and optional anti-join tuning. Rows whose key already exists in the store are skipped on re-run. Requires the Ray runner. @@ -269,7 +269,7 @@ def merge_columns( read_columns=read_columns, reader_schema=reader_schema, storage_options=dataset_handle.storage_options, - namespace_kwargs=dataset_handle.namespace_kwargs, + namespace_kwargs=dataset_handle.commit_kwargs, daft_remote_args=daft_remote_args, concurrency=concurrency, ) @@ -383,7 +383,7 @@ def merge_columns_df( read_columns=read_columns, reader_schema=reader_schema, storage_options=dataset_handle.storage_options, - namespace_kwargs=dataset_handle.namespace_kwargs, + namespace_kwargs=dataset_handle.commit_kwargs, daft_remote_args=daft_remote_args, concurrency=concurrency, left_on=left_on, @@ -536,6 +536,7 @@ def create_scalar_index( replace=replace, storage_options=dataset_handle.storage_options, namespace_kwargs=dataset_handle.namespace_kwargs, + namespace_commit_kwargs=dataset_handle.commit_kwargs, fragment_group_size=fragment_group_size, num_partitions=num_partitions, max_concurrency=max_concurrency, @@ -662,6 +663,12 @@ def write_lance( Returns: DataFrame: write statistics (num_fragments, num_deleted_rows, num_small_files, version). + Note: + Target-dependent validation is performed when the returned DataFrame is + executed (for example, by ``collect()``), not while the logical write + plan is constructed. This includes missing/duplicate targets, append + schema compatibility, and storage-version conflicts. + Examples: >>> import daft, daft_lance >>> df = daft.from_pydict({"id": [1, 2]}) diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index bda33d1..946cd3c 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -28,6 +28,7 @@ ) from daft_lance.namespace import ( ResolvedNamespaceTable, + get_namespace_commit_kwargs, get_namespace_kwargs, get_write_fragments_kwargs, merge_storage_options, @@ -99,6 +100,7 @@ def __init__( # Construction must stay side-effect free: no namespace calls, no dataset opens. self._table_uri: str | None = None self._storage_options: dict[str, str] | None = None + self._managed_versioning = False self._data_storage_version: LanceStorageVersion | None = None self._effective_pyarrow_schema: pa.Schema | None = None self._version: int = 0 @@ -123,6 +125,7 @@ def start(self) -> None: resolved = self._resolve_table() self._table_uri = resolved.uri self._storage_options = self._merged_storage_options(resolved) + self._managed_versioning = resolved.managed_versioning existing = self._absorb_existing_dataset() existing_version = getattr(existing, "data_storage_version", None) if existing is not None else None @@ -144,6 +147,15 @@ def start(self) -> None: def _namespace_kwargs(self) -> dict[str, Any]: return get_namespace_kwargs(self._namespace_impl, self._namespace_properties, self._table_id) + @property + def _namespace_commit_kwargs(self) -> dict[str, Any]: + return get_namespace_commit_kwargs( + self._namespace_impl, + self._namespace_properties, + self._table_id, + self._managed_versioning, + ) + @property def _dataset_uri_arg(self) -> str | None: return None if self._namespace_impl is not None and self._table_id is not None else self._table_uri @@ -165,14 +177,12 @@ def _resolve_table(self) -> ResolvedNamespaceTable: def _merged_storage_options(self, resolved: ResolvedNamespaceTable) -> dict[str, str] | None: """Layer storage options: io_config-derived < user-provided < namespace-vended. - For a plain uri, non-empty user-provided options replace the - io_config-derived ones entirely (falsy fallthrough, matching the read - entry points via construct_lance_dataset). + For a plain URI, explicitly provided options (including ``{}``) replace + the io_config-derived ones, preserving the historical sink behavior. """ io_derived = io_config_to_storage_options(self._io_config, resolved.uri) if self._uri is not None: - base = self._user_storage_options or io_derived - return merge_storage_options(base, resolved.storage_options) + return self._user_storage_options if self._user_storage_options is not None else io_derived return merge_storage_options(io_derived, self._user_storage_options, resolved.storage_options) @staticmethod @@ -415,7 +425,7 @@ def _finalize_cow(self, write_results: list[WriteResult[list[FragmentMetadata]]] operation, read_version=self._version, storage_options=self._storage_options, - **self._namespace_kwargs, + **self._namespace_commit_kwargs, ) stats = dataset.stats.dataset_stats() stats_dict = MicroPartition.from_pydict( diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index 6d12969..298bc40 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -140,6 +140,7 @@ def create_scalar_index_internal( replace: bool = False, storage_options: dict[str, Any] | None = None, namespace_kwargs: dict[str, Any] | None = None, + namespace_commit_kwargs: dict[str, Any] | None = None, fragment_group_size: int | None = None, num_partitions: int | None = None, max_concurrency: int | None = None, @@ -314,6 +315,7 @@ def create_scalar_index_internal( replace=replace, storage_options=storage_options, namespace_kwargs=namespace_kwargs, + namespace_commit_kwargs=namespace_commit_kwargs, fragment_data=fragment_data, fragment_ids_to_use=fragment_ids_to_use, num_partitions=num_partitions, @@ -430,6 +432,7 @@ def _create_partitioned_index( replace: bool, storage_options: dict[str, Any] | None, namespace_kwargs: dict[str, Any] | None, + namespace_commit_kwargs: dict[str, Any] | None, fragment_data: list[dict[str, list[int]]], fragment_ids_to_use: list[int], num_partitions: int | None, @@ -519,7 +522,7 @@ def _create_partitioned_index( create_index_op, read_version=lance_ds.version, storage_options=storage_options, - **(namespace_kwargs or {}), + **(namespace_commit_kwargs or namespace_kwargs or {}), ) logger.info("Index %s created successfully with ID %s", name, index_id) diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index a88a4cd..0b3a713 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -113,6 +113,19 @@ def get_namespace_kwargs( return {"namespace_client": namespace, "table_id": table_id} +def get_namespace_commit_kwargs( + namespace_impl: str | None, + namespace_properties: dict[str, str] | None, + table_id: list[str] | None, + managed_versioning: bool, +) -> dict[str, Any]: + """Namespace kwargs for commit APIs, including catalog-managed versioning.""" + kwargs = get_namespace_kwargs(namespace_impl, namespace_properties, table_id) + if kwargs: + kwargs["namespace_client_managed_versioning"] = managed_versioning + return kwargs + + # `lance.fragment.write_fragments` accepts the same namespace kwargs as `lance.dataset`. get_write_fragments_kwargs = get_namespace_kwargs @@ -133,16 +146,18 @@ def _response_location(response: Any) -> str: @dataclass(frozen=True) class ResolvedNamespaceTable: - """A namespace table resolved to a physical location and storage options.""" + """A namespace table resolved to its physical access and versioning policy.""" uri: str storage_options: dict[str, str] | None = None + managed_versioning: bool = False def _resolved_from_response(response: Any) -> ResolvedNamespaceTable: return ResolvedNamespaceTable( uri=_response_location(response), storage_options=_storage_options(response), + managed_versioning=getattr(response, "managed_versioning", None) is True, ) diff --git a/daft_lance/utils.py b/daft_lance/utils.py index 7447fd1..dc6cd34 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -10,6 +10,7 @@ from daft.io.object_store_options import io_config_to_storage_options from daft.logical.schema import Schema as DaftSchema from daft_lance.namespace import ( + get_namespace_commit_kwargs, get_namespace_kwargs, has_namespace_params, merge_storage_options, @@ -37,6 +38,7 @@ class LanceDatasetHandle: dataset: lance.LanceDataset uri: str open_kwargs: dict[str, Any] = field(repr=False) + managed_versioning: bool = False default_scan_options: dict[str, Any] | None = field(default=None, repr=False) @property @@ -52,6 +54,15 @@ def namespace_kwargs(self) -> dict[str, Any]: self.open_kwargs.get("table_id"), ) + @property + def commit_kwargs(self) -> dict[str, Any]: + return get_namespace_commit_kwargs( + self.open_kwargs.get("namespace_impl"), + self.open_kwargs.get("namespace_properties"), + self.open_kwargs.get("table_id"), + self.managed_versioning, + ) + def distribute_fragments_balanced(fragments: list[Any], fragment_group_size: int) -> list[dict[str, list[int]]]: """Distribute fragments across workers using a balanced algorithm considering fragment sizes.""" @@ -140,6 +151,7 @@ def construct_lance_dataset_handle( validate_uri_or_namespace(uri, namespace_impl, table_id, namespace_properties) resolved_uri = str(uri) if uri is not None else None namespace_storage_options = None + managed_versioning = False if resolved_uri is None: resolved = resolve_namespace_table( namespace_impl=namespace_impl, @@ -150,6 +162,7 @@ def construct_lance_dataset_handle( if resolved is not None: resolved_uri = resolved.uri namespace_storage_options = resolved.storage_options + managed_versioning = resolved.managed_versioning if resolved_uri is None: raise ValueError("Unable to resolve Lance dataset URI.") @@ -193,6 +206,7 @@ def construct_lance_dataset_handle( dataset=dataset, uri=resolved_uri, open_kwargs=effective_kwargs, + managed_versioning=managed_versioning, # Preserve the full user-provided defaults (including nearest) for # Daft's planning even if keys were stripped before calling Lance. default_scan_options=original_default_scan_options if isinstance(original_default_scan_options, dict) else None, diff --git a/tests/io/lancedb/test_lance_data_sink_internals.py b/tests/io/lancedb/test_lance_data_sink_internals.py index e6983c6..c559106 100644 --- a/tests/io/lancedb/test_lance_data_sink_internals.py +++ b/tests/io/lancedb/test_lance_data_sink_internals.py @@ -1,7 +1,7 @@ """Unit-level tests for LanceDataSink internals. -Phase 1 covers the _load_existing_dataset typed-exception path and the Lance -message-format pin. +These tests also pin the sink lifecycle contract: construction validates local +arguments only, while target-dependent validation runs in ``start()``. """ from __future__ import annotations @@ -30,6 +30,26 @@ def test_load_existing_dataset_missing_returns_none_for_create(tmp_path): assert sink is not None +def test_uri_create_existing_error_is_deferred_until_start(tmp_path): + uri = str(tmp_path / "existing") + lance.write_dataset(pa.table({"a": [1]}), uri) + + sink = LanceDataSink(uri=uri, schema=pa.schema([("a", pa.int64())]), mode="create") + + with pytest.raises(ValueError, match="already exists"): + sink.start() + + +def test_uri_append_schema_error_is_deferred_until_start(tmp_path): + uri = str(tmp_path / "schema-mismatch") + lance.write_dataset(pa.table({"a": pa.array([1], type=pa.int64())}), uri) + + sink = LanceDataSink(uri=uri, schema=pa.schema([("a", pa.struct([("value", pa.string())]))]), mode="append") + + with pytest.raises(ValueError, match="Schema of data does not match"): + sink.start() + + def test_mode_merge_rejected_at_construction(tmp_path): schema = pa.schema([("a", pa.int64())]) with pytest.raises(ValueError, match='mode="merge" is no longer supported'): diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index 83b02f7..2a939e1 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -357,6 +357,35 @@ def declare_table(self, request: Any) -> Any: assert requests[0].vend_credentials is True +def test_namespace_response_preserves_managed_versioning(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + response = SimpleNamespace( + location=str(tmp_path / "managed.lance"), + storage_options={"token": "temporary"}, + managed_versioning=True, + ) + + resolved = namespace_mod._resolved_from_response(response) + + assert resolved.managed_versioning is True + + +def test_namespace_commit_kwargs_include_managed_versioning(monkeypatch: pytest.MonkeyPatch) -> None: + namespace_client = object() + monkeypatch.setattr( + namespace_mod, + "get_namespace_kwargs", + lambda *args: {"namespace_client": namespace_client, "table_id": ["catalog", "table"]}, + ) + + kwargs = namespace_mod.get_namespace_commit_kwargs("rest", {}, ["catalog", "table"], True) + + assert kwargs == { + "namespace_client": namespace_client, + "table_id": ["catalog", "table"], + "namespace_client_managed_versioning": True, + } + + def test_construct_lance_dataset_empty_storage_options_falls_back_to_io_config( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -436,8 +465,8 @@ def declare_table(self, request: Any) -> Any: assert all(request.vend_credentials is True for request in requests) -def test_sink_empty_storage_options_falls_back_to_io_config(tmp_path: Path) -> None: - """storage_options={} must behave like the read entry points: fall through to io_config.""" +def test_sink_empty_storage_options_remain_explicit_for_uri(tmp_path: Path) -> None: + """The URI sink historically treats storage_options={} as explicitly empty.""" from daft.io import IOConfig, S3Config io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) @@ -445,8 +474,7 @@ def test_sink_empty_storage_options_falls_back_to_io_config(tmp_path: Path) -> N sink = LanceDataSink("s3://bucket/t.lance", schema, "create", io_config, storage_options={}) merged = sink._merged_storage_options(namespace_mod.ResolvedNamespaceTable(uri="s3://bucket/t.lance")) - assert merged is not None - assert merged["access_key_id"] == "io-key" + assert merged == {} def test_construct_lance_dataset_storage_options_priority(monkeypatch: pytest.MonkeyPatch) -> None: @@ -471,6 +499,7 @@ def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **k lambda **kwargs: namespace_mod.ResolvedNamespaceTable( uri="s3://bucket/t.lance", storage_options={"access_key_id": "vended-key", "session_token": "vended-token"}, + managed_versioning=True, ), ) @@ -493,6 +522,7 @@ def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **k assert captured["uri"] is None # namespace addressing passes uri=None assert handle.storage_options == merged assert handle.uri == "s3://bucket/t.lance" + assert handle.managed_versioning is True assert not hasattr(handle.dataset, "_lance_open_kwargs") diff --git a/tests/io/lancedb/test_namespace_rest_e2e.py b/tests/io/lancedb/test_namespace_rest_e2e.py index 2445cc3..8dafb6a 100644 --- a/tests/io/lancedb/test_namespace_rest_e2e.py +++ b/tests/io/lancedb/test_namespace_rest_e2e.py @@ -15,13 +15,32 @@ ) -def test_rest_namespace_write_read_append_roundtrip() -> None: +def test_rest_namespace_write_read_append_roundtrip(monkeypatch: pytest.MonkeyPatch) -> None: import lance import lance_namespace as ln from lance_namespace import CreateNamespaceRequest, DescribeTableRequest, NamespaceExistsRequest from lance_namespace.errors import TableAlreadyExistsError namespace_properties = {"uri": os.environ["DAFT_LANCE_REST_URI"]} + + # Ensure object-store access in this test cannot silently succeed through + # ambient AWS credentials instead of credentials vended by the catalog. + for name in ( + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_ROLE_ARN", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + "AWS_CONTAINER_AUTHORIZATION_TOKEN", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", os.devnull) + monkeypatch.setenv("AWS_CONFIG_FILE", os.devnull) catalog = os.environ.get("DAFT_LANCE_REST_CATALOG", "lance_catalog") schema = os.environ.get("DAFT_LANCE_REST_SCHEMA", "daft_ns_e2e") table_id = [catalog, schema, f"orders_{uuid.uuid4().hex[:8]}"] @@ -47,9 +66,10 @@ def test_rest_namespace_write_read_append_roundtrip() -> None: **ns, ).collect() - describe = namespace.describe_table(DescribeTableRequest(id=table_id)) + describe = namespace.describe_table(DescribeTableRequest(id=table_id, vend_credentials=True)) location = getattr(describe, "location", None) or getattr(describe, "table_uri", None) assert location + assert describe.storage_options, "catalog must vend storage options for this integration test" result = daft_lance.read_lance(table_id=table_id, **ns).to_pydict() assert result == { From 2b97eb1de0321fd325da2f5f18bda80252029522 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Tue, 21 Jul 2026 12:49:34 +0900 Subject: [PATCH 19/25] refactor(namespace): tighten integration types --- daft_lance/_lance.py | 5 +++-- daft_lance/namespace.py | 19 ++++++++++++------- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index a58cc5d..ff682ea 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -8,8 +8,10 @@ from daft.api_annotations import PublicAPI from daft.daft import IOConfig, ScanOperatorHandle from daft.dataframe import DataFrame +from daft.dependencies import pa from daft.io._checkpoint import attach_checkpoint from daft.logical.builder import LogicalPlanBuilder +from daft.schema import Schema from .lance_compaction import compact_files_internal from .lance_data_sink import LanceDataSink @@ -24,7 +26,6 @@ from lance.udf import BatchUDF from daft.checkpoint import CheckpointConfig - from daft.dependencies import pa @PublicAPI @@ -639,7 +640,7 @@ def write_lance( uri: str | pathlib.Path | None = None, mode: Literal["create", "append", "overwrite"] = "create", io_config: IOConfig | None = None, - schema: Any | None = None, + schema: Schema | pa.Schema | None = None, *, table_id: list[str] | None = None, namespace_impl: str | None = None, diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index 0b3a713..6105c45 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -25,6 +25,7 @@ from urllib.parse import urlparse import lance +from lance_namespace import DeclareTableResponse, DescribeTableResponse, LanceNamespace _NAMESPACE_CACHE_SIZE = int(os.environ.get("DAFT_LANCE_NAMESPACE_CACHE_SIZE", "16")) @@ -73,14 +74,18 @@ def _normalize_file_uri(location: str) -> str: @lru_cache(maxsize=_NAMESPACE_CACHE_SIZE) -def _get_cached_namespace(namespace_impl: str, namespace_properties_tuple: tuple[tuple[str, str], ...] | None) -> Any: +def _get_cached_namespace( + namespace_impl: str, namespace_properties_tuple: tuple[tuple[str, str], ...] | None +) -> LanceNamespace: import lance_namespace as ln namespace_properties = dict(namespace_properties_tuple) if namespace_properties_tuple else {} return ln.connect(namespace_impl, namespace_properties) -def get_or_create_namespace(namespace_impl: str | None, namespace_properties: dict[str, str] | None) -> Any | None: +def get_or_create_namespace( + namespace_impl: str | None, namespace_properties: dict[str, str] | None +) -> LanceNamespace | None: """Per-process namespace client pool. ``ln.connect`` may build HTTP clients / perform auth, and workers re-derive @@ -130,14 +135,14 @@ def get_namespace_commit_kwargs( get_write_fragments_kwargs = get_namespace_kwargs -def _storage_options(response: Any) -> dict[str, str] | None: +def _storage_options(response: DescribeTableResponse | DeclareTableResponse) -> dict[str, str] | None: storage_options = getattr(response, "storage_options", None) if storage_options is None: return None return dict(storage_options) -def _response_location(response: Any) -> str: +def _response_location(response: DescribeTableResponse | DeclareTableResponse) -> str: location = getattr(response, "location", None) or getattr(response, "table_uri", None) if not location: raise ValueError("Namespace response did not include a table location.") @@ -153,7 +158,7 @@ class ResolvedNamespaceTable: managed_versioning: bool = False -def _resolved_from_response(response: Any) -> ResolvedNamespaceTable: +def _resolved_from_response(response: DescribeTableResponse | DeclareTableResponse) -> ResolvedNamespaceTable: return ResolvedNamespaceTable( uri=_response_location(response), storage_options=_storage_options(response), @@ -161,7 +166,7 @@ def _resolved_from_response(response: Any) -> ResolvedNamespaceTable: ) -def _describe_table(namespace: Any, table_id: list[str]) -> Any: +def _describe_table(namespace: LanceNamespace, table_id: list[str]) -> DescribeTableResponse: from lance_namespace import DescribeTableRequest # vend_credentials is explicit: when unset, whether the namespace returns @@ -169,7 +174,7 @@ def _describe_table(namespace: Any, table_id: list[str]) -> Any: return namespace.describe_table(DescribeTableRequest(id=table_id, vend_credentials=True)) -def _declare_table(namespace: Any, table_id: list[str]) -> ResolvedNamespaceTable: +def _declare_table(namespace: LanceNamespace, table_id: list[str]) -> ResolvedNamespaceTable: from lance_namespace import DeclareTableRequest response = namespace.declare_table(DeclareTableRequest(id=table_id, location=None, vend_credentials=True)) From 9db4115d583186f5f7f874741a4be3528d32dd40 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Wed, 22 Jul 2026 08:45:09 +0900 Subject: [PATCH 20/25] fix(namespace): percent-decode file:// locations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Namespaces vend URI-encoded locations: a table under `daft lance/` comes back from the dir namespace as `file:///.../daft%20lance/t.lance`. _normalize_file_uri took parsed.path verbatim, so writes landed in a literally-named `daft%20lance` directory while the subsequent describe_table resolved to the decoded one — create appeared to succeed and reads found nothing. Claude-Session: https://claude.ai/code/session_018YLUk2vAhX3jXFL9D7za7u --- daft_lance/namespace.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index 6105c45..82c2739 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -22,7 +22,7 @@ from dataclasses import dataclass from functools import lru_cache from typing import Any -from urllib.parse import urlparse +from urllib.parse import unquote, urlparse import lance from lance_namespace import DeclareTableResponse, DescribeTableResponse, LanceNamespace @@ -66,10 +66,15 @@ def _normalize_file_uri(location: str) -> str: ``file:///...``), but the location is later used where a plain path is required: ``write_fragments(dataset_uri=...)``, ``LanceDataset.commit``, and ``pathlib`` checks in the sink. Object-store URIs pass through as-is. + + The path is percent-decoded: namespaces vend URI-encoded locations (a table + under ``daft lance/`` comes back as ``daft%20lance/``), and writing to the + literal encoded form would create a different directory than the one + ``describe_table`` later resolves to. """ parsed = urlparse(location) if parsed.scheme == "file": - return parsed.path + return unquote(parsed.path) return location From 4596b22a380676bbdfd09e06ece29620540c90ac Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Wed, 22 Jul 2026 08:45:20 +0900 Subject: [PATCH 21/25] refactor(utils): drop construct_lance_dataset compat wrapper The wrapper returned only .dataset, discarding the resolved uri, open_kwargs, managed_versioning and default_scan_options that construct_lance_dataset_handle resolves. Callers going through it lost the reusable open context that distributed workers need to reopen the dataset, and nearest-vector scan defaults silently became None. It had no callers inside daft_lance and was never exported, so remove it rather than reconstructing the lost state from dataset private attributes. This also keeps a single explicit context type as the foundation for the serializable access context that credential refresh will need. Claude-Session: https://claude.ai/code/session_018YLUk2vAhX3jXFL9D7za7u --- daft_lance/utils.py | 28 ---------------------------- 1 file changed, 28 deletions(-) diff --git a/daft_lance/utils.py b/daft_lance/utils.py index dc6cd34..f23025b 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -213,34 +213,6 @@ def construct_lance_dataset_handle( ) -def construct_lance_dataset( - uri: str | pathlib.Path | None, - version: int | str | None = None, - storage_options: dict[str, Any] | None = None, - io_config: IOConfig | None = None, - namespace_impl: str | None = None, - namespace_properties: dict[str, str] | None = None, - table_id: list[str] | None = None, - **kwargs: Any, -) -> lance.LanceDataset: - """Construct a Lance dataset with common options. - - This compatibility wrapper preserves the historical return type. Internal - callers that also need the resolved URI or reusable open arguments should - use :func:`construct_lance_dataset_handle`. - """ - return construct_lance_dataset_handle( - uri, - version=version, - storage_options=storage_options, - io_config=io_config, - namespace_impl=namespace_impl, - namespace_properties=namespace_properties, - table_id=table_id, - **kwargs, - ).dataset - - def combine_filters_to_arrow(predicates: list[Any] | None) -> pa.compute.Expression | None: """Combine a list of Daft PyExpr predicates into a single Arrow Expression. From 8af22c75c995fd9996c114f51e51efaca6085bf5 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Wed, 22 Jul 2026 08:45:20 +0900 Subject: [PATCH 22/25] feat(sink): reject namespace writes combined with use_mem_wal The namespace write path declares the table up front as a metadata-only reservation, then the mem-WAL path asked Lance for a namespace-aware write_dataset(mode="create"), which declared the same table again and raised TableAlreadyExistsError. The combination could never succeed, and _ensure_mem_wal_dataset's except clause does not cover the namespace error types either. Fail at construction time with an actionable message instead of claiming support. Making it work requires the mem-WAL path to skip our declare and let the native create own table creation, which is left to a follow-up. Claude-Session: https://claude.ai/code/session_018YLUk2vAhX3jXFL9D7za7u --- daft_lance/lance_data_sink.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index 946cd3c..dcc2cc0 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -69,6 +69,7 @@ def __init__( compact_after_write: bool = True, ) -> None: self._reject_unsupported_modes(mode, use_legacy_format) + self._reject_namespace_mem_wal(namespace_impl, table_id, use_mem_wal) validate_uri_or_namespace(uri, namespace_impl, table_id, namespace_properties) if uri is not None and not isinstance(uri, (str, pathlib.Path)): raise TypeError(f"Expected URI to be str or pathlib.Path, got {type(uri)}") @@ -206,6 +207,24 @@ def _reject_unsupported_modes( stacklevel=3, ) + @staticmethod + def _reject_namespace_mem_wal(namespace_impl: str | None, table_id: list[str] | None, use_mem_wal: bool) -> None: + """Reject the namespace + mem-WAL combination instead of failing mid-write. + + The namespace write path declares the table up front (a metadata-only + reservation), but the mem-WAL path then asks Lance for a namespace-aware + ``write_dataset(mode="create")``, which declares the same table a second + time and raises ``TableAlreadyExistsError``. Supporting this needs the + mem-WAL path to skip our own declare and let the native create own it; + until then, fail loudly at construction time. + """ + if use_mem_wal and namespace_impl is not None and table_id is not None: + raise ValueError( + "use_mem_wal=True is not supported with namespace-addressed tables " + "('namespace_impl' + 'table_id'). Write to a 'uri' instead, or set " + "use_mem_wal=False." + ) + def _init_lance_knobs( self, *, From 0cbe59a859532582273f2b88755c795c803857a5 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Wed, 22 Jul 2026 08:45:20 +0900 Subject: [PATCH 23/25] test(namespace): cover uri decoding, handle API and mem-WAL rejection Adds unit coverage for _normalize_file_uri (spaces, non-ASCII, literal percent, object-store passthrough) and an end-to-end roundtrip through a namespace root and table id that both percent-encode, asserting data lands under the decoded name. Also asserts namespace + use_mem_wal is rejected, and moves the two construct_lance_dataset call sites to construct_lance_dataset_handle. Claude-Session: https://claude.ai/code/session_018YLUk2vAhX3jXFL9D7za7u --- tests/io/lancedb/test_namespace.py | 56 ++++++++++++++++++++++++++++-- 1 file changed, 53 insertions(+), 3 deletions(-) diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index 2a939e1..c149d9f 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -23,6 +23,43 @@ def _double_score(batch: Any) -> Any: return pa.RecordBatch.from_arrays([pc.multiply(batch["score"], 2)], ["doubled"]) +@pytest.mark.parametrize( + ("location", "expected"), + [ + ("file:///tmp/plain/t.lance", "/tmp/plain/t.lance"), + ("file:///tmp/root%20space/t.lance", "/tmp/root space/t.lance"), + ("file:///tmp/%E4%B8%AD%E6%96%87/t.lance", "/tmp/中文/t.lance"), + ("file:///tmp/100%25/t.lance", "/tmp/100%/t.lance"), + # Object-store URIs pass through untouched, encoding included. + ("s3://bucket/root%20space/t.lance", "s3://bucket/root%20space/t.lance"), + ], +) +def test_normalize_file_uri_percent_decodes_paths(location: str, expected: str) -> None: + assert namespace_mod._normalize_file_uri(location) == expected + + +@pytest.mark.parametrize("root_name", ["daft lance", "中文目录"]) +def test_namespace_roundtrip_with_encoded_location(tmp_path: Path, root_name: str) -> None: + """A namespace root/table whose location percent-encodes must still round-trip. + + The dir namespace vends ``file://.../daft%20lance/table%20space.lance``; writing + to the literal encoded form would create a directory that the later + ``describe_table`` never resolves to. + """ + root = tmp_path / root_name + root.mkdir() + ns = _dir_ns(root) + table_id = ["table space"] + + df = daft.from_pydict({"id": [1, 2], "label": ["a", "b"]}) + daft_lance.write_lance(df, table_id=table_id, mode="create", **ns).collect() + + assert daft_lance.read_lance(table_id=table_id, **ns).to_pydict() == {"id": [1, 2], "label": ["a", "b"]} + # The data landed under the decoded name, not a literal "%20" sibling. + assert (root / "table space.lance").is_dir() + assert not (root / "table%20space.lance").exists() + + def test_namespace_write_read_append_roundtrip(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) table_id = ["roundtrip"] @@ -386,6 +423,19 @@ def test_namespace_commit_kwargs_include_managed_versioning(monkeypatch: pytest. } +def test_namespace_with_mem_wal_is_rejected(tmp_path: Path) -> None: + schema = daft.from_pydict({"id": [1]}).schema() + with pytest.raises(ValueError, match="use_mem_wal=True is not supported"): + LanceDataSink( + uri=None, + schema=schema, + mode="create", + table_id=["memwal"], + use_mem_wal=True, + **_dir_ns(tmp_path), + ) + + def test_construct_lance_dataset_empty_storage_options_falls_back_to_io_config( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -404,10 +454,10 @@ def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **k monkeypatch.setattr("daft_lance.utils.lance.dataset", fake_dataset) io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) - dataset = utils_mod.construct_lance_dataset("s3://bucket/t.lance", storage_options={}, io_config=io_config) + handle = utils_mod.construct_lance_dataset_handle("s3://bucket/t.lance", storage_options={}, io_config=io_config) merged = captured["storage_options"] - assert isinstance(dataset, FakeDataset) + assert isinstance(handle.dataset, FakeDataset) assert merged is not None assert merged["access_key_id"] == "io-key" @@ -548,7 +598,7 @@ def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **k ) io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) - utils_mod.construct_lance_dataset( + utils_mod.construct_lance_dataset_handle( None, io_config=io_config, namespace_impl="rest", From d1424f36403c700b192721d94396970e2e452172 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Wed, 22 Jul 2026 14:44:41 +0900 Subject: [PATCH 24/25] refactor(namespace): reopen datasets on workers via DatasetOpenContext LanceDataset.__reduce__ carries only (uri, storage_options, version, manifest, ...); _namespace_client, _table_id and _namespace_client_managed_versioning are assigned after construction and are dropped by pickle. Compaction, scalar index and merge all shipped the driver's live dataset into their UDFs, so workers silently lost the table's namespace identity and committed as if it were a uri table. Introduce a frozen, serializable DatasetOpenContext carrying the physical uri, the driver's resolved numeric version, the effective initial storage options, the namespace triple, managed_versioning and the worker read options. Workers rebuild the namespace client per process and reopen; the driver keeps its live dataset for planning, validation and commits. Workers open through the low-level LanceDataset constructor with the physical uri rather than lance.dataset(None, namespace_client=...), which resolves the location with a describe_table on every call. Measured with pylance's ops_metrics: the low-level open costs zero namespace calls, the high-level one costs a round-trip per task. Because the context supplies uri, storage options and namespace kwargs, the internal entry points drop those four parameters instead of carrying a second copy of the same state; FastPathFragmentWriter loses its private uri/storage_options fields for the same reason. Each UDF instance opens once, lazily -- the reopen costs a pinned-manifest read, so it must not sit on the per-call path. The serialized manifest is deliberately not carried, keeping the task payload independent of fragment count at the cost of that manifest read. Version semantics are preserved: workers pin the snapshot the driver planned against, and only the index coordinator steps that must observe worker output reopen at latest. Known limitation: workers start from the storage options the driver resolved. pylance installs its namespace refresh provider only when those options are present, and refresh fires on expires_at_millis, so tables whose catalog vends no expiry metadata will not refresh mid-task. Full credential refresh is #53. Mem-WAL still rejects namespace writes (#54) and builds a uri-only context. Scan keeps its existing open_kwargs path untouched. Claude-Session: https://claude.ai/code/session_01Ss3gSAbK3FwfPmWqk1CSbJ --- daft_lance/_lance.py | 14 +- daft_lance/lance_compaction.py | 28 ++- daft_lance/lance_data_sink.py | 14 +- daft_lance/lance_merge_column.py | 109 +++++------ daft_lance/lance_scalar_index.py | 81 ++++---- daft_lance/namespace.py | 133 ++++++++++++- daft_lance/utils.py | 24 +++ tests/io/lancedb/test_fast_path_merge.py | 81 ++++---- tests/io/lancedb/test_lancedb_scalar_index.py | 31 ++- tests/io/lancedb/test_namespace.py | 176 ++++++++++++++++++ 10 files changed, 532 insertions(+), 159 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index ff682ea..0ad2231 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -265,12 +265,10 @@ def merge_columns( return merge_columns_internal( dataset_handle.dataset, - dataset_handle.uri, + dataset_handle.worker_open_context(), transform=transform, read_columns=read_columns, reader_schema=reader_schema, - storage_options=dataset_handle.storage_options, - namespace_kwargs=dataset_handle.commit_kwargs, daft_remote_args=daft_remote_args, concurrency=concurrency, ) @@ -380,11 +378,9 @@ def merge_columns_df( return merge_columns_from_df( df, lance_ds=dataset_handle.dataset, - uri=dataset_handle.uri, + open_context=dataset_handle.worker_open_context(), read_columns=read_columns, reader_schema=reader_schema, - storage_options=dataset_handle.storage_options, - namespace_kwargs=dataset_handle.commit_kwargs, daft_remote_args=daft_remote_args, concurrency=concurrency, left_on=left_on, @@ -530,14 +526,11 @@ def create_scalar_index( create_scalar_index_internal( lance_ds=dataset_handle.dataset, - uri=dataset_handle.uri, + open_context=dataset_handle.worker_open_context(), column=column, index_type=index_type, name=name, replace=replace, - storage_options=dataset_handle.storage_options, - namespace_kwargs=dataset_handle.namespace_kwargs, - namespace_commit_kwargs=dataset_handle.commit_kwargs, fragment_group_size=fragment_group_size, num_partitions=num_partitions, max_concurrency=max_concurrency, @@ -628,6 +621,7 @@ def compact_files( return compact_files_internal( lance_ds=dataset_handle.dataset, + open_context=dataset_handle.worker_open_context(), compaction_options=compaction_options, partition_num=partition_num, concurrency=concurrency, diff --git a/daft_lance/lance_compaction.py b/daft_lance/lance_compaction.py index 4575cb7..2ca2b41 100644 --- a/daft_lance/lance_compaction.py +++ b/daft_lance/lance_compaction.py @@ -1,13 +1,16 @@ from __future__ import annotations import logging -from typing import Any +from typing import TYPE_CHECKING, Any from lance import LanceDataset from lance.optimize import Compaction, CompactionMetrics, CompactionOptions, CompactionTask, RewriteResult import daft +if TYPE_CHECKING: + from daft_lance.namespace import DatasetOpenContext + logger = logging.getLogger(__name__) @@ -16,23 +19,36 @@ class CompactionTaskUDF: def __init__( self, - lance_ds: LanceDataset, + open_context: DatasetOpenContext, ) -> None: - self.lance_ds = lance_ds + self.open_context = open_context + self._lance_ds: LanceDataset | None = None + + def _dataset(self) -> LanceDataset: + # Opened once per UDF instance, not per task: the reopen costs a pinned + # manifest read, so it must not sit on the per-row path. + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds def __call__(self, task: CompactionTask) -> RewriteResult: - rewrite = task.execute(self.lance_ds) + rewrite = task.execute(self._dataset()) return rewrite def compact_files_internal( lance_ds: LanceDataset, + open_context: DatasetOpenContext, *, compaction_options: dict[str, Any] | None = None, partition_num: int | None = None, concurrency: int | None = None, ) -> CompactionMetrics | None: - """Execute Lance file compaction in distributed environment using Daft UDF style.""" + """Execute Lance file compaction in distributed environment using Daft UDF style. + + ``lance_ds`` is the driver's live dataset and stays on the driver for + planning and the final commit; ``open_context`` is what workers reopen from. + """ logger.info("Starting UDF-style distributed compaction") plan = Compaction.plan( lance_ds, @@ -59,7 +75,7 @@ def compact_files_internal( CompactionTaskUDF, max_concurrency=concurrency, ) - df = df.select(WrappedRunner(lance_ds)(df["task"]).alias("rewrite")) + df = df.select(WrappedRunner(open_context)(df["task"]).alias("rewrite")) results = df.to_pandas() metrics = Compaction.commit(lance_ds, results["rewrite"].to_list()) diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index dcc2cc0..3feef7e 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -467,8 +467,18 @@ def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadat self._mem_wal_total_bytes, ) from daft_lance.lance_compaction import compact_files_internal - - compact_files_internal(dataset) + from daft_lance.namespace import DatasetOpenContext + + # Mem-WAL rejects namespace addressing up front (issue #54), so this + # is always a uri-only table and the context carries no triple. + compact_files_internal( + dataset, + DatasetOpenContext.from_dataset( + dataset, + str(self._dataset_uri_arg), + storage_options=self._storage_options, + ), + ) dataset = lance.dataset( self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs ) diff --git a/daft_lance/lance_merge_column.py b/daft_lance/lance_merge_column.py index 44f44ea..6546459 100644 --- a/daft_lance/lance_merge_column.py +++ b/daft_lance/lance_merge_column.py @@ -15,10 +15,10 @@ from daft.udf import method if TYPE_CHECKING: - import pathlib from collections.abc import Callable from daft.dependencies import pa + from daft_lance.namespace import DatasetOpenContext _FRAGMENT_HANDLER_RETURN_DTYPE = DataType.struct({"fragment_meta": DataType.binary(), "schema": DataType.binary()}) @@ -28,21 +28,28 @@ class FragmentHandler: def __init__( self, - lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, transform: dict[str, str] | lance.udf.BatchUDF | Callable[[pa.lib.RecordBatch], pa.lib.RecordBatch], read_columns: list[str] | None, reader_schema: pa.Schema | None = None, ): - self.lance_ds = lance_ds + self.open_context = open_context self.transform = transform self.read_columns = read_columns self.reader_schema = reader_schema + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds @method.batch(return_dtype=_FRAGMENT_HANDLER_RETURN_DTYPE) def __call__(self, fragment_ids: Any) -> list[dict[str, bytes]]: results = [] + lance_ds = self._dataset() for fragment_id in fragment_ids: - fragment = self.lance_ds.get_fragment(fragment_id) + fragment = lance_ds.get_fragment(fragment_id) if fragment is None: raise ValueError(f"Fragment {fragment_id} not found in dataset") fragment_meta, schema = fragment.merge_columns(self.transform, self.read_columns, None, self.reader_schema) @@ -52,13 +59,11 @@ def __call__(self, fragment_ids: Any) -> list[dict[str, bytes]]: def merge_columns_internal( lance_ds: lance.LanceDataset, - url: str | pathlib.Path, + open_context: DatasetOpenContext, *, transform: dict[str, str] | lance.udf.BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch], read_columns: list[str] | None = None, reader_schema: pa.Schema | None = None, - storage_options: dict[str, Any] | None = None, - namespace_kwargs: dict[str, Any] | None = None, daft_remote_args: dict[str, Any] | None = None, concurrency: int | None = None, ) -> lance.LanceDataset: @@ -73,7 +78,7 @@ def merge_columns_internal( # Instantiate the Daft class with Lance-specific state and apply the # batch method over the fragment_id column. - handler = FragmentHandler(lance_ds, transform, read_columns, reader_schema) + handler = FragmentHandler(open_context, transform, read_columns, reader_schema) df = df.with_column("commit_message", handler(df["fragment_id"])) # type: ignore[arg-type] commit_messages = df.collect().to_pydict()["commit_message"] @@ -90,11 +95,11 @@ def merge_columns_internal( raise ValueError("No schema for new fragment found") op = lance.LanceOperation.Merge(fragment_metas, new_schema) return lance_ds.commit( - url, + open_context.uri, op, read_version=lance_ds.version, - storage_options=storage_options, - **(namespace_kwargs or {}), + storage_options=open_context.storage_options, + **open_context.commit_kwargs, ) @@ -102,7 +107,7 @@ def merge_columns_internal( class GroupFragmentMergeUDF: def __init__( self, - lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, left_on: str | None = "_rowaddr", right_on: str | None = None, read_columns: list[str] | None = None, @@ -112,19 +117,25 @@ def __init__( """Per-group merge handler that directly invokes Lance fragment.merge with keyed join. Args: - lance_ds: Target Lance dataset. + open_context: Serializable handle the worker reopens the target dataset from. left_on: Key column on the Lance fragment (default "_rowaddr"). right_on: Key column name present in the provided reader data (defaults to left_on). read_columns: Names for columns provided to the handler via map_groups (must include right_on). reader_schema: Optional Arrow schema for the reader. batch_size: Optional batch size when building RecordBatchReader from the provided data. """ - self.lance_ds = lance_ds + self.open_context = open_context self.left_on = left_on or "_rowaddr" self.right_on = right_on or self.left_on self.read_columns = read_columns or [] self.reader_schema = reader_schema self.batch_size = batch_size + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds @method.batch(return_dtype=_FRAGMENT_HANDLER_RETURN_DTYPE) def __call__(self, *cols: Any) -> list[dict[str, bytes]]: @@ -194,16 +205,17 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: # Enforce that reader stream contains only join key + new columns (exclude existing dataset fields) df_schema = tbl.schema + lance_ds = self._dataset() existing_fields: set[str] = set() try: - existing_fields = {getattr(f, "name", str(f)) for f in self.lance_ds.schema} + existing_fields = {getattr(f, "name", str(f)) for f in lance_ds.schema} except Exception: names = [] try: - names = list(getattr(self.lance_ds.schema, "names", [])) + names = list(getattr(lance_ds.schema, "names", [])) except Exception: try: - names = [getattr(f, "name", str(f)) for f in getattr(self.lance_ds.schema, "fields", [])] + names = [getattr(f, "name", str(f)) for f in getattr(lance_ds.schema, "fields", [])] except Exception: names = [] existing_fields = set(names) @@ -221,7 +233,7 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: batches = tbl.to_batches(max_chunksize=self.batch_size) if self.batch_size is not None else tbl.to_batches() reader = _pa.RecordBatchReader.from_batches(tbl.schema, batches) - fragment = self.lance_ds.get_fragment(frag_id) + fragment = lance_ds.get_fragment(frag_id) if fragment is None: raise ValueError(f"Fragment {frag_id} not found in dataset") # Build schema argument: use the table's schema (including join key and new columns) unless an explicit reader_schema is provided @@ -241,15 +253,17 @@ class FastPathFragmentWriter: def __init__( self, - lance_ds: lance.LanceDataset, - uri: str, + open_context: DatasetOpenContext, new_column_names: list[str], - storage_options: dict[str, str] | None = None, ): - self.lance_ds = lance_ds - self.uri = str(uri) + self.open_context = open_context self.new_column_names = new_column_names - self.storage_options = storage_options + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds @method.batch(return_dtype=_FRAGMENT_HANDLER_RETURN_DTYPE) def __call__(self, *cols: Any) -> list[dict[str, bytes]]: @@ -293,7 +307,8 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: # Determine the existing file format version so the new file matches. # Lance commit rejects fragments whose files mix major/minor versions. - fragment = self.lance_ds.get_fragment(frag_id) + lance_ds = self._dataset() + fragment = lance_ds.get_fragment(frag_id) if fragment is None: raise ValueError(f"Fragment {frag_id} not found in dataset") meta = dict(fragment.metadata.to_json()) @@ -305,12 +320,12 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: # Write raw .lance file with only new columns filename = uuid.uuid4().hex + ".lance" - filepath = os.path.join(self.uri, "data", filename) + filepath = os.path.join(self.open_context.uri, "data", filename) with LanceFileWriter( filepath, tbl.schema, version=f"{file_major}.{file_minor}", - storage_options=self.storage_options, + storage_options=self.open_context.storage_options, ) as writer: for b in tbl.to_batches(): writer.write_batch(b) @@ -319,7 +334,7 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: # Determine field IDs for the new columns. Lance's manifest-level # max_field_id includes nested child fields and field IDs from dropped # columns, so it is the correct high-water mark for dataset evolution. - next_fid = self.lance_ds.max_field_id + 1 + next_fid = lance_ds.max_field_id + 1 # Stitch new data file into fragment metadata new_file_entry = { @@ -335,7 +350,7 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: new_frag_meta = FragmentMetadata.from_json(json.dumps(meta)) # Build new schema (original + new columns) - new_schema = self.lance_ds.schema + new_schema = lance_ds.schema for col_name in self.new_column_names: col_idx = tbl.schema.get_field_index(col_name) new_schema = new_schema.append(_pa.field(col_name, tbl.schema.field(col_idx).type)) @@ -366,12 +381,10 @@ def _can_use_fast_path( def merge_columns_from_df( df: daft.DataFrame, lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, *, read_columns: list[str] | None = None, reader_schema: pa.Schema | None = None, - storage_options: dict[str, Any] | None = None, - namespace_kwargs: dict[str, Any] | None = None, daft_remote_args: dict[str, Any] | None = None, concurrency: int | None = None, left_on: str | None = "_rowaddr", @@ -419,36 +432,30 @@ def merge_columns_from_df( return _merge_fast_path( df, lance_ds, - uri, + open_context, new_cols, - storage_options=storage_options, - namespace_kwargs=namespace_kwargs, ) else: return _merge_slow_path( df, lance_ds, - uri, + open_context, read_columns, left_on, right_on, reader_schema, batch_size, - storage_options=storage_options, - namespace_kwargs=namespace_kwargs, ) def _merge_fast_path( df: daft.DataFrame, lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, new_column_names: list[str], - storage_options: dict[str, Any] | None = None, - namespace_kwargs: dict[str, Any] | None = None, ) -> lance.LanceDataset: """Metadata-only add_columns: write raw .lance files and stitch into fragment metadata.""" - handler = FastPathFragmentWriter(lance_ds, str(uri), new_column_names, storage_options=storage_options) + handler = FastPathFragmentWriter(open_context, new_column_names) grouped = df.groupby("fragment_id").map_groups( handler(*(df[c] for c in new_column_names), df["_rowaddr"], df["fragment_id"]).alias("commit_message") # type: ignore[attr-defined] @@ -481,29 +488,27 @@ def _merge_fast_path( op = lance.LanceOperation.Merge(fragment_metas, LanceSchema.from_pyarrow(new_schema)) return lance.LanceDataset.commit( - str(uri), + open_context.uri, op, read_version=lance_ds.version, - storage_options=storage_options, - **(namespace_kwargs or {}), + storage_options=open_context.storage_options, + **open_context.commit_kwargs, ) def _merge_slow_path( df: daft.DataFrame, lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, read_columns: list[str], left_on: str | None, right_on: str | None, reader_schema: pa.Schema | None, batch_size: int | None, - storage_options: dict[str, Any] | None = None, - namespace_kwargs: dict[str, Any] | None = None, ) -> lance.LanceDataset: """Original keyed-join merge path: rewrites fragment data.""" handler_udf = GroupFragmentMergeUDF( - lance_ds, + open_context, left_on, right_on, read_columns, @@ -531,9 +536,9 @@ def _merge_slow_path( return lance_ds op = lance.LanceOperation.Merge(fragment_metas, new_schema) return lance_ds.commit( - uri, + open_context.uri, op, read_version=lance_ds.version, - storage_options=storage_options, - **(namespace_kwargs or {}), + storage_options=open_context.storage_options, + **open_context.commit_kwargs, ) diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index 298bc40..d29fd23 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -9,7 +9,7 @@ from daft import execution_config_ctx, from_pylist if TYPE_CHECKING: - import pathlib + from daft_lance.namespace import DatasetOpenContext import lance @@ -27,7 +27,7 @@ class FragmentIndexHandler: def __init__( self, - lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, column: str, index_type: str, name: str, @@ -35,13 +35,19 @@ def __init__( replace: bool, **kwargs: Any, ) -> None: - self.lance_ds = lance_ds + self.open_context = open_context self.column = column self.index_type = index_type self.name = name self.fragment_uuid = fragment_uuid self.replace = replace self.kwargs = kwargs + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds def __call__(self, fragment_ids: list[int]) -> bool: """Process a batch of fragment IDs for scalar index creation.""" @@ -50,7 +56,7 @@ def __call__(self, fragment_ids: list[int]) -> bool: fragment_ids, ) - self.lance_ds.create_scalar_index( + self._dataset().create_scalar_index( column=self.column, index_type=self.index_type, # type: ignore[arg-type] name=self.name, @@ -74,17 +80,23 @@ class SegmentedFragmentIndexHandler: def __init__( self, - lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, column: str, index_type: str, name: str, **kwargs: Any, ) -> None: - self.lance_ds = lance_ds + self.open_context = open_context self.column = column self.index_type = index_type self.name = name self.kwargs = kwargs + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds def __call__(self, fragment_ids: list[int], shard_id: int | None = None) -> bytes: """Build an independent index segment and return its pickled metadata.""" @@ -104,7 +116,7 @@ def __call__(self, fragment_ids: list[int], shard_id: int | None = None) -> byte # scalar index segments through this public API. Segment creation always # uses ``replace=False`` because replacement, if supported, must happen # in the final manifest commit rather than independently in each worker. - index_meta = self.lance_ds.create_index_uncommitted( + index_meta = self._dataset().create_index_uncommitted( column=self.column, index_type=self.index_type, name=self.name, @@ -132,15 +144,12 @@ def _existing_index_names(lance_ds: lance.LanceDataset) -> set[str]: def create_scalar_index_internal( lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, *, column: str, index_type: str = "INVERTED", name: str | None = None, replace: bool = False, - storage_options: dict[str, Any] | None = None, - namespace_kwargs: dict[str, Any] | None = None, - namespace_commit_kwargs: dict[str, Any] | None = None, fragment_group_size: int | None = None, num_partitions: int | None = None, max_concurrency: int | None = None, @@ -149,6 +158,10 @@ def create_scalar_index_internal( ) -> None: """Internal implementation of distributed scalar index creation. + ``lance_ds`` is the driver's live dataset (planning, validation, commits); + ``open_context`` is the serializable handle workers reopen from and the + single source of uri, storage options and namespace kwargs. + When ``segmented=True``, ``BITMAP``, ``BTREE``, and ``INVERTED`` use Lance's public segment-index workflow: each worker builds a fully independent index segment, and the coordinator commits them atomically with @@ -274,7 +287,7 @@ def create_scalar_index_internal( # Configure maximum concurrency for fragment batches if not fragment_data: - logger.info("No fragments found for dataset at %s; skipping scalar index creation.", uri) + logger.info("No fragments found for dataset at %s; skipping scalar index creation.", open_context.uri) return logger.info( @@ -292,13 +305,10 @@ def create_scalar_index_internal( # older/unsupported distributed scalar index types. if segmented: _create_segmented_index( - lance_ds=lance_ds, - uri=uri, + open_context=open_context, column=column, index_type=index_type, name=name, - storage_options=storage_options, - namespace_kwargs=namespace_kwargs, fragment_data=fragment_data, fragment_ids_to_use=fragment_ids_to_use, num_partitions=num_partitions, @@ -307,15 +317,11 @@ def create_scalar_index_internal( ) else: _create_partitioned_index( - lance_ds=lance_ds, - uri=uri, + open_context=open_context, column=column, index_type=index_type, name=name, replace=replace, - storage_options=storage_options, - namespace_kwargs=namespace_kwargs, - namespace_commit_kwargs=namespace_commit_kwargs, fragment_data=fragment_data, fragment_ids_to_use=fragment_ids_to_use, num_partitions=num_partitions, @@ -325,14 +331,11 @@ def create_scalar_index_internal( def _create_segmented_index( - lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, *, column: str, index_type: str, name: str, - storage_options: dict[str, Any] | None, - namespace_kwargs: dict[str, Any] | None, fragment_data: list[dict[str, list[int]]], fragment_ids_to_use: list[int], num_partitions: int | None, @@ -351,7 +354,7 @@ def _create_segmented_index( max_concurrency=max_concurrency, ) handler = handler_cls( - lance_ds=lance_ds, + open_context=open_context, column=column, index_type=index_type, name=name, @@ -392,11 +395,7 @@ def _create_segmented_index( # Reload dataset to pick up the latest version (segment files were written # by workers against the version that was current at their invocation time). - lance_ds = lance.dataset( - None if namespace_kwargs else uri, - storage_options=storage_options, - **(namespace_kwargs or {}), - ) + lance_ds = open_context.open_latest() index_metas = _prepare_index_segments_for_commit(lance_ds, index_type, index_metas) logger.info( @@ -423,16 +422,12 @@ def _prepare_index_segments_for_commit( def _create_partitioned_index( - lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, *, column: str, index_type: str, name: str, replace: bool, - storage_options: dict[str, Any] | None, - namespace_kwargs: dict[str, Any] | None, - namespace_commit_kwargs: dict[str, Any] | None, fragment_data: list[dict[str, list[int]]], fragment_ids_to_use: list[int], num_partitions: int | None, @@ -453,7 +448,7 @@ def _create_partitioned_index( max_concurrency=max_concurrency, ) handler = handler_cls( - lance_ds=lance_ds, + open_context=open_context, column=column, index_type=index_type, name=name, @@ -472,11 +467,7 @@ def _create_partitioned_index( df.collect() logger.info("Starting index metadata merging by reloading dataset to get latest state") - lance_ds = lance.dataset( - None if namespace_kwargs else uri, - storage_options=storage_options, - **(namespace_kwargs or {}), - ) + lance_ds = open_context.open_latest() lance_ds.merge_index_metadata(index_id, index_type) logger.info("Starting atomic index creation and commit") @@ -518,11 +509,11 @@ def _create_partitioned_index( # Commit the index operation atomically lance.LanceDataset.commit( - uri, + open_context.uri, create_index_op, read_version=lance_ds.version, - storage_options=storage_options, - **(namespace_commit_kwargs or namespace_kwargs or {}), + storage_options=open_context.storage_options, + **open_context.commit_kwargs, ) logger.info("Index %s created successfully with ID %s", name, index_id) diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index 82c2739..4945b64 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -14,12 +14,21 @@ the atomic ``declare_table`` operation, while overwrite describes first and declares only when the namespace raises its typed ``TableNotFoundError``. Other namespace failures propagate unchanged. + +3. **Distributed workers reopen, they never receive an open dataset.** + ``LanceDataset.__reduce__`` only carries (uri, storage_options, version, + manifest, ...); ``_namespace_client``, ``_table_id`` and + ``_namespace_client_managed_versioning`` are assigned after construction and + are therefore dropped by pickle. A worker that received a pickled dataset + would silently lose its namespace identity and its managed-versioning commit + handler. :class:`DatasetOpenContext` carries the serializable facts instead + and reopens on the worker. """ from __future__ import annotations import os -from dataclasses import dataclass +from dataclasses import dataclass, field from functools import lru_cache from typing import Any from urllib.parse import unquote, urlparse @@ -221,6 +230,13 @@ def merge_storage_options(*layers: dict[str, Any] | None) -> dict[str, Any] | No Callers order layers from lowest to highest priority, conventionally: io_config-derived < user-provided ``storage_options`` < namespace-vended. + + Returning ``None`` rather than ``{}`` for an empty result is load-bearing: + pylance only installs the namespace credential-refresh provider when the + initial storage options are ``Some`` (``python/src/dataset.rs``). An empty + dict would install a provider whose initial options carry no + ``expires_at_millis``, i.e. one that is treated as never expiring and so + never refreshes. """ merged: dict[str, Any] = {} for layer in layers: @@ -238,6 +254,121 @@ def pop_namespace_params(open_kwargs: dict[str, Any]) -> tuple[str | None, dict[ ) +@dataclass(frozen=True) +class DatasetOpenContext: + """Everything a distributed worker needs to reopen a table, and nothing else. + + Maintenance operations (compaction, scalar index, merge columns) plan on the + driver and execute on workers. The driver's live ``LanceDataset`` cannot be + shipped: pickling it drops the namespace client, the table id and the + managed-versioning flag (see this module's docstring), so a worker would + commit as if the table were an unmanaged uri table. + + This context is the serializable substitute. It deliberately excludes the + live dataset, the namespace client, ``serialized_manifest``, ``session``, + ``commit_lock`` and ``asof``: + + * the client is rebuilt per process from the triple via the ``lru_cache``; + * the manifest is omitted so the payload stays independent of fragment + count -- workers pay one pinned-manifest read instead; + * ``asof``/tags are already resolved by the driver into :attr:`version`, so + workers cannot drift onto a different snapshot. + """ + + uri: str + version: int + storage_options: dict[str, Any] | None = field(default=None, repr=False) + + namespace_impl: str | None = None + namespace_properties: dict[str, str] | None = None + table_id: list[str] | None = None + managed_versioning: bool = False + + block_size: int | None = None + index_cache_size: int | None = None + metadata_cache_size_bytes: int | None = None + default_scan_options: dict[str, Any] | None = field(default=None, repr=False) + + @classmethod + def from_dataset( + cls, + dataset: lance.LanceDataset, + uri: str, + *, + storage_options: dict[str, Any] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, + table_id: list[str] | None = None, + managed_versioning: bool = False, + block_size: int | None = None, + index_cache_size: int | None = None, + metadata_cache_size_bytes: int | None = None, + default_scan_options: dict[str, Any] | None = None, + ) -> DatasetOpenContext: + """Derive a context from the snapshot the driver actually opened.""" + return cls( + uri=uri, + version=dataset.version, + storage_options=storage_options, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, + managed_versioning=managed_versioning, + block_size=block_size, + index_cache_size=index_cache_size, + metadata_cache_size_bytes=metadata_cache_size_bytes, + default_scan_options=default_scan_options, + ) + + @property + def namespace_kwargs(self) -> dict[str, Any]: + return get_namespace_kwargs(self.namespace_impl, self.namespace_properties, self.table_id) + + @property + def commit_kwargs(self) -> dict[str, Any]: + return get_namespace_commit_kwargs( + self.namespace_impl, + self.namespace_properties, + self.table_id, + self.managed_versioning, + ) + + def _open(self, version: int | None) -> lance.LanceDataset: + """Open through the low-level constructor, passing the physical uri. + + ``lance.dataset(None, namespace_client=..., table_id=...)`` resolves the + location with a Python-side ``describe_table`` on *every* call, which + would put one namespace round-trip on every task. The driver already + resolved the location, so workers pass the uri directly and the + constructor still installs the namespace wiring (credential provider and, + when the catalog manages versioning, the commit handler). + """ + return lance.LanceDataset( + self.uri, + version=version, + storage_options=self.storage_options, + block_size=self.block_size, + index_cache_size=self.index_cache_size, + metadata_cache_size_bytes=self.metadata_cache_size_bytes, + default_scan_options=self.default_scan_options, + namespace_client=get_or_create_namespace(self.namespace_impl, self.namespace_properties), + table_id=self.table_id, + namespace_client_managed_versioning=self.managed_versioning, + ) + + def open_pinned(self) -> lance.LanceDataset: + """Open the exact snapshot the driver planned against. + + Every worker must see the same fragment ids and schema the coordinator + used to build its plan, so this is the default for worker code. + """ + return self._open(self.version) + + def open_latest(self) -> lance.LanceDataset: + """Open the newest version; only for coordinator steps that must observe worker output.""" + return self._open(None) + + def open_dataset_from_open_kwargs(ds_uri: str | None, open_kwargs: dict[str, Any] | None) -> lance.LanceDataset: """Re-open a dataset on a worker from serialized open arguments. diff --git a/daft_lance/utils.py b/daft_lance/utils.py index f23025b..e02c58a 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -10,6 +10,7 @@ from daft.io.object_store_options import io_config_to_storage_options from daft.logical.schema import Schema as DaftSchema from daft_lance.namespace import ( + DatasetOpenContext, get_namespace_commit_kwargs, get_namespace_kwargs, has_namespace_params, @@ -63,6 +64,29 @@ def commit_kwargs(self) -> dict[str, Any]: self.managed_versioning, ) + def worker_open_context(self) -> DatasetOpenContext: + """Derive the serializable context distributed maintenance workers reopen from. + + Note this pins :attr:`dataset`'s *resolved* version rather than the + caller's ``version``/``asof`` request, so every worker lands on the exact + snapshot the driver planned against. + """ + return DatasetOpenContext.from_dataset( + self.dataset, + self.uri, + storage_options=self.storage_options, + namespace_impl=self.open_kwargs.get("namespace_impl"), + namespace_properties=self.open_kwargs.get("namespace_properties"), + table_id=self.open_kwargs.get("table_id"), + managed_versioning=self.managed_versioning, + block_size=self.open_kwargs.get("block_size"), + index_cache_size=self.open_kwargs.get("index_cache_size"), + metadata_cache_size_bytes=self.open_kwargs.get("metadata_cache_size_bytes"), + # The stripped variant that was actually handed to pylance; the + # unstripped copy on this handle is for Daft planning only. + default_scan_options=self.open_kwargs.get("default_scan_options"), + ) + def distribute_fragments_balanced(fragments: list[Any], fragment_group_size: int) -> list[dict[str, list[int]]]: """Distribute fragments across workers using a balanced algorithm considering fragment sizes.""" diff --git a/tests/io/lancedb/test_fast_path_merge.py b/tests/io/lancedb/test_fast_path_merge.py index 2947ad5..147eef2 100644 --- a/tests/io/lancedb/test_fast_path_merge.py +++ b/tests/io/lancedb/test_fast_path_merge.py @@ -17,6 +17,7 @@ import daft from daft_lance.lance_merge_column import _can_use_fast_path, merge_columns_from_df +from daft_lance.namespace import DatasetOpenContext # --------------------------------------------------------------------------- # Fixtures and helpers @@ -28,6 +29,11 @@ def ds_path(tmp_path_factory): yield str(tmp_path_factory.mktemp("fast_path")) +def open_ctx(ds: lance.LanceDataset, path: str) -> DatasetOpenContext: + """Uri-only worker context pinned to the snapshot the caller opened.""" + return DatasetOpenContext.from_dataset(ds, path) + + def create_dataset(path: str, fragments: list[dict]) -> lance.LanceDataset: for i, data in enumerate(fragments): table = pa.table(data) @@ -61,7 +67,7 @@ def test_fast_path_single_column_int(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2, 3], "val": [10, 20, 30]}]) df = read_with_metadata(ds_path) df = df.with_column("doubled", daft.col("val").cast(daft.DataType.int64()) * 2) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().to_pydict() assert result["doubled"] == [2 * v for v in result["val"]] @@ -69,7 +75,7 @@ def test_fast_path_single_column_float(self, ds_path): ds = create_dataset(ds_path, [{"x": [1.0, 2.0, 3.0]}]) df = read_with_metadata(ds_path) df = df.with_column("half", daft.col("x").cast(daft.DataType.float64()) / 2.0) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().to_pydict() for x, h in zip(result["x"], result["half"]): assert pytest.approx(x / 2.0, rel=1e-6) == h @@ -78,7 +84,7 @@ def test_fast_path_single_column_string(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2], "name": ["alice", "bob"]}]) df = read_with_metadata(ds_path) df = df.with_column("greeting", daft.lit("hello_") + daft.col("name")) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().to_pydict() assert result["greeting"] == ["hello_alice", "hello_bob"] @@ -88,7 +94,7 @@ def test_fast_path_multiple_new_columns(self, ds_path): df = df.with_column("a", daft.col("id").cast(daft.DataType.int64()) * 10) df = df.with_column("b", daft.col("id").cast(daft.DataType.float64()) + 0.5) df = df.with_column("c", daft.lit("row_") + daft.col("id").cast(daft.DataType.string())) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().to_pydict() assert result["a"] == [10, 20, 30] for i, b in zip(result["id"], result["b"]): @@ -106,7 +112,7 @@ def test_fast_path_multi_fragment(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("score", daft.col("val").cast(daft.DataType.float64()) * 1.5) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() for v, s in zip(result["val"], result["score"]): assert pytest.approx(v * 1.5, rel=1e-6) == s @@ -122,7 +128,7 @@ def test_fast_path_computed_column(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("z", daft.col("x").cast(daft.DataType.int64()) + daft.col("y").cast(daft.DataType.int64())) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("x").to_pydict() for x, y, z in zip(result["x"], result["y"], result["z"]): assert x + y == z @@ -141,7 +147,7 @@ def test_existing_columns_unchanged(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("new_col", daft.lit(999)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) after = ds2.to_table().sort_by("id").to_pydict() assert after["id"] == before["id"] @@ -158,7 +164,7 @@ def test_fragment_file_count(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("b", daft.lit(0)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) for frag in ds2.get_fragments(): assert len(list(frag.data_files())) == 2 @@ -172,7 +178,7 @@ def test_original_files_not_rewritten(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("b", daft.lit(42)) - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) for fname, old_hash in original_files.items(): fpath = os.path.join(data_dir, fname) @@ -192,7 +198,7 @@ def test_row_count_preserved(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("w", daft.lit(0)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.count_rows() == original_count for frag, expected in zip(ds2.get_fragments(), per_frag_counts): @@ -202,7 +208,7 @@ def test_schema_evolution_correct(self, ds_path): ds = create_dataset(ds_path, [{"id": [1], "name": ["x"]}]) df = read_with_metadata(ds_path) df = df.with_column("score", daft.lit(3.14)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) schema_names = set(ds2.schema.names) assert "id" in schema_names assert "name" in schema_names @@ -221,7 +227,7 @@ def test_rowaddr_sorting_restores_order(self, ds_path): df = read_with_metadata(ds_path) # Daft doesn't guarantee order, but let's force a known computation df = df.with_column("doubled", daft.col("val").cast(daft.DataType.int64()) * 2) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert result["doubled"] == [20, 40, 60, 80, 100] @@ -236,7 +242,7 @@ def test_multi_fragment_ordering(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("neg", daft.col("id").cast(daft.DataType.int64()) * -1) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() for i, n in zip(result["id"], result["neg"]): assert n == -i @@ -310,12 +316,12 @@ def test_two_sequential_merges(self, ds_path): # First merge: add column A df = read_with_metadata(ds_path) df = df.with_column("a", daft.col("id").cast(daft.DataType.int64()) * 10) - ds = merge_columns_from_df(df, ds, ds_path) + ds = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) # Second merge: add column B df2 = read_with_metadata(ds_path) df2 = df2.with_column("b", daft.col("id").cast(daft.DataType.int64()) * 100) - ds2 = merge_columns_from_df(df2, ds, ds_path) + ds2 = merge_columns_from_df(df2, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert result["a"] == [10, 20, 30] @@ -331,7 +337,7 @@ def test_merge_after_append(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("flag", daft.lit(True)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert result["flag"] == [True, True, True, True] assert len(result["id"]) == 4 @@ -342,14 +348,14 @@ def test_merge_preserves_previous_merge(self, ds_path): # Merge A df = read_with_metadata(ds_path) df = df.with_column("a", daft.col("val").cast(daft.DataType.int64()) + 1) - ds = merge_columns_from_df(df, ds, ds_path) + ds = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) check_a = ds.to_table().sort_by("id").to_pydict() assert check_a["a"] == [11, 21] # Merge B df2 = read_with_metadata(ds_path) df2 = df2.with_column("b", daft.col("val").cast(daft.DataType.int64()) + 2) - ds2 = merge_columns_from_df(df2, ds, ds_path) + ds2 = merge_columns_from_df(df2, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert result["a"] == [11, 21], "Column A corrupted by second merge" assert result["b"] == [12, 22] @@ -366,7 +372,7 @@ def test_type_int64(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2]}]) df = read_with_metadata(ds_path) df = df.with_column("x", daft.lit(42).cast(daft.DataType.int64())) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.schema.field("x").type == pa.int64() assert ds2.to_table().column("x").to_pylist() == [42, 42] @@ -376,7 +382,7 @@ def test_type_float32(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2]}]) df = read_with_metadata(ds_path) df = df.with_column("x", daft.lit(3.14).cast(daft.DataType.float32())) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) # Value is preserved even if type widens vals = ds2.to_table().column("x").to_pylist() assert all(pytest.approx(v, rel=1e-5) == 3.14 for v in vals) @@ -385,21 +391,21 @@ def test_type_float64(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2]}]) df = read_with_metadata(ds_path) df = df.with_column("x", daft.lit(2.718)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.schema.field("x").type == pa.float64() def test_type_string(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2]}]) df = read_with_metadata(ds_path) df = df.with_column("label", daft.lit("hello")) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.to_table().column("label").to_pylist() == ["hello", "hello"] def test_type_bool(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2, 3]}]) df = read_with_metadata(ds_path) df = df.with_column("flag", daft.col("id").cast(daft.DataType.int64()) > 1) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").column("flag").to_pylist() assert result == [False, True, True] @@ -412,7 +418,7 @@ def test_type_nullable(self, ds_path): (daft.col("id").cast(daft.DataType.int64()) % 2 != 0).cast(daft.DataType.int64()) * daft.col("id").cast(daft.DataType.int64()), ) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").column("maybe").to_pylist() # id=1 → odd → 1*1=1, id=2 → even → 0*2=0, id=3 → odd → 1*3=3, id=4 → even → 0*4=0 assert result == [1, 0, 3, 0] @@ -428,7 +434,7 @@ def test_single_row_fragment(self, ds_path): ds = create_dataset(ds_path, [{"id": [1]}]) df = read_with_metadata(ds_path) df = df.with_column("x", daft.lit(99)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.to_table().to_pydict() == {"id": [1], "x": [99]} def test_large_fragment(self, ds_path): @@ -436,7 +442,7 @@ def test_large_fragment(self, ds_path): ds = create_dataset(ds_path, [{"id": list(range(n))}]) df = read_with_metadata(ds_path) df = df.with_column("neg", daft.col("id").cast(daft.DataType.int64()) * -1) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert len(result["id"]) == n for i, neg in zip(result["id"], result["neg"]): @@ -449,7 +455,7 @@ def test_many_fragments(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("doubled", daft.col("id").cast(daft.DataType.int64()) * 2) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert len(result["id"]) == 20 for i, d in zip(result["id"], result["doubled"]): @@ -460,14 +466,14 @@ def test_empty_new_columns_raises(self, ds_path): df = read_with_metadata(ds_path) # No new columns added — only existing + metadata columns with pytest.raises(ValueError, match="No new columns"): - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) def test_dataset_version_incremented(self, ds_path): ds = create_dataset(ds_path, [{"id": [1]}]) v_before = ds.version df = read_with_metadata(ds_path) df = df.with_column("x", daft.lit(1)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.version == v_before + 1 @@ -491,7 +497,7 @@ def test_fast_vs_slow_identical_results(self, tmp_path_factory): # Fast path: read with _rowaddr + fragment_id df_fast = read_with_metadata(fast_path) df_fast = df_fast.with_column("score", daft.col("val").cast(daft.DataType.float64()) * 2.5) - ds_fast = merge_columns_from_df(df_fast, ds_fast, fast_path) + ds_fast = merge_columns_from_df(df_fast, ds_fast, open_ctx(ds_fast, fast_path)) # Slow path: read with fragment_id only, use business key df_slow = daft.read_lance(slow_path, include_fragment_id=True, default_scan_options={"with_row_address": True}) @@ -501,13 +507,12 @@ def test_fast_vs_slow_identical_results(self, tmp_path_factory): ds_slow = _merge_slow_path( df_slow, ds_slow, - slow_path, + open_ctx(ds_slow, slow_path), read_columns=["_rowaddr", "score"], left_on="_rowaddr", right_on="_rowaddr", reader_schema=None, batch_size=None, - storage_options=None, ) fast_result = ds_fast.to_table().sort_by("id").to_pydict() @@ -534,7 +539,7 @@ def test_read_after_merge_with_filter(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("score", daft.col("val").cast(daft.DataType.int64()) * 2) - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) ds2 = lance.dataset(ds_path) filtered = ds2.to_table(filter="score > 40") @@ -546,7 +551,7 @@ def test_read_after_merge_with_projection(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2], "val": [10, 20]}]) df = read_with_metadata(ds_path) df = df.with_column("new_col", daft.lit(42)) - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) ds2 = lance.dataset(ds_path) projected = ds2.to_table(columns=["new_col"]) @@ -558,7 +563,7 @@ def test_read_after_merge_select_original_only(self, ds_path): ds = create_dataset(ds_path, [original]) df = read_with_metadata(ds_path) df = df.with_column("extra", daft.lit("x")) - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) ds2 = lance.dataset(ds_path) result = ds2.to_table(columns=["id", "name"]).sort_by("id").to_pydict() @@ -575,7 +580,7 @@ def test_scan_fragments_individually(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("doubled", daft.col("id").cast(daft.DataType.int64()) * 2) - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) ds2 = lance.dataset(ds_path) all_ids = [] @@ -634,7 +639,7 @@ def _make_vec(ids): return [np.array([float(i)] * N, dtype=np.float32) for i in ids.to_pylist()] df = df.with_column("embedding", _make_vec(daft.col("id"))) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) field = ds2.schema.field("embedding") # Type must be preserved: fixed_size_list[N], NOT list @@ -676,7 +681,7 @@ def test_next_fid_uses_manifest_max_field_id(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("score", daft.col("id").cast(daft.DataType.int64()) * 10) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert result["score"] == [10, 20, 30], f"score is null or wrong (field ID collision): {result['score']}" diff --git a/tests/io/lancedb/test_lancedb_scalar_index.py b/tests/io/lancedb/test_lancedb_scalar_index.py index 1b325e6..eeb4280 100644 --- a/tests/io/lancedb/test_lancedb_scalar_index.py +++ b/tests/io/lancedb/test_lancedb_scalar_index.py @@ -17,6 +17,27 @@ _prepare_index_segments_for_commit, create_scalar_index_internal, ) +from daft_lance.namespace import DatasetOpenContext + + +class FakeOpenContext: + """Stands in for DatasetOpenContext so handlers can be driven against a fake dataset. + + Also counts opens, which is what proves a handler reopens once per instance + rather than once per call. + """ + + def __init__(self, dataset, uri="memory://fake"): + self.dataset = dataset + self.uri = uri + self.opens = 0 + + def open_pinned(self): + self.opens += 1 + return self.dataset + + def open_latest(self): + return self.open_pinned() @pytest.fixture @@ -621,7 +642,7 @@ def create_index_uncommitted(self, **kwargs): fake_ds = FakeLanceDataset() handler = SegmentedFragmentIndexHandler( - lance_ds=fake_ds, + open_context=FakeOpenContext(fake_ds), column="price", index_type="BTREE", name="price_idx", @@ -704,7 +725,7 @@ def create_index_uncommitted(self, **kwargs): fake_ds = FakeLanceDataset() handler = SegmentedFragmentIndexHandler( - lance_ds=fake_ds, + open_context=FakeOpenContext(fake_ds), column="flag", index_type="BITMAP", name="flag_idx", @@ -772,7 +793,7 @@ def fake_create_segmented_index(**kwargs): create_scalar_index_internal( lance_ds=FakeLanceDataset(), - uri="memory://bitmap", + open_context=DatasetOpenContext(uri="memory://bitmap", version=1), column="flag", index_type="BITMAP", name="flag_bitmap_idx", @@ -804,7 +825,7 @@ def create_scalar_index(self, **kwargs): create_scalar_index_internal( lance_ds=fake_ds, - uri="memory://bitmap", + open_context=DatasetOpenContext(uri="memory://bitmap", version=1), column="flag", index_type="BITMAP", name="flag_bitmap_idx", @@ -1067,7 +1088,7 @@ def fake_create_partitioned_index(**kwargs): create_scalar_index_internal( lance_ds=FakeLanceDataset(), - uri="memory://btree", + open_context=DatasetOpenContext(uri="memory://btree", version=1), column="price", index_type="BTREE", name="price_btree_idx", diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index c149d9f..0056b68 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -10,6 +10,7 @@ import daft_lance import daft_lance.namespace as namespace_mod from daft_lance.lance_data_sink import LanceDataSink +from daft_lance.utils import construct_lance_dataset_handle def _dir_ns(tmp_path: Path) -> dict[str, Any]: @@ -670,3 +671,178 @@ def test_namespace_requires_impl() -> None: def test_namespace_requires_uri_or_namespace() -> None: with pytest.raises(ValueError, match="Must provide either 'uri' OR"): daft_lance.read_lance() + + +# --------------------------------------------------------------------------- +# Distributed worker reopen (DatasetOpenContext) +# +# LanceDataset.__reduce__ drops _namespace_client / _table_id / +# _namespace_client_managed_versioning, so a pickled dataset reaches a worker +# stripped of its namespace identity. Maintenance paths therefore ship a +# DatasetOpenContext and reopen. These tests lock that contract down. +# --------------------------------------------------------------------------- + + +def _ns_handle(tmp_path: Path, table: str = "ctx_tbl"): + daft_lance.write_lance( + daft.from_pydict({"score": [1, 2, 3, 4]}), + table_id=[table], + mode="create", + **_dir_ns(tmp_path), + ).collect() + return construct_lance_dataset_handle(None, table_id=[table], **_dir_ns(tmp_path)) + + +def test_worker_open_context_is_free_of_live_objects(tmp_path: Path) -> None: + """The context must survive pickling without dragging unpicklable state along.""" + import pickle + + import lance + from lance_namespace import LanceNamespace + + context = _ns_handle(tmp_path).worker_open_context() + restored = pickle.loads(pickle.dumps(context)) + + for value in vars(restored).values(): + assert not isinstance(value, lance.LanceDataset) + assert not isinstance(value, LanceNamespace) + assert not hasattr(restored, "serialized_manifest") + + assert restored.table_id == ["ctx_tbl"] + assert restored.namespace_impl == "dir" + assert restored.version == context.version + + +def test_worker_open_restores_namespace_identity(tmp_path: Path) -> None: + """Reopening on a worker must rebuild the wiring pickle would have dropped.""" + context = _ns_handle(tmp_path, "identity_tbl").worker_open_context() + + worker_ds = context.open_pinned() + + assert worker_ds._namespace_client is not None + assert worker_ds._table_id == ["identity_tbl"] + assert worker_ds.version == context.version + assert worker_ds.count_rows() == 4 + + +def test_worker_open_propagates_managed_versioning(tmp_path: Path) -> None: + """A catalog-managed table must not be reopened as an unmanaged uri table.""" + import dataclasses + + context = dataclasses.replace( + _ns_handle(tmp_path, "managed_tbl").worker_open_context(), + managed_versioning=True, + ) + + assert context.commit_kwargs["namespace_client_managed_versioning"] is True + assert context.open_pinned()._namespace_client_managed_versioning is True + + +def test_worker_open_does_not_describe_the_table_location(tmp_path: Path) -> None: + """Workers must not put a namespace round-trip on every task. + + ``lance.dataset(None, namespace_client=..., table_id=...)`` resolves the + location with a describe_table on every call; the context passes the uri the + driver already resolved, so the worker open costs zero namespace calls. + """ + import lance + + metered = { + "namespace_impl": "dir", + "namespace_properties": {"root": str(tmp_path), "ops_metrics_enabled": "true"}, + } + daft_lance.write_lance( + daft.from_pydict({"score": [1, 2, 3, 4]}), table_id=["metered_tbl"], mode="create", **metered + ).collect() + handle = construct_lance_dataset_handle(None, table_id=["metered_tbl"], **metered) + context = handle.worker_open_context() + client = namespace_mod.get_or_create_namespace(metered["namespace_impl"], metered["namespace_properties"]) + + client.reset_ops_metrics() + worker_ds = context.open_pinned() + worker_ds.count_rows() + assert client.retrieve_ops_metrics().get("describe_table", 0) == 0 + + # The high-level entry point is what the low-level open exists to avoid. + client.reset_ops_metrics() + lance.dataset(None, namespace_client=client, table_id=["metered_tbl"]) + assert client.retrieve_ops_metrics().get("describe_table", 0) == 1 + + +def test_open_latest_sees_versions_written_after_pinning(tmp_path: Path) -> None: + """Workers stay on the planned snapshot; only coordinator steps move forward.""" + handle = _ns_handle(tmp_path, "versioned_tbl") + context = handle.worker_open_context() + + daft_lance.write_lance( + daft.from_pydict({"score": [5, 6]}), table_id=["versioned_tbl"], mode="append", **_dir_ns(tmp_path) + ).collect() + + assert context.open_pinned().version == context.version + assert context.open_pinned().count_rows() == 4 + assert context.open_latest().version > context.version + assert context.open_latest().count_rows() == 6 + + +def test_maintenance_udfs_hold_a_context_not_a_dataset(tmp_path: Path) -> None: + """No maintenance UDF may capture the driver's live dataset.""" + import pickle + + import lance + + from daft_lance.lance_compaction import CompactionTaskUDF + from daft_lance.lance_merge_column import ( + FastPathFragmentWriter, + FragmentHandler, + GroupFragmentMergeUDF, + ) + from daft_lance.lance_scalar_index import ( + FragmentIndexHandler, + SegmentedFragmentIndexHandler, + ) + + context = _ns_handle(tmp_path, "udf_tbl").worker_open_context() + + plain = [ + CompactionTaskUDF(context), + FragmentIndexHandler(context, "score", "BTREE", "idx", "uuid", False), + SegmentedFragmentIndexHandler(context, "score", "BTREE", "idx"), + ] + # daft.cls wraps these, so reach through to the instance it actually holds. + wrapped = [ + FragmentHandler(context, {"doubled": "score * 2"}, ["score"]), + GroupFragmentMergeUDF(context), + FastPathFragmentWriter(context, ["doubled"]), + ] + + instances = plain + [udf._daft_get_instance() for udf in wrapped] + for udf in instances: + state = vars(udf) + assert not any(isinstance(value, lance.LanceDataset) for value in state.values()), ( + f"{type(udf).__name__} captured a live LanceDataset" + ) + assert state["open_context"] is context + + for udf in plain: + assert isinstance(pickle.loads(pickle.dumps(udf)), type(udf)) + + +def test_worker_udf_opens_the_dataset_once_per_instance(tmp_path: Path) -> None: + """The pinned-manifest read is per UDF instance, never per call.""" + from daft_lance.lance_scalar_index import SegmentedFragmentIndexHandler + + context = _ns_handle(tmp_path, "reopen_tbl").worker_open_context() + opens = [] + + class CountingContext: + uri = context.uri + + def open_pinned(self): + opens.append(1) + return context.open_pinned() + + handler = SegmentedFragmentIndexHandler(CountingContext(), "score", "BTREE", "idx") + handler._dataset() + handler._dataset() + + assert len(opens) == 1 From 918a43a2af33f743f47ca9784bc47677cefb7a1f Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Wed, 22 Jul 2026 16:08:25 +0900 Subject: [PATCH 25/25] fix(namespace): pin scan workers to the planned version, recover lost declare race Two review findings. Scan workers were not pinned to the version the driver planned against. construct_lance_dataset_handle stored the caller's `version` argument in open_kwargs, which defaults to None, and open_kwargs is exactly what crosses to scan workers -- so `version=None` there means every task independently opens latest. Reproduced both failure modes: a compaction landing between planning and execution leaves workers looking for fragment ids that no longer exist (planned [0,1,2], worker sees [3], get_fragment(0) returns None), and an overwrite silently substitutes different data with no error at all (driver planned 3 rows, worker read 1 different row). open_kwargs now carries the resolved numeric `dataset.version`, and `asof` is dropped since it is an input to that resolution rather than something workers should re-evaluate. pylance lets `version` win when both are passed, so dropping it is belt-and-braces rather than load bearing. This is the same invariant DatasetOpenContext already enforces for the maintenance paths; scan was left out when that change deliberately avoided touching the scan state model. Overwrite could also lose a declare race: when the initial describe raises TableNotFoundError but a rival writer declares before our own declare_table lands, TableAlreadyExistsError propagated and the overwrite failed. Overwrite targets whatever exists now, so it re-describes on that conflict. `create` is unchanged -- for it the conflict is the correct answer, and a test pins that distinction. Claude-Session: https://claude.ai/code/session_01Ss3gSAbK3FwfPmWqk1CSbJ --- daft_lance/namespace.py | 18 +++-- daft_lance/utils.py | 11 ++- tests/io/lancedb/test_namespace.py | 104 ++++++++++++++++++++++++++++- 3 files changed, 124 insertions(+), 9 deletions(-) diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index 4945b64..ea7c8a4 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -205,8 +205,9 @@ def resolve_namespace_table( """Resolve a namespace table to a :class:`ResolvedNamespaceTable`. ``mode="create"`` atomically declares a new table. ``mode="overwrite"`` - declares the table only when a typed ``TableNotFoundError`` is raised; - other modes require the table to exist. + declares the table only when a typed ``TableNotFoundError`` is raised, and + falls back to a second describe if it loses the declare race; other modes + require the table to exist. """ namespace = get_or_create_namespace(namespace_impl, namespace_properties) if namespace is None or table_id is None: @@ -215,14 +216,21 @@ def resolve_namespace_table( if mode == "create": return _declare_table(namespace, table_id) - from lance_namespace.errors import TableNotFoundError + from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError try: return _resolved_from_response(_describe_table(namespace, table_id)) except TableNotFoundError: - if mode == "overwrite": + if mode != "overwrite": + raise + try: return _declare_table(namespace, table_id) - raise + except TableAlreadyExistsError: + # Lost a declare race: another writer declared the table between our + # describe and our declare. Overwrite targets whatever exists now, so + # re-describing is the correct recovery. ``create`` deliberately does + # not reach here -- for it the conflict is the answer. + return _resolved_from_response(_describe_table(namespace, table_id)) def merge_storage_options(*layers: dict[str, Any] | None) -> dict[str, Any] | None: diff --git a/daft_lance/utils.py b/daft_lance/utils.py index e02c58a..c3cf79d 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -220,12 +220,21 @@ def construct_lance_dataset_handle( effective_kwargs = { "storage_options": merged_storage_options, - "version": version, + # Pin the snapshot the driver resolved, not the caller's request. These + # kwargs cross to scan workers, and ``version=None`` there means "open + # latest": a compaction landing between planning and execution leaves + # workers looking for fragment ids that no longer exist, and an + # overwrite silently feeds them different data entirely. + "version": dataset.version, "namespace_impl": namespace_impl, "namespace_properties": namespace_properties, "table_id": table_id, } effective_kwargs.update(kwargs or {}) + # ``asof``/tags are inputs to the version resolution above; carrying them + # further is redundant at best. (pylance lets ``version`` win when both are + # passed, so this is belt-and-braces.) + effective_kwargs.pop("asof", None) return LanceDatasetHandle( dataset=dataset, uri=resolved_uri, diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py index 0056b68..991fce0 100644 --- a/tests/io/lancedb/test_namespace.py +++ b/tests/io/lancedb/test_namespace.py @@ -446,7 +446,8 @@ def test_construct_lance_dataset_empty_storage_options_falls_back_to_io_config( captured = {} class FakeDataset: - pass + # construct_lance_dataset_handle pins this into open_kwargs["version"]. + version = 1 def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **kwargs: Any) -> Any: captured["storage_options"] = storage_options @@ -535,7 +536,8 @@ def test_construct_lance_dataset_storage_options_priority(monkeypatch: pytest.Mo captured = {} class FakeDataset: - pass + # construct_lance_dataset_handle pins this into open_kwargs["version"]. + version = 1 def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **kwargs: Any) -> Any: captured["uri"] = uri @@ -584,7 +586,8 @@ def test_construct_lance_dataset_io_config_reaches_namespace_location(monkeypatc captured = {} class FakeDataset: - pass + # construct_lance_dataset_handle pins this into open_kwargs["version"]. + version = 1 def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **kwargs: Any) -> Any: captured["storage_options"] = storage_options @@ -846,3 +849,98 @@ def open_pinned(self): handler._dataset() assert len(opens) == 1 + + +def test_scan_open_kwargs_pin_the_planned_version(tmp_path: Path) -> None: + """Workers must read the snapshot the driver planned against, not latest. + + ``open_kwargs`` crosses to scan workers, so a ``version=None`` there means + each task independently opens latest. A compaction landing in between makes + planned fragment ids vanish; an overwrite silently substitutes the data. + """ + import lance + import pyarrow as pa + + ds_path = str(tmp_path / "pinned.lance") + lance.write_dataset(pa.table({"a": [1, 2, 3]}), ds_path) + handle = construct_lance_dataset_handle(ds_path) + planned_version = handle.dataset.version + + assert handle.open_kwargs["version"] == planned_version + + lance.write_dataset(pa.table({"a": [99]}), ds_path, mode="overwrite") + + worker_ds = namespace_mod.open_dataset_from_open_kwargs(handle.uri, handle.open_kwargs) + assert worker_ds.version == planned_version + assert worker_ds.to_table().to_pydict() == {"a": [1, 2, 3]} + + +def test_scan_open_kwargs_resolve_asof_into_a_numeric_version(tmp_path: Path) -> None: + """A tag/asof request is resolved on the driver; workers get the number.""" + import lance + import pyarrow as pa + + ds_path = str(tmp_path / "asof.lance") + lance.write_dataset(pa.table({"a": [1]}), ds_path) + lance.write_dataset(pa.table({"a": [2]}), ds_path, mode="append") + + handle = construct_lance_dataset_handle(ds_path, asof="2030-01-01 00:00:00") + + assert "asof" not in handle.open_kwargs + assert handle.open_kwargs["version"] == handle.dataset.version + + +def test_namespace_overwrite_recovers_from_a_lost_declare_race(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Overwrite must survive another writer declaring the table first.""" + from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError + + class RacingNamespace: + def __init__(self) -> None: + self.describes = 0 + self.declares = 0 + + def describe_table(self, request: Any) -> Any: + self.describes += 1 + # First describe loses the race; by the second, the rival's table exists. + if self.describes == 1: + raise TableNotFoundError("not found") + return SimpleNamespace(location=str(tmp_path / "rival.lance"), storage_options=None) + + def declare_table(self, request: Any) -> Any: + self.declares += 1 + raise TableAlreadyExistsError("declared by another writer") + + namespace = RacingNamespace() + monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: namespace) + + resolved = namespace_mod.resolve_namespace_table( + namespace_impl="rest", + namespace_properties={"uri": "http://namespace.example"}, + table_id=["raced"], + mode="overwrite", + ) + + assert resolved.uri == str(tmp_path / "rival.lance") + assert (namespace.describes, namespace.declares) == (2, 1) + + +def test_namespace_create_still_fails_on_a_lost_declare_race(monkeypatch: pytest.MonkeyPatch) -> None: + """For create, losing the race is the answer -- it must not be swallowed.""" + from lance_namespace.errors import TableAlreadyExistsError + + class RacingNamespace: + def describe_table(self, request: Any) -> Any: + raise AssertionError("create must not describe") + + def declare_table(self, request: Any) -> Any: + raise TableAlreadyExistsError("declared by another writer") + + monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: RacingNamespace()) + + with pytest.raises(TableAlreadyExistsError): + namespace_mod.resolve_namespace_table( + namespace_impl="rest", + namespace_properties={"uri": "http://namespace.example"}, + table_id=["raced"], + mode="create", + )