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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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], + }