diff --git a/README.md b/README.md index 11e7c3f..106d6bd 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,63 @@ from daft_lance import merge_columns_df 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 + +table_id = ["my_table"] +namespace = {"namespace_impl": "dir", "namespace_properties": {"root": "/tmp/lance_tables"}} + +daft_lance.write_lance( + daft.from_pydict({"id": [1, 2, 3]}), + table_id=table_id, + mode="create", + **namespace, +).collect() + +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. 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). + +#### Daft's own entry points + +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 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..1a0323d 100644 --- a/daft_lance/__init__.py +++ b/daft_lance/__init__.py @@ -10,6 +10,7 @@ merge_columns, merge_columns_df, read_lance, + write_lance, ) __all__ = [ @@ -19,4 +20,5 @@ "merge_columns_df", "read_lance", "take_blobs", + "write_lance", ] diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 1bb7430..0ad2231 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -2,46 +2,50 @@ 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 from daft.daft import IOConfig, ScanOperatorHandle from daft.dataframe import DataFrame +from daft.dependencies import pa from daft.io._checkpoint import attach_checkpoint -from daft.io.object_store_options import io_config_to_storage_options from daft.logical.builder import LogicalPlanBuilder +from daft.schema import Schema 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 .utils import construct_lance_dataset +from .namespace import validate_uri_or_namespace +from .utils import construct_lance_dataset_handle if TYPE_CHECKING: from lance.dataset import LanceDataset from lance.udf import BatchUDF from daft.checkpoint import CheckpointConfig - from daft.dependencies import pa @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, 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, 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. @@ -49,6 +53,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. @@ -93,7 +102,7 @@ def read_lance( each fragment will be processed individually (default behavior). include_fragment_id : Optional, bool Whether to display fragment_id. - if you have the behavior of 'merge_columns_df' or 'write_lance(mode = 'merge')', the `include_fragment_id` must be set to True + Set this to True when preparing input for ``merge_columns_df``. checkpoint: Optional :class:`daft.CheckpointConfig` for progress tracking across runs. Bundles the checkpoint store, the source key column (``on=``), and optional anti-join tuning. Rows whose key already exists in the store are skipped on re-run. Requires the Ray runner. @@ -123,21 +132,25 @@ 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 " - "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 - storage_options = io_config_to_storage_options(io_config, uri_str) - ds = construct_lance_dataset( + dataset_handle = construct_lance_dataset_handle( uri_str, - storage_options=storage_options, + io_config=io_config, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, version=version, asof=asof, block_size=block_size, @@ -148,9 +161,11 @@ def read_lance( ) lance_operator = LanceDBScanOperator( - ds, + dataset_handle.dataset, fragment_group_size=fragment_group_size, include_fragment_id=include_fragment_id, + open_kwargs=dataset_handle.open_kwargs, + default_scan_options=dataset_handle.default_scan_options, ) handle = ScanOperatorHandle.from_python_scan_operator(lance_operator) @@ -161,9 +176,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, @@ -186,6 +204,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. @@ -222,12 +245,15 @@ 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) # Build Lance dataset handle for committing - lance_ds = construct_lance_dataset( + dataset_handle = construct_lance_dataset_handle( uri, storage_options=storage_options, + io_config=io_config, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, version=version, asof=asof, block_size=block_size, @@ -238,12 +264,11 @@ def merge_columns( ) return merge_columns_internal( - lance_ds, - uri, + dataset_handle.dataset, + dataset_handle.worker_open_context(), transform=transform, read_columns=read_columns, reader_schema=reader_schema, - storage_options=storage_options, daft_remote_args=daft_remote_args, concurrency=concurrency, ) @@ -252,9 +277,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, @@ -280,6 +308,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. @@ -318,12 +351,15 @@ 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) # Build Lance dataset handle for committing - lance_ds = construct_lance_dataset( + dataset_handle = construct_lance_dataset_handle( uri, storage_options=storage_options, + io_config=io_config, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, version=version, asof=asof, block_size=block_size, @@ -341,11 +377,10 @@ def merge_columns_df( return merge_columns_from_df( df, - lance_ds=lance_ds, - uri=uri, + lance_ds=dataset_handle.dataset, + open_context=dataset_handle.worker_open_context(), read_columns=read_columns, reader_schema=reader_schema, - storage_options=storage_options, daft_remote_args=daft_remote_args, concurrency=concurrency, left_on=left_on, @@ -356,9 +391,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, @@ -386,6 +424,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". @@ -464,11 +507,14 @@ 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)) - lance_ds = construct_lance_dataset( + dataset_handle = construct_lance_dataset_handle( uri, storage_options=storage_options, + io_config=io_config, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, version=version, asof=asof, block_size=block_size, @@ -479,13 +525,12 @@ def create_scalar_index( ) create_scalar_index_internal( - lance_ds=lance_ds, - uri=uri, + lance_ds=dataset_handle.dataset, + open_context=dataset_handle.worker_open_context(), column=column, index_type=index_type, name=name, replace=replace, - storage_options=storage_options, fragment_group_size=fragment_group_size, num_partitions=num_partitions, max_concurrency=max_concurrency, @@ -496,9 +541,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, @@ -519,6 +567,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. @@ -549,13 +602,14 @@ 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 - ) - lance_ds = lance.dataset( + dataset_handle = construct_lance_dataset_handle( uri, storage_options=storage_options, + io_config=io_config, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, version=version, asof=asof, block_size=block_size, @@ -566,8 +620,70 @@ def compact_files( ) return compact_files_internal( - lance_ds=lance_ds, + lance_ds=dataset_handle.dataset, + open_context=dataset_handle.worker_open_context(), compaction_options=compaction_options, partition_num=partition_num, concurrency=concurrency, ) + + +@PublicAPI +def write_lance( + df: DataFrame, + uri: str | pathlib.Path | None = None, + mode: Literal["create", "append", "overwrite"] = "create", + io_config: IOConfig | None = None, + schema: Schema | pa.Schema | 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", 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. + 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). + + Note: + Target-dependent validation is performed when the returned DataFrame is + executed (for example, by ``collect()``), not while the logical write + plan is constructed. This includes missing/duplicate targets, append + schema compatibility, and storage-version conflicts. + + Examples: + >>> import daft, daft_lance + >>> df = daft.from_pydict({"id": [1, 2]}) + >>> 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, namespace_properties) + + if schema is None: + schema = df.schema() + + sink = LanceDataSink( + uri, + schema, + mode, + io_config, + table_id=table_id, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + **kwargs, + ) + return df.write_sink(sink) diff --git a/daft_lance/lance_compaction.py b/daft_lance/lance_compaction.py index 4575cb7..2ca2b41 100644 --- a/daft_lance/lance_compaction.py +++ b/daft_lance/lance_compaction.py @@ -1,13 +1,16 @@ from __future__ import annotations import logging -from typing import Any +from typing import TYPE_CHECKING, Any from lance import LanceDataset from lance.optimize import Compaction, CompactionMetrics, CompactionOptions, CompactionTask, RewriteResult import daft +if TYPE_CHECKING: + from daft_lance.namespace import DatasetOpenContext + logger = logging.getLogger(__name__) @@ -16,23 +19,36 @@ class CompactionTaskUDF: def __init__( self, - lance_ds: LanceDataset, + open_context: DatasetOpenContext, ) -> None: - self.lance_ds = lance_ds + self.open_context = open_context + self._lance_ds: LanceDataset | None = None + + def _dataset(self) -> LanceDataset: + # Opened once per UDF instance, not per task: the reopen costs a pinned + # manifest read, so it must not sit on the per-row path. + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds def __call__(self, task: CompactionTask) -> RewriteResult: - rewrite = task.execute(self.lance_ds) + rewrite = task.execute(self._dataset()) return rewrite def compact_files_internal( lance_ds: LanceDataset, + open_context: DatasetOpenContext, *, compaction_options: dict[str, Any] | None = None, partition_num: int | None = None, concurrency: int | None = None, ) -> CompactionMetrics | None: - """Execute Lance file compaction in distributed environment using Daft UDF style.""" + """Execute Lance file compaction in distributed environment using Daft UDF style. + + ``lance_ds`` is the driver's live dataset and stays on the driver for + planning and the final commit; ``open_context`` is what workers reopen from. + """ logger.info("Starting UDF-style distributed compaction") plan = Compaction.plan( lance_ds, @@ -59,7 +75,7 @@ def compact_files_internal( CompactionTaskUDF, max_concurrency=concurrency, ) - df = df.select(WrappedRunner(lance_ds)(df["task"]).alias("rewrite")) + df = df.select(WrappedRunner(open_context)(df["task"]).alias("rewrite")) results = df.to_pandas() metrics = Compaction.commit(lance_ds, results["rewrite"].to_list()) diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index 62c36bb..3feef7e 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 @@ -26,6 +26,15 @@ detect_blob_v2_columns, resolve_storage_version, ) +from daft_lance.namespace import ( + ResolvedNamespaceTable, + get_namespace_commit_kwargs, + 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 +49,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 +69,18 @@ def __init__( compact_after_write: bool = True, ) -> None: self._reject_unsupported_modes(mode, use_legacy_format) - if not isinstance(uri, (str, pathlib.Path)): + self._reject_namespace_mem_wal(namespace_impl, table_id, use_mem_wal) + validate_uri_or_namespace(uri, namespace_impl, table_id, namespace_properties) + if uri is not None and not isinstance(uri, (str, pathlib.Path)): raise TypeError(f"Expected URI to be str or pathlib.Path, got {type(uri)}") - self._table_uri = str(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 - self._storage_options = ( - storage_options - if storage_options is not None - else io_config_to_storage_options(self._io_config, self._table_uri) - ) + 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, @@ -77,18 +90,48 @@ 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._managed_versioning = False + 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._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._storage_options = self._merged_storage_options(resolved) + self._managed_versioning = resolved.managed_versioning + existing = self._absorb_existing_dataset() existing_version = getattr(existing, "data_storage_version", None) if existing is not None else None self._data_storage_version = resolve_storage_version( - self._blob.apply_blob_v2_default(data_storage_version), + self._requested_storage_version, existing_version, self._mode, ) @@ -100,14 +143,48 @@ 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, Any]: + return get_namespace_kwargs(self._namespace_impl, self._namespace_properties, self._table_id) + + @property + def _namespace_commit_kwargs(self) -> dict[str, Any]: + return get_namespace_commit_kwargs( + self._namespace_impl, + self._namespace_properties, + self._table_id, + self._managed_versioning, + ) + + @property + def _dataset_uri_arg(self) -> str | None: + return None if self._namespace_impl is not None and self._table_id is not None else self._table_uri + + 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 is None: + raise ValueError("Unable to resolve Lance dataset URI from namespace.") + return resolved + + 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, explicitly provided options (including ``{}``) replace + the io_config-derived ones, preserving the historical sink behavior. + """ + io_derived = io_config_to_storage_options(self._io_config, resolved.uri) + if self._uri is not None: + return self._user_storage_options if self._user_storage_options is not None else io_derived + return merge_storage_options(io_derived, self._user_storage_options, resolved.storage_options) @staticmethod def _reject_unsupported_modes( @@ -130,6 +207,24 @@ def _reject_unsupported_modes( stacklevel=3, ) + @staticmethod + def _reject_namespace_mem_wal(namespace_impl: str | None, table_id: list[str] | None, use_mem_wal: bool) -> None: + """Reject the namespace + mem-WAL combination instead of failing mid-write. + + The namespace write path declares the table up front (a metadata-only + reservation), but the mem-WAL path then asks Lance for a namespace-aware + ``write_dataset(mode="create")``, which declares the same table a second + time and raises ``TableAlreadyExistsError``. Supporting this needs the + mem-WAL path to skip our own declare and let the native create own it; + until then, fail loudly at construction time. + """ + if use_mem_wal and namespace_impl is not None and table_id is not None: + raise ValueError( + "use_mem_wal=True is not supported with namespace-addressed tables " + "('namespace_impl' + 'table_id'). Write to a 'uri' instead, or set " + "use_mem_wal=False." + ) + def _init_lance_knobs( self, *, @@ -167,7 +262,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): @@ -179,21 +276,25 @@ 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.") 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": - raise ValueError("Cannot create a Lance dataset at a location where one already exists.") + 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, 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" @@ -218,6 +319,7 @@ 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) fragments = lance.fragment.write_fragments( wrapped, @@ -230,6 +332,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). @@ -239,8 +342,9 @@ 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._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 +354,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() @@ -324,17 +429,22 @@ 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, read_version=self._version, storage_options=self._storage_options, + **self._namespace_commit_kwargs, ) stats = dataset.stats.dataset_stats() stats_dict = MicroPartition.from_pydict( @@ -348,7 +458,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( @@ -357,9 +467,21 @@ def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadat self._mem_wal_total_bytes, ) from daft_lance.lance_compaction import compact_files_internal - - compact_files_internal(dataset) - dataset = lance.dataset(self._table_uri, storage_options=self._storage_options) + from daft_lance.namespace import DatasetOpenContext + + # Mem-WAL rejects namespace addressing up front (issue #54), so this + # is always a uri-only table and the context carries no triple. + compact_files_internal( + dataset, + DatasetOpenContext.from_dataset( + dataset, + str(self._dataset_uri_arg), + storage_options=self._storage_options, + ), + ) + dataset = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + ) stats = dataset.stats.dataset_stats() return MicroPartition.from_pydict( diff --git a/daft_lance/lance_merge_column.py b/daft_lance/lance_merge_column.py index 8984120..6546459 100644 --- a/daft_lance/lance_merge_column.py +++ b/daft_lance/lance_merge_column.py @@ -15,10 +15,10 @@ from daft.udf import method if TYPE_CHECKING: - import pathlib from collections.abc import Callable from daft.dependencies import pa + from daft_lance.namespace import DatasetOpenContext _FRAGMENT_HANDLER_RETURN_DTYPE = DataType.struct({"fragment_meta": DataType.binary(), "schema": DataType.binary()}) @@ -28,21 +28,28 @@ class FragmentHandler: def __init__( self, - lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, transform: dict[str, str] | lance.udf.BatchUDF | Callable[[pa.lib.RecordBatch], pa.lib.RecordBatch], read_columns: list[str] | None, reader_schema: pa.Schema | None = None, ): - self.lance_ds = lance_ds + self.open_context = open_context self.transform = transform self.read_columns = read_columns self.reader_schema = reader_schema + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds @method.batch(return_dtype=_FRAGMENT_HANDLER_RETURN_DTYPE) def __call__(self, fragment_ids: Any) -> list[dict[str, bytes]]: results = [] + lance_ds = self._dataset() for fragment_id in fragment_ids: - fragment = self.lance_ds.get_fragment(fragment_id) + fragment = lance_ds.get_fragment(fragment_id) if fragment is None: raise ValueError(f"Fragment {fragment_id} not found in dataset") fragment_meta, schema = fragment.merge_columns(self.transform, self.read_columns, None, self.reader_schema) @@ -52,12 +59,11 @@ def __call__(self, fragment_ids: Any) -> list[dict[str, bytes]]: def merge_columns_internal( lance_ds: lance.LanceDataset, - url: str | pathlib.Path, + open_context: DatasetOpenContext, *, transform: dict[str, str] | lance.udf.BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch], read_columns: list[str] | None = None, reader_schema: pa.Schema | None = None, - storage_options: dict[str, Any] | None = None, daft_remote_args: dict[str, Any] | None = None, concurrency: int | None = None, ) -> lance.LanceDataset: @@ -72,7 +78,7 @@ def merge_columns_internal( # Instantiate the Daft class with Lance-specific state and apply the # batch method over the fragment_id column. - handler = FragmentHandler(lance_ds, transform, read_columns, reader_schema) + handler = FragmentHandler(open_context, transform, read_columns, reader_schema) df = df.with_column("commit_message", handler(df["fragment_id"])) # type: ignore[arg-type] commit_messages = df.collect().to_pydict()["commit_message"] @@ -89,10 +95,11 @@ def merge_columns_internal( raise ValueError("No schema for new fragment found") op = lance.LanceOperation.Merge(fragment_metas, new_schema) return lance_ds.commit( - url, + open_context.uri, op, read_version=lance_ds.version, - storage_options=storage_options, + storage_options=open_context.storage_options, + **open_context.commit_kwargs, ) @@ -100,7 +107,7 @@ def merge_columns_internal( class GroupFragmentMergeUDF: def __init__( self, - lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, left_on: str | None = "_rowaddr", right_on: str | None = None, read_columns: list[str] | None = None, @@ -110,19 +117,25 @@ def __init__( """Per-group merge handler that directly invokes Lance fragment.merge with keyed join. Args: - lance_ds: Target Lance dataset. + open_context: Serializable handle the worker reopens the target dataset from. left_on: Key column on the Lance fragment (default "_rowaddr"). right_on: Key column name present in the provided reader data (defaults to left_on). read_columns: Names for columns provided to the handler via map_groups (must include right_on). reader_schema: Optional Arrow schema for the reader. batch_size: Optional batch size when building RecordBatchReader from the provided data. """ - self.lance_ds = lance_ds + self.open_context = open_context self.left_on = left_on or "_rowaddr" self.right_on = right_on or self.left_on self.read_columns = read_columns or [] self.reader_schema = reader_schema self.batch_size = batch_size + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds @method.batch(return_dtype=_FRAGMENT_HANDLER_RETURN_DTYPE) def __call__(self, *cols: Any) -> list[dict[str, bytes]]: @@ -192,16 +205,17 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: # Enforce that reader stream contains only join key + new columns (exclude existing dataset fields) df_schema = tbl.schema + lance_ds = self._dataset() existing_fields: set[str] = set() try: - existing_fields = {getattr(f, "name", str(f)) for f in self.lance_ds.schema} + existing_fields = {getattr(f, "name", str(f)) for f in lance_ds.schema} except Exception: names = [] try: - names = list(getattr(self.lance_ds.schema, "names", [])) + names = list(getattr(lance_ds.schema, "names", [])) except Exception: try: - names = [getattr(f, "name", str(f)) for f in getattr(self.lance_ds.schema, "fields", [])] + names = [getattr(f, "name", str(f)) for f in getattr(lance_ds.schema, "fields", [])] except Exception: names = [] existing_fields = set(names) @@ -219,7 +233,7 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: batches = tbl.to_batches(max_chunksize=self.batch_size) if self.batch_size is not None else tbl.to_batches() reader = _pa.RecordBatchReader.from_batches(tbl.schema, batches) - fragment = self.lance_ds.get_fragment(frag_id) + fragment = lance_ds.get_fragment(frag_id) if fragment is None: raise ValueError(f"Fragment {frag_id} not found in dataset") # Build schema argument: use the table's schema (including join key and new columns) unless an explicit reader_schema is provided @@ -239,15 +253,17 @@ class FastPathFragmentWriter: def __init__( self, - lance_ds: lance.LanceDataset, - uri: str, + open_context: DatasetOpenContext, new_column_names: list[str], - storage_options: dict[str, str] | None = None, ): - self.lance_ds = lance_ds - self.uri = str(uri) + self.open_context = open_context self.new_column_names = new_column_names - self.storage_options = storage_options + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds @method.batch(return_dtype=_FRAGMENT_HANDLER_RETURN_DTYPE) def __call__(self, *cols: Any) -> list[dict[str, bytes]]: @@ -291,7 +307,8 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: # Determine the existing file format version so the new file matches. # Lance commit rejects fragments whose files mix major/minor versions. - fragment = self.lance_ds.get_fragment(frag_id) + lance_ds = self._dataset() + fragment = lance_ds.get_fragment(frag_id) if fragment is None: raise ValueError(f"Fragment {frag_id} not found in dataset") meta = dict(fragment.metadata.to_json()) @@ -303,12 +320,12 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: # Write raw .lance file with only new columns filename = uuid.uuid4().hex + ".lance" - filepath = os.path.join(self.uri, "data", filename) + filepath = os.path.join(self.open_context.uri, "data", filename) with LanceFileWriter( filepath, tbl.schema, version=f"{file_major}.{file_minor}", - storage_options=self.storage_options, + storage_options=self.open_context.storage_options, ) as writer: for b in tbl.to_batches(): writer.write_batch(b) @@ -317,7 +334,7 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: # Determine field IDs for the new columns. Lance's manifest-level # max_field_id includes nested child fields and field IDs from dropped # columns, so it is the correct high-water mark for dataset evolution. - next_fid = self.lance_ds.max_field_id + 1 + next_fid = lance_ds.max_field_id + 1 # Stitch new data file into fragment metadata new_file_entry = { @@ -333,7 +350,7 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: new_frag_meta = FragmentMetadata.from_json(json.dumps(meta)) # Build new schema (original + new columns) - new_schema = self.lance_ds.schema + new_schema = lance_ds.schema for col_name in self.new_column_names: col_idx = tbl.schema.get_field_index(col_name) new_schema = new_schema.append(_pa.field(col_name, tbl.schema.field(col_idx).type)) @@ -364,11 +381,10 @@ def _can_use_fast_path( def merge_columns_from_df( df: daft.DataFrame, lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, *, read_columns: list[str] | None = None, reader_schema: pa.Schema | None = None, - storage_options: dict[str, Any] | None = None, daft_remote_args: dict[str, Any] | None = None, concurrency: int | None = None, left_on: str | None = "_rowaddr", @@ -413,30 +429,33 @@ def merge_columns_from_df( use_fast_path = _can_use_fast_path(df, lance_ds, join_key) if use_fast_path: - return _merge_fast_path(df, lance_ds, uri, new_cols, storage_options=storage_options) + return _merge_fast_path( + df, + lance_ds, + open_context, + new_cols, + ) else: return _merge_slow_path( df, lance_ds, - uri, + open_context, read_columns, left_on, right_on, reader_schema, batch_size, - storage_options=storage_options, ) def _merge_fast_path( df: daft.DataFrame, lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, new_column_names: list[str], - storage_options: dict[str, Any] | None = None, ) -> lance.LanceDataset: """Metadata-only add_columns: write raw .lance files and stitch into fragment metadata.""" - handler = FastPathFragmentWriter(lance_ds, str(uri), new_column_names, storage_options=storage_options) + handler = FastPathFragmentWriter(open_context, new_column_names) grouped = df.groupby("fragment_id").map_groups( handler(*(df[c] for c in new_column_names), df["_rowaddr"], df["fragment_id"]).alias("commit_message") # type: ignore[attr-defined] @@ -469,27 +488,27 @@ def _merge_fast_path( op = lance.LanceOperation.Merge(fragment_metas, LanceSchema.from_pyarrow(new_schema)) return lance.LanceDataset.commit( - str(uri), + open_context.uri, op, read_version=lance_ds.version, - storage_options=storage_options, + storage_options=open_context.storage_options, + **open_context.commit_kwargs, ) def _merge_slow_path( df: daft.DataFrame, lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, read_columns: list[str], left_on: str | None, right_on: str | None, reader_schema: pa.Schema | None, batch_size: int | None, - storage_options: dict[str, Any] | None = None, ) -> lance.LanceDataset: """Original keyed-join merge path: rewrites fragment data.""" handler_udf = GroupFragmentMergeUDF( - lance_ds, + open_context, left_on, right_on, read_columns, @@ -517,8 +536,9 @@ def _merge_slow_path( return lance_ds op = lance.LanceOperation.Merge(fragment_metas, new_schema) return lance_ds.commit( - uri, + open_context.uri, op, read_version=lance_ds.version, - storage_options=storage_options, + storage_options=open_context.storage_options, + **open_context.commit_kwargs, ) diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index bd8e60c..d29fd23 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -9,7 +9,7 @@ from daft import execution_config_ctx, from_pylist if TYPE_CHECKING: - import pathlib + from daft_lance.namespace import DatasetOpenContext import lance @@ -27,7 +27,7 @@ class FragmentIndexHandler: def __init__( self, - lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, column: str, index_type: str, name: str, @@ -35,13 +35,19 @@ def __init__( replace: bool, **kwargs: Any, ) -> None: - self.lance_ds = lance_ds + self.open_context = open_context self.column = column self.index_type = index_type self.name = name self.fragment_uuid = fragment_uuid self.replace = replace self.kwargs = kwargs + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds def __call__(self, fragment_ids: list[int]) -> bool: """Process a batch of fragment IDs for scalar index creation.""" @@ -50,7 +56,7 @@ def __call__(self, fragment_ids: list[int]) -> bool: fragment_ids, ) - self.lance_ds.create_scalar_index( + self._dataset().create_scalar_index( column=self.column, index_type=self.index_type, # type: ignore[arg-type] name=self.name, @@ -74,17 +80,23 @@ class SegmentedFragmentIndexHandler: def __init__( self, - lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, column: str, index_type: str, name: str, **kwargs: Any, ) -> None: - self.lance_ds = lance_ds + self.open_context = open_context self.column = column self.index_type = index_type self.name = name self.kwargs = kwargs + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds def __call__(self, fragment_ids: list[int], shard_id: int | None = None) -> bytes: """Build an independent index segment and return its pickled metadata.""" @@ -104,7 +116,7 @@ def __call__(self, fragment_ids: list[int], shard_id: int | None = None) -> byte # scalar index segments through this public API. Segment creation always # uses ``replace=False`` because replacement, if supported, must happen # in the final manifest commit rather than independently in each worker. - index_meta = self.lance_ds.create_index_uncommitted( + index_meta = self._dataset().create_index_uncommitted( column=self.column, index_type=self.index_type, name=self.name, @@ -132,13 +144,12 @@ def _existing_index_names(lance_ds: lance.LanceDataset) -> set[str]: def create_scalar_index_internal( lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, *, column: str, index_type: str = "INVERTED", name: str | None = None, replace: bool = False, - storage_options: dict[str, Any] | None = None, fragment_group_size: int | None = None, num_partitions: int | None = None, max_concurrency: int | None = None, @@ -147,6 +158,10 @@ def create_scalar_index_internal( ) -> None: """Internal implementation of distributed scalar index creation. + ``lance_ds`` is the driver's live dataset (planning, validation, commits); + ``open_context`` is the serializable handle workers reopen from and the + single source of uri, storage options and namespace kwargs. + When ``segmented=True``, ``BITMAP``, ``BTREE``, and ``INVERTED`` use Lance's public segment-index workflow: each worker builds a fully independent index segment, and the coordinator commits them atomically with @@ -272,7 +287,7 @@ def create_scalar_index_internal( # Configure maximum concurrency for fragment batches if not fragment_data: - logger.info("No fragments found for dataset at %s; skipping scalar index creation.", uri) + logger.info("No fragments found for dataset at %s; skipping scalar index creation.", open_context.uri) return logger.info( @@ -290,12 +305,10 @@ def create_scalar_index_internal( # older/unsupported distributed scalar index types. if segmented: _create_segmented_index( - lance_ds=lance_ds, - uri=uri, + open_context=open_context, column=column, index_type=index_type, name=name, - storage_options=storage_options, fragment_data=fragment_data, fragment_ids_to_use=fragment_ids_to_use, num_partitions=num_partitions, @@ -304,13 +317,11 @@ def create_scalar_index_internal( ) else: _create_partitioned_index( - lance_ds=lance_ds, - uri=uri, + open_context=open_context, column=column, index_type=index_type, name=name, replace=replace, - storage_options=storage_options, fragment_data=fragment_data, fragment_ids_to_use=fragment_ids_to_use, num_partitions=num_partitions, @@ -320,13 +331,11 @@ def create_scalar_index_internal( def _create_segmented_index( - lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, *, column: str, index_type: str, name: str, - storage_options: dict[str, Any] | None, fragment_data: list[dict[str, list[int]]], fragment_ids_to_use: list[int], num_partitions: int | None, @@ -345,7 +354,7 @@ def _create_segmented_index( max_concurrency=max_concurrency, ) handler = handler_cls( - lance_ds=lance_ds, + open_context=open_context, column=column, index_type=index_type, name=name, @@ -386,7 +395,7 @@ def _create_segmented_index( # Reload dataset to pick up the latest version (segment files were written # by workers against the version that was current at their invocation time). - lance_ds = lance.LanceDataset(uri, storage_options=storage_options) + lance_ds = open_context.open_latest() index_metas = _prepare_index_segments_for_commit(lance_ds, index_type, index_metas) logger.info( @@ -413,14 +422,12 @@ def _prepare_index_segments_for_commit( def _create_partitioned_index( - lance_ds: lance.LanceDataset, - uri: str | pathlib.Path, + open_context: DatasetOpenContext, *, column: str, index_type: str, name: str, replace: bool, - storage_options: dict[str, Any] | None, fragment_data: list[dict[str, list[int]]], fragment_ids_to_use: list[int], num_partitions: int | None, @@ -441,7 +448,7 @@ def _create_partitioned_index( max_concurrency=max_concurrency, ) handler = handler_cls( - lance_ds=lance_ds, + open_context=open_context, column=column, index_type=index_type, name=name, @@ -460,7 +467,7 @@ 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) + lance_ds = open_context.open_latest() lance_ds.merge_index_metadata(index_id, index_type) logger.info("Starting atomic index creation and commit") @@ -502,10 +509,11 @@ def _create_partitioned_index( # Commit the index operation atomically lance.LanceDataset.commit( - uri, + open_context.uri, create_index_op, read_version=lance_ds.version, - storage_options=storage_options, + storage_options=open_context.storage_options, + **open_context.commit_kwargs, ) logger.info("Index %s created successfully with ID %s", name, index_id) diff --git a/daft_lance/lance_scan.py b/daft_lance/lance_scan.py index e0ea8fb..18cb957 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 open_dataset_from_open_kwargs from .point_lookup import detect_point_lookup_columns from .utils import combine_filters_to_arrow @@ -40,7 +41,7 @@ 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 {})) + 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 @@ -117,7 +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.""" - ds = lance.dataset(ds_uri, **(open_kwargs or {})) + ds = open_dataset_from_open_kwargs(ds_uri, open_kwargs) logger.debug("Using metadata for counting all rows") count = ds.count_rows(filter=filter) @@ -134,8 +135,12 @@ def __init__( ds: lance.LanceDataset, fragment_group_size: int | None = None, include_fragment_id: bool | None = False, + open_kwargs: dict[str, Any] | None = None, + default_scan_options: dict[str, Any] | None = None, ): self._ds = ds + self._open_kwargs = dict(open_kwargs or {}) + self._default_scan_options = default_scan_options self._pushed_filters: list[PyExpr] | None = None self._remaining_filters: list[PyExpr] | None = None self._fragment_group_size = fragment_group_size @@ -255,11 +260,10 @@ def _create_count_rows_scan_task(self, pushdowns: PyPushdowns) -> Iterator[ScanT """Create scan task for counting rows.""" fields = pushdowns.aggregation_required_column_names() new_schema = Schema.from_pyarrow_schema(pa.schema([pa.field(fields[0], pa.uint64())])) - open_kwargs = getattr(self._ds, "_lance_open_kwargs", None) yield ScanTask.python_factory_func_scan_task( module=_lancedb_count_result_function.__module__, func_name=_lancedb_count_result_function.__name__, - func_args=(self._ds.uri, open_kwargs, fields[0], self._combine_filters_to_arrow()), + func_args=(self._ds.uri, self._open_kwargs, fields[0], self._combine_filters_to_arrow()), schema=new_schema._schema, num_rows=1, size_bytes=None, @@ -275,7 +279,6 @@ def _create_scan_tasks_with_limit_and_no_filters( assert self._pushed_filters is None, "Expected no filters when creating scan tasks with limit and no filters" assert pushdowns.limit is not None, "Expected a limit when creating scan tasks with limit and no filters" - open_kwargs = getattr(self._ds, "_lance_open_kwargs", None) fragments = self._ds.get_fragments() remaining_limit = pushdowns.limit @@ -300,7 +303,7 @@ def _create_scan_tasks_with_limit_and_no_filters( func_name=_lancedb_table_factory_function.__name__, func_args=( self._ds.uri, - open_kwargs, + self._open_kwargs, [fragment.fragment_id], required_columns, None, @@ -320,7 +323,6 @@ def _create_regular_scan_tasks( self, pushdowns: PyPushdowns, required_columns: list[str] | None, nearest_option: dict[str, Any] | None = None ) -> Iterator[ScanTask]: """Create regular scan tasks without count pushdown.""" - open_kwargs = getattr(self._ds, "_lance_open_kwargs", None) fragments = self._ds.get_fragments() pushed_expr = self._combine_filters_to_arrow() @@ -335,7 +337,7 @@ def _python_factory_func_scan_task( func_name=_lancedb_table_factory_function.__name__, func_args=( self._ds.uri, - open_kwargs, + self._open_kwargs, fragment_ids, required_columns, pushed_expr, @@ -459,15 +461,15 @@ def _should_use_index_for_point_lookup(self) -> bool: def _nearest_default_option(self) -> dict[str, Any] | None: """Return the default nearest option configured on the Lance dataset, if any. - Prefer Daft-specific `_daft_default_scan_options` to preserve options stripped before `lance.dataset` (e.g., `nearest`). + Daft keeps the original options explicitly because unsupported planning + keys such as ``nearest`` are stripped before calling ``lance.dataset``. """ - default_opts = getattr(self._ds, "_daft_default_scan_options", None) + default_opts = self._default_scan_options if not isinstance(default_opts, dict): + # Preserve direct ``LanceDBScanOperator(ds)`` compatibility. This + # is Lance's own stored constructor option, not Daft metadata + # attached to the third-party object. default_opts = getattr(self._ds, "_default_scan_options", None) - if not isinstance(default_opts, dict): - open_kwargs = getattr(self._ds, "_lance_open_kwargs", None) - if isinstance(open_kwargs, dict): - default_opts = open_kwargs.get("default_scan_options") if not isinstance(default_opts, dict): return None diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py new file mode 100644 index 0000000..ea7c8a4 --- /dev/null +++ b/daft_lance/namespace.py @@ -0,0 +1,392 @@ +"""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 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. + +3. **Distributed workers reopen, they never receive an open dataset.** + ``LanceDataset.__reduce__`` only carries (uri, storage_options, version, + manifest, ...); ``_namespace_client``, ``_table_id`` and + ``_namespace_client_managed_versioning`` are assigned after construction and + are therefore dropped by pickle. A worker that received a pickled dataset + would silently lose its namespace identity and its managed-versioning commit + handler. :class:`DatasetOpenContext` carries the serializable facts instead + and reopens on the worker. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from functools import lru_cache +from typing import Any +from urllib.parse import unquote, urlparse + +import lance +from lance_namespace import DeclareTableResponse, DescribeTableResponse, LanceNamespace + +_NAMESPACE_CACHE_SIZE = int(os.environ.get("DAFT_LANCE_NAMESPACE_CACHE_SIZE", "16")) + + +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 + + +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: + """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) + + 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: + 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: + """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. + + The path is percent-decoded: namespaces vend URI-encoded locations (a table + under ``daft lance/`` comes back as ``daft%20lance/``), and writing to the + literal encoded form would create a different directory than the one + ``describe_table`` later resolves to. + """ + parsed = urlparse(location) + if parsed.scheme == "file": + return unquote(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 +) -> LanceNamespace: + import lance_namespace as ln + + namespace_properties = dict(namespace_properties_tuple) if namespace_properties_tuple else {} + return ln.connect(namespace_impl, namespace_properties) + + +def get_or_create_namespace( + namespace_impl: str | None, namespace_properties: dict[str, str] | None +) -> LanceNamespace | 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 + return _get_cached_namespace(namespace_impl, namespace_properties_tuple) + + +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 {} + return {"namespace_client": namespace, "table_id": table_id} + + +def get_namespace_commit_kwargs( + namespace_impl: str | None, + namespace_properties: dict[str, str] | None, + table_id: list[str] | None, + managed_versioning: bool, +) -> dict[str, Any]: + """Namespace kwargs for commit APIs, including catalog-managed versioning.""" + kwargs = get_namespace_kwargs(namespace_impl, namespace_properties, table_id) + if kwargs: + kwargs["namespace_client_managed_versioning"] = managed_versioning + return kwargs + + +# `lance.fragment.write_fragments` accepts the same namespace kwargs as `lance.dataset`. +get_write_fragments_kwargs = get_namespace_kwargs + + +def _storage_options(response: DescribeTableResponse | DeclareTableResponse) -> dict[str, str] | None: + storage_options = getattr(response, "storage_options", None) + if storage_options is None: + return None + return dict(storage_options) + + +def _response_location(response: DescribeTableResponse | DeclareTableResponse) -> str: + location = getattr(response, "location", None) or getattr(response, "table_uri", None) + if not location: + raise ValueError("Namespace response did not include a table location.") + return _normalize_file_uri(str(location)) + + +@dataclass(frozen=True) +class ResolvedNamespaceTable: + """A namespace table resolved to its physical access and versioning policy.""" + + uri: str + storage_options: dict[str, str] | None = None + managed_versioning: bool = False + + +def _resolved_from_response(response: DescribeTableResponse | DeclareTableResponse) -> ResolvedNamespaceTable: + return ResolvedNamespaceTable( + uri=_response_location(response), + storage_options=_storage_options(response), + managed_versioning=getattr(response, "managed_versioning", None) is True, + ) + + +def _describe_table(namespace: LanceNamespace, table_id: list[str]) -> DescribeTableResponse: + from lance_namespace import DescribeTableRequest + + # vend_credentials is explicit: when unset, whether the namespace returns + # storage credentials is implementation-defined. + return namespace.describe_table(DescribeTableRequest(id=table_id, vend_credentials=True)) + + +def _declare_table(namespace: LanceNamespace, table_id: list[str]) -> ResolvedNamespaceTable: + from lance_namespace import DeclareTableRequest + + response = namespace.declare_table(DeclareTableRequest(id=table_id, location=None, vend_credentials=True)) + return _resolved_from_response(response) + + +def resolve_namespace_table( + *, + namespace_impl: str | None, + namespace_properties: dict[str, str] | None, + table_id: list[str] | None, + mode: str = "read", +) -> ResolvedNamespaceTable | None: + """Resolve a namespace table to a :class:`ResolvedNamespaceTable`. + + ``mode="create"`` atomically declares a new table. ``mode="overwrite"`` + declares the table only when a typed ``TableNotFoundError`` is raised, and + falls back to a second describe if it loses the declare race; other modes + require the table to exist. + """ + namespace = get_or_create_namespace(namespace_impl, namespace_properties) + if namespace is None or table_id is None: + return None + + if mode == "create": + return _declare_table(namespace, table_id) + + from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError + + try: + return _resolved_from_response(_describe_table(namespace, table_id)) + except TableNotFoundError: + if mode != "overwrite": + raise + try: + return _declare_table(namespace, table_id) + except TableAlreadyExistsError: + # Lost a declare race: another writer declared the table between our + # describe and our declare. Overwrite targets whatever exists now, so + # re-describing is the correct recovery. ``create`` deliberately does + # not reach here -- for it the conflict is the answer. + return _resolved_from_response(_describe_table(namespace, table_id)) + + +def merge_storage_options(*layers: dict[str, Any] | None) -> dict[str, Any] | None: + """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. + + Returning ``None`` rather than ``{}`` for an empty result is load-bearing: + pylance only installs the namespace credential-refresh provider when the + initial storage options are ``Some`` (``python/src/dataset.rs``). An empty + dict would install a provider whose initial options carry no + ``expires_at_millis``, i.e. one that is treated as never expiring and so + never refreshes. + """ + merged: dict[str, Any] = {} + for layer in layers: + if layer: + merged.update(layer) + return merged or None + + +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 serialized open arguments.""" + return ( + open_kwargs.pop("namespace_impl", None), + open_kwargs.pop("namespace_properties", None), + open_kwargs.pop("table_id", None), + ) + + +@dataclass(frozen=True) +class DatasetOpenContext: + """Everything a distributed worker needs to reopen a table, and nothing else. + + Maintenance operations (compaction, scalar index, merge columns) plan on the + driver and execute on workers. The driver's live ``LanceDataset`` cannot be + shipped: pickling it drops the namespace client, the table id and the + managed-versioning flag (see this module's docstring), so a worker would + commit as if the table were an unmanaged uri table. + + This context is the serializable substitute. It deliberately excludes the + live dataset, the namespace client, ``serialized_manifest``, ``session``, + ``commit_lock`` and ``asof``: + + * the client is rebuilt per process from the triple via the ``lru_cache``; + * the manifest is omitted so the payload stays independent of fragment + count -- workers pay one pinned-manifest read instead; + * ``asof``/tags are already resolved by the driver into :attr:`version`, so + workers cannot drift onto a different snapshot. + """ + + uri: str + version: int + storage_options: dict[str, Any] | None = field(default=None, repr=False) + + namespace_impl: str | None = None + namespace_properties: dict[str, str] | None = None + table_id: list[str] | None = None + managed_versioning: bool = False + + block_size: int | None = None + index_cache_size: int | None = None + metadata_cache_size_bytes: int | None = None + default_scan_options: dict[str, Any] | None = field(default=None, repr=False) + + @classmethod + def from_dataset( + cls, + dataset: lance.LanceDataset, + uri: str, + *, + storage_options: dict[str, Any] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, + table_id: list[str] | None = None, + managed_versioning: bool = False, + block_size: int | None = None, + index_cache_size: int | None = None, + metadata_cache_size_bytes: int | None = None, + default_scan_options: dict[str, Any] | None = None, + ) -> DatasetOpenContext: + """Derive a context from the snapshot the driver actually opened.""" + return cls( + uri=uri, + version=dataset.version, + storage_options=storage_options, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, + managed_versioning=managed_versioning, + block_size=block_size, + index_cache_size=index_cache_size, + metadata_cache_size_bytes=metadata_cache_size_bytes, + default_scan_options=default_scan_options, + ) + + @property + def namespace_kwargs(self) -> dict[str, Any]: + return get_namespace_kwargs(self.namespace_impl, self.namespace_properties, self.table_id) + + @property + def commit_kwargs(self) -> dict[str, Any]: + return get_namespace_commit_kwargs( + self.namespace_impl, + self.namespace_properties, + self.table_id, + self.managed_versioning, + ) + + def _open(self, version: int | None) -> lance.LanceDataset: + """Open through the low-level constructor, passing the physical uri. + + ``lance.dataset(None, namespace_client=..., table_id=...)`` resolves the + location with a Python-side ``describe_table`` on *every* call, which + would put one namespace round-trip on every task. The driver already + resolved the location, so workers pass the uri directly and the + constructor still installs the namespace wiring (credential provider and, + when the catalog manages versioning, the commit handler). + """ + return lance.LanceDataset( + self.uri, + version=version, + storage_options=self.storage_options, + block_size=self.block_size, + index_cache_size=self.index_cache_size, + metadata_cache_size_bytes=self.metadata_cache_size_bytes, + default_scan_options=self.default_scan_options, + namespace_client=get_or_create_namespace(self.namespace_impl, self.namespace_properties), + table_id=self.table_id, + namespace_client_managed_versioning=self.managed_versioning, + ) + + def open_pinned(self) -> lance.LanceDataset: + """Open the exact snapshot the driver planned against. + + Every worker must see the same fragment ids and schema the coordinator + used to build its plan, so this is the default for worker code. + """ + return self._open(self.version) + + def open_latest(self) -> lance.LanceDataset: + """Open the newest version; only for coordinator steps that must observe worker output.""" + return self._open(None) + + +def open_dataset_from_open_kwargs(ds_uri: str | None, open_kwargs: dict[str, Any] | None) -> lance.LanceDataset: + """Re-open a dataset on a worker from serialized open arguments. + + 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 ds_uri, + **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), + **open_kwargs, + ) diff --git a/daft_lance/utils.py b/daft_lance/utils.py index 0eb32c7..c3cf79d 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -1,19 +1,93 @@ from __future__ import annotations import logging +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any 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 ( + DatasetOpenContext, + get_namespace_commit_kwargs, + get_namespace_kwargs, + has_namespace_params, + merge_storage_options, + resolve_namespace_table, + validate_uri_or_namespace, +) if TYPE_CHECKING: import pathlib + from daft.daft import IOConfig + logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class LanceDatasetHandle: + """An opened Lance dataset together with the context needed to reuse it. + + ``lance.LanceDataset`` does not expose all arguments used to open it. Keep + those arguments in this Daft-owned value object instead of attaching + private attributes to the third-party dataset instance. + """ + + dataset: lance.LanceDataset + uri: str + open_kwargs: dict[str, Any] = field(repr=False) + managed_versioning: bool = False + default_scan_options: dict[str, Any] | None = field(default=None, repr=False) + + @property + def storage_options(self) -> dict[str, Any] | None: + options = self.open_kwargs.get("storage_options") + return options if isinstance(options, dict) else None + + @property + def namespace_kwargs(self) -> dict[str, Any]: + return get_namespace_kwargs( + self.open_kwargs.get("namespace_impl"), + self.open_kwargs.get("namespace_properties"), + self.open_kwargs.get("table_id"), + ) + + @property + def commit_kwargs(self) -> dict[str, Any]: + return get_namespace_commit_kwargs( + self.open_kwargs.get("namespace_impl"), + self.open_kwargs.get("namespace_properties"), + self.open_kwargs.get("table_id"), + self.managed_versioning, + ) + + def worker_open_context(self) -> DatasetOpenContext: + """Derive the serializable context distributed maintenance workers reopen from. + + Note this pins :attr:`dataset`'s *resolved* version rather than the + caller's ``version``/``asof`` request, so every worker lands on the exact + snapshot the driver planned against. + """ + return DatasetOpenContext.from_dataset( + self.dataset, + self.uri, + storage_options=self.storage_options, + namespace_impl=self.open_kwargs.get("namespace_impl"), + namespace_properties=self.open_kwargs.get("namespace_properties"), + table_id=self.open_kwargs.get("table_id"), + managed_versioning=self.managed_versioning, + block_size=self.open_kwargs.get("block_size"), + index_cache_size=self.open_kwargs.get("index_cache_size"), + metadata_cache_size_bytes=self.open_kwargs.get("metadata_cache_size_bytes"), + # The stripped variant that was actually handed to pylance; the + # unstripped copy on this handle is for Daft planning only. + default_scan_options=self.open_kwargs.get("default_scan_options"), + ) + + def distribute_fragments_balanced(fragments: list[Any], fragment_group_size: int) -> list[dict[str, list[int]]]: """Distribute fragments across workers using a balanced algorithm considering fragment sizes.""" if fragment_group_size <= 0: @@ -81,13 +155,50 @@ def distribute_fragments_balanced(fragments: list[Any], fragment_group_size: int return non_empty_batches -def construct_lance_dataset( - uri: str | pathlib.Path, +def construct_lance_dataset_handle( + 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.""" +) -> LanceDatasetHandle: + """Construct a Lance dataset and retain its reusable open context. + + Storage options are layered from lowest to highest priority: + io_config-derived < user-provided ``storage_options`` < namespace-vended. + 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, namespace_properties) + resolved_uri = str(uri) if uri is not None else None + namespace_storage_options = None + managed_versioning = False + if resolved_uri is None: + resolved = resolve_namespace_table( + namespace_impl=namespace_impl, + 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 + managed_versioning = resolved.managed_versioning + if resolved_uri is None: + raise ValueError("Unable to resolve Lance dataset URI.") + + 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: + # 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) + original_default_scan_options = kwargs.pop("default_scan_options", None) safe_default_scan_options = None if isinstance(original_default_scan_options, dict): @@ -98,26 +209,41 @@ 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 + dataset = 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, - "version": version, + "storage_options": merged_storage_options, + # Pin the snapshot the driver resolved, not the caller's request. These + # kwargs cross to scan workers, and ``version=None`` there means "open + # latest": a compaction landing between planning and execution leaves + # workers looking for fragment ids that no longer exist, and an + # overwrite silently feeds them different data entirely. + "version": dataset.version, + "namespace_impl": namespace_impl, + "namespace_properties": namespace_properties, + "table_id": table_id, } effective_kwargs.update(kwargs or {}) - try: - ds._lance_open_kwargs = effective_kwargs # type: ignore[attr-defined] - except Exception: - pass - - # Preserve the full user-provided defaults (including nearest) for Daft's planning - # even if we stripped keys out before calling `lance.dataset`. - try: - ds._daft_default_scan_options = original_default_scan_options # type: ignore[attr-defined] - except Exception: - pass - - return ds + # ``asof``/tags are inputs to the version resolution above; carrying them + # further is redundant at best. (pylance lets ``version`` win when both are + # passed, so this is belt-and-braces.) + effective_kwargs.pop("asof", None) + return LanceDatasetHandle( + dataset=dataset, + uri=resolved_uri, + open_kwargs=effective_kwargs, + managed_versioning=managed_versioning, + # Preserve the full user-provided defaults (including nearest) for + # Daft's planning even if keys were stripped before calling Lance. + default_scan_options=original_default_scan_options if isinstance(original_default_scan_options, dict) else None, + ) def combine_filters_to_arrow(predicates: list[Any] | None) -> pa.compute.Expression | None: diff --git a/tests/io/lancedb/test_fast_path_merge.py b/tests/io/lancedb/test_fast_path_merge.py index 2947ad5..147eef2 100644 --- a/tests/io/lancedb/test_fast_path_merge.py +++ b/tests/io/lancedb/test_fast_path_merge.py @@ -17,6 +17,7 @@ import daft from daft_lance.lance_merge_column import _can_use_fast_path, merge_columns_from_df +from daft_lance.namespace import DatasetOpenContext # --------------------------------------------------------------------------- # Fixtures and helpers @@ -28,6 +29,11 @@ def ds_path(tmp_path_factory): yield str(tmp_path_factory.mktemp("fast_path")) +def open_ctx(ds: lance.LanceDataset, path: str) -> DatasetOpenContext: + """Uri-only worker context pinned to the snapshot the caller opened.""" + return DatasetOpenContext.from_dataset(ds, path) + + def create_dataset(path: str, fragments: list[dict]) -> lance.LanceDataset: for i, data in enumerate(fragments): table = pa.table(data) @@ -61,7 +67,7 @@ def test_fast_path_single_column_int(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2, 3], "val": [10, 20, 30]}]) df = read_with_metadata(ds_path) df = df.with_column("doubled", daft.col("val").cast(daft.DataType.int64()) * 2) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().to_pydict() assert result["doubled"] == [2 * v for v in result["val"]] @@ -69,7 +75,7 @@ def test_fast_path_single_column_float(self, ds_path): ds = create_dataset(ds_path, [{"x": [1.0, 2.0, 3.0]}]) df = read_with_metadata(ds_path) df = df.with_column("half", daft.col("x").cast(daft.DataType.float64()) / 2.0) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().to_pydict() for x, h in zip(result["x"], result["half"]): assert pytest.approx(x / 2.0, rel=1e-6) == h @@ -78,7 +84,7 @@ def test_fast_path_single_column_string(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2], "name": ["alice", "bob"]}]) df = read_with_metadata(ds_path) df = df.with_column("greeting", daft.lit("hello_") + daft.col("name")) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().to_pydict() assert result["greeting"] == ["hello_alice", "hello_bob"] @@ -88,7 +94,7 @@ def test_fast_path_multiple_new_columns(self, ds_path): df = df.with_column("a", daft.col("id").cast(daft.DataType.int64()) * 10) df = df.with_column("b", daft.col("id").cast(daft.DataType.float64()) + 0.5) df = df.with_column("c", daft.lit("row_") + daft.col("id").cast(daft.DataType.string())) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().to_pydict() assert result["a"] == [10, 20, 30] for i, b in zip(result["id"], result["b"]): @@ -106,7 +112,7 @@ def test_fast_path_multi_fragment(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("score", daft.col("val").cast(daft.DataType.float64()) * 1.5) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() for v, s in zip(result["val"], result["score"]): assert pytest.approx(v * 1.5, rel=1e-6) == s @@ -122,7 +128,7 @@ def test_fast_path_computed_column(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("z", daft.col("x").cast(daft.DataType.int64()) + daft.col("y").cast(daft.DataType.int64())) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("x").to_pydict() for x, y, z in zip(result["x"], result["y"], result["z"]): assert x + y == z @@ -141,7 +147,7 @@ def test_existing_columns_unchanged(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("new_col", daft.lit(999)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) after = ds2.to_table().sort_by("id").to_pydict() assert after["id"] == before["id"] @@ -158,7 +164,7 @@ def test_fragment_file_count(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("b", daft.lit(0)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) for frag in ds2.get_fragments(): assert len(list(frag.data_files())) == 2 @@ -172,7 +178,7 @@ def test_original_files_not_rewritten(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("b", daft.lit(42)) - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) for fname, old_hash in original_files.items(): fpath = os.path.join(data_dir, fname) @@ -192,7 +198,7 @@ def test_row_count_preserved(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("w", daft.lit(0)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.count_rows() == original_count for frag, expected in zip(ds2.get_fragments(), per_frag_counts): @@ -202,7 +208,7 @@ def test_schema_evolution_correct(self, ds_path): ds = create_dataset(ds_path, [{"id": [1], "name": ["x"]}]) df = read_with_metadata(ds_path) df = df.with_column("score", daft.lit(3.14)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) schema_names = set(ds2.schema.names) assert "id" in schema_names assert "name" in schema_names @@ -221,7 +227,7 @@ def test_rowaddr_sorting_restores_order(self, ds_path): df = read_with_metadata(ds_path) # Daft doesn't guarantee order, but let's force a known computation df = df.with_column("doubled", daft.col("val").cast(daft.DataType.int64()) * 2) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert result["doubled"] == [20, 40, 60, 80, 100] @@ -236,7 +242,7 @@ def test_multi_fragment_ordering(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("neg", daft.col("id").cast(daft.DataType.int64()) * -1) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() for i, n in zip(result["id"], result["neg"]): assert n == -i @@ -310,12 +316,12 @@ def test_two_sequential_merges(self, ds_path): # First merge: add column A df = read_with_metadata(ds_path) df = df.with_column("a", daft.col("id").cast(daft.DataType.int64()) * 10) - ds = merge_columns_from_df(df, ds, ds_path) + ds = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) # Second merge: add column B df2 = read_with_metadata(ds_path) df2 = df2.with_column("b", daft.col("id").cast(daft.DataType.int64()) * 100) - ds2 = merge_columns_from_df(df2, ds, ds_path) + ds2 = merge_columns_from_df(df2, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert result["a"] == [10, 20, 30] @@ -331,7 +337,7 @@ def test_merge_after_append(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("flag", daft.lit(True)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert result["flag"] == [True, True, True, True] assert len(result["id"]) == 4 @@ -342,14 +348,14 @@ def test_merge_preserves_previous_merge(self, ds_path): # Merge A df = read_with_metadata(ds_path) df = df.with_column("a", daft.col("val").cast(daft.DataType.int64()) + 1) - ds = merge_columns_from_df(df, ds, ds_path) + ds = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) check_a = ds.to_table().sort_by("id").to_pydict() assert check_a["a"] == [11, 21] # Merge B df2 = read_with_metadata(ds_path) df2 = df2.with_column("b", daft.col("val").cast(daft.DataType.int64()) + 2) - ds2 = merge_columns_from_df(df2, ds, ds_path) + ds2 = merge_columns_from_df(df2, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert result["a"] == [11, 21], "Column A corrupted by second merge" assert result["b"] == [12, 22] @@ -366,7 +372,7 @@ def test_type_int64(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2]}]) df = read_with_metadata(ds_path) df = df.with_column("x", daft.lit(42).cast(daft.DataType.int64())) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.schema.field("x").type == pa.int64() assert ds2.to_table().column("x").to_pylist() == [42, 42] @@ -376,7 +382,7 @@ def test_type_float32(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2]}]) df = read_with_metadata(ds_path) df = df.with_column("x", daft.lit(3.14).cast(daft.DataType.float32())) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) # Value is preserved even if type widens vals = ds2.to_table().column("x").to_pylist() assert all(pytest.approx(v, rel=1e-5) == 3.14 for v in vals) @@ -385,21 +391,21 @@ def test_type_float64(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2]}]) df = read_with_metadata(ds_path) df = df.with_column("x", daft.lit(2.718)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.schema.field("x").type == pa.float64() def test_type_string(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2]}]) df = read_with_metadata(ds_path) df = df.with_column("label", daft.lit("hello")) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.to_table().column("label").to_pylist() == ["hello", "hello"] def test_type_bool(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2, 3]}]) df = read_with_metadata(ds_path) df = df.with_column("flag", daft.col("id").cast(daft.DataType.int64()) > 1) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").column("flag").to_pylist() assert result == [False, True, True] @@ -412,7 +418,7 @@ def test_type_nullable(self, ds_path): (daft.col("id").cast(daft.DataType.int64()) % 2 != 0).cast(daft.DataType.int64()) * daft.col("id").cast(daft.DataType.int64()), ) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").column("maybe").to_pylist() # id=1 → odd → 1*1=1, id=2 → even → 0*2=0, id=3 → odd → 1*3=3, id=4 → even → 0*4=0 assert result == [1, 0, 3, 0] @@ -428,7 +434,7 @@ def test_single_row_fragment(self, ds_path): ds = create_dataset(ds_path, [{"id": [1]}]) df = read_with_metadata(ds_path) df = df.with_column("x", daft.lit(99)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.to_table().to_pydict() == {"id": [1], "x": [99]} def test_large_fragment(self, ds_path): @@ -436,7 +442,7 @@ def test_large_fragment(self, ds_path): ds = create_dataset(ds_path, [{"id": list(range(n))}]) df = read_with_metadata(ds_path) df = df.with_column("neg", daft.col("id").cast(daft.DataType.int64()) * -1) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert len(result["id"]) == n for i, neg in zip(result["id"], result["neg"]): @@ -449,7 +455,7 @@ def test_many_fragments(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("doubled", daft.col("id").cast(daft.DataType.int64()) * 2) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert len(result["id"]) == 20 for i, d in zip(result["id"], result["doubled"]): @@ -460,14 +466,14 @@ def test_empty_new_columns_raises(self, ds_path): df = read_with_metadata(ds_path) # No new columns added — only existing + metadata columns with pytest.raises(ValueError, match="No new columns"): - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) def test_dataset_version_incremented(self, ds_path): ds = create_dataset(ds_path, [{"id": [1]}]) v_before = ds.version df = read_with_metadata(ds_path) df = df.with_column("x", daft.lit(1)) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) assert ds2.version == v_before + 1 @@ -491,7 +497,7 @@ def test_fast_vs_slow_identical_results(self, tmp_path_factory): # Fast path: read with _rowaddr + fragment_id df_fast = read_with_metadata(fast_path) df_fast = df_fast.with_column("score", daft.col("val").cast(daft.DataType.float64()) * 2.5) - ds_fast = merge_columns_from_df(df_fast, ds_fast, fast_path) + ds_fast = merge_columns_from_df(df_fast, ds_fast, open_ctx(ds_fast, fast_path)) # Slow path: read with fragment_id only, use business key df_slow = daft.read_lance(slow_path, include_fragment_id=True, default_scan_options={"with_row_address": True}) @@ -501,13 +507,12 @@ def test_fast_vs_slow_identical_results(self, tmp_path_factory): ds_slow = _merge_slow_path( df_slow, ds_slow, - slow_path, + open_ctx(ds_slow, slow_path), read_columns=["_rowaddr", "score"], left_on="_rowaddr", right_on="_rowaddr", reader_schema=None, batch_size=None, - storage_options=None, ) fast_result = ds_fast.to_table().sort_by("id").to_pydict() @@ -534,7 +539,7 @@ def test_read_after_merge_with_filter(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("score", daft.col("val").cast(daft.DataType.int64()) * 2) - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) ds2 = lance.dataset(ds_path) filtered = ds2.to_table(filter="score > 40") @@ -546,7 +551,7 @@ def test_read_after_merge_with_projection(self, ds_path): ds = create_dataset(ds_path, [{"id": [1, 2], "val": [10, 20]}]) df = read_with_metadata(ds_path) df = df.with_column("new_col", daft.lit(42)) - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) ds2 = lance.dataset(ds_path) projected = ds2.to_table(columns=["new_col"]) @@ -558,7 +563,7 @@ def test_read_after_merge_select_original_only(self, ds_path): ds = create_dataset(ds_path, [original]) df = read_with_metadata(ds_path) df = df.with_column("extra", daft.lit("x")) - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) ds2 = lance.dataset(ds_path) result = ds2.to_table(columns=["id", "name"]).sort_by("id").to_pydict() @@ -575,7 +580,7 @@ def test_scan_fragments_individually(self, ds_path): ) df = read_with_metadata(ds_path) df = df.with_column("doubled", daft.col("id").cast(daft.DataType.int64()) * 2) - merge_columns_from_df(df, ds, ds_path) + merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) ds2 = lance.dataset(ds_path) all_ids = [] @@ -634,7 +639,7 @@ def _make_vec(ids): return [np.array([float(i)] * N, dtype=np.float32) for i in ids.to_pylist()] df = df.with_column("embedding", _make_vec(daft.col("id"))) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) field = ds2.schema.field("embedding") # Type must be preserved: fixed_size_list[N], NOT list @@ -676,7 +681,7 @@ def test_next_fid_uses_manifest_max_field_id(self, ds_path): df = read_with_metadata(ds_path) df = df.with_column("score", daft.col("id").cast(daft.DataType.int64()) * 10) - ds2 = merge_columns_from_df(df, ds, ds_path) + ds2 = merge_columns_from_df(df, ds, open_ctx(ds, ds_path)) result = ds2.to_table().sort_by("id").to_pydict() assert result["score"] == [10, 20, 30], f"score is null or wrong (field ID collision): {result['score']}" diff --git a/tests/io/lancedb/test_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..c559106 100644 --- a/tests/io/lancedb/test_lance_data_sink_internals.py +++ b/tests/io/lancedb/test_lance_data_sink_internals.py @@ -1,7 +1,7 @@ """Unit-level tests for LanceDataSink internals. -Phase 1 covers the _load_existing_dataset typed-exception path and the Lance -message-format pin. +These tests also pin the sink lifecycle contract: construction validates local +arguments only, while target-dependent validation runs in ``start()``. """ from __future__ import annotations @@ -17,17 +17,39 @@ 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 +def test_uri_create_existing_error_is_deferred_until_start(tmp_path): + uri = str(tmp_path / "existing") + lance.write_dataset(pa.table({"a": [1]}), uri) + + sink = LanceDataSink(uri=uri, schema=pa.schema([("a", pa.int64())]), mode="create") + + with pytest.raises(ValueError, match="already exists"): + sink.start() + + +def test_uri_append_schema_error_is_deferred_until_start(tmp_path): + uri = str(tmp_path / "schema-mismatch") + lance.write_dataset(pa.table({"a": pa.array([1], type=pa.int64())}), uri) + + sink = LanceDataSink(uri=uri, schema=pa.schema([("a", pa.struct([("value", pa.string())]))]), mode="append") + + with pytest.raises(ValueError, match="Schema of data does not match"): + sink.start() + + def test_mode_merge_rejected_at_construction(tmp_path): schema = pa.schema([("a", pa.int64())]) with pytest.raises(ValueError, match='mode="merge" is no longer supported'): @@ -143,6 +165,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 +184,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_lancedb_point_lookup.py b/tests/io/lancedb/test_lancedb_point_lookup.py index 92b07ff..63d2dd7 100644 --- a/tests/io/lancedb/test_lancedb_point_lookup.py +++ b/tests/io/lancedb/test_lancedb_point_lookup.py @@ -104,7 +104,7 @@ def test_scanner_without_fragments(lance_dataset, idx_type): # Invoke factory with fragment_ids=None to exercise index-driven fragment selection gen = lance_scan._lancedb_table_factory_function( ds.uri, - getattr(ds, "_lance_open_kwargs", None), + {}, None, ["id", "value"], arrow_filter, diff --git a/tests/io/lancedb/test_lancedb_scalar_index.py b/tests/io/lancedb/test_lancedb_scalar_index.py index 1b325e6..eeb4280 100644 --- a/tests/io/lancedb/test_lancedb_scalar_index.py +++ b/tests/io/lancedb/test_lancedb_scalar_index.py @@ -17,6 +17,27 @@ _prepare_index_segments_for_commit, create_scalar_index_internal, ) +from daft_lance.namespace import DatasetOpenContext + + +class FakeOpenContext: + """Stands in for DatasetOpenContext so handlers can be driven against a fake dataset. + + Also counts opens, which is what proves a handler reopens once per instance + rather than once per call. + """ + + def __init__(self, dataset, uri="memory://fake"): + self.dataset = dataset + self.uri = uri + self.opens = 0 + + def open_pinned(self): + self.opens += 1 + return self.dataset + + def open_latest(self): + return self.open_pinned() @pytest.fixture @@ -621,7 +642,7 @@ def create_index_uncommitted(self, **kwargs): fake_ds = FakeLanceDataset() handler = SegmentedFragmentIndexHandler( - lance_ds=fake_ds, + open_context=FakeOpenContext(fake_ds), column="price", index_type="BTREE", name="price_idx", @@ -704,7 +725,7 @@ def create_index_uncommitted(self, **kwargs): fake_ds = FakeLanceDataset() handler = SegmentedFragmentIndexHandler( - lance_ds=fake_ds, + open_context=FakeOpenContext(fake_ds), column="flag", index_type="BITMAP", name="flag_idx", @@ -772,7 +793,7 @@ def fake_create_segmented_index(**kwargs): create_scalar_index_internal( lance_ds=FakeLanceDataset(), - uri="memory://bitmap", + open_context=DatasetOpenContext(uri="memory://bitmap", version=1), column="flag", index_type="BITMAP", name="flag_bitmap_idx", @@ -804,7 +825,7 @@ def create_scalar_index(self, **kwargs): create_scalar_index_internal( lance_ds=fake_ds, - uri="memory://bitmap", + open_context=DatasetOpenContext(uri="memory://bitmap", version=1), column="flag", index_type="BITMAP", name="flag_bitmap_idx", @@ -1067,7 +1088,7 @@ def fake_create_partitioned_index(**kwargs): create_scalar_index_internal( lance_ds=FakeLanceDataset(), - uri="memory://btree", + open_context=DatasetOpenContext(uri="memory://btree", version=1), column="price", index_type="BTREE", name="price_btree_idx", diff --git a/tests/io/lancedb/test_lancedb_vector_search.py b/tests/io/lancedb/test_lancedb_vector_search.py index 54d4077..babd899 100755 --- a/tests/io/lancedb/test_lancedb_vector_search.py +++ b/tests/io/lancedb/test_lancedb_vector_search.py @@ -9,6 +9,7 @@ import daft from daft import col +from daft_lance.lance_scan import LanceDBScanOperator def build_single_fragment_dataset(tmp_path_factory) -> str: @@ -115,6 +116,16 @@ def test_nearest_global_single_scan_task(tmp_path_factory) -> None: assert "Num Scan Tasks = 1" in explain_output +def test_direct_scan_operator_uses_lance_default_nearest(tmp_path_factory) -> None: + dataset_path = build_multi_fragment_dataset(tmp_path_factory) + nearest = {"column": "vector", "q": pa.array([0.0, 0.0], type=pa.float32()), "k": 1} + dataset = lance.dataset(dataset_path, default_scan_options={"nearest": nearest}) + + scan = LanceDBScanOperator(dataset) + + assert scan._nearest_default_option() == nearest + + def test_nearest_metric_cosine_k1(tmp_path_factory) -> None: dataset_path = build_metric_dataset(tmp_path_factory) diff --git a/tests/io/lancedb/test_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 new file mode 100644 index 0000000..991fce0 --- /dev/null +++ b/tests/io/lancedb/test_namespace.py @@ -0,0 +1,946 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +import daft +import daft_lance +import daft_lance.namespace as namespace_mod +from daft_lance.lance_data_sink import LanceDataSink +from daft_lance.utils import construct_lance_dataset_handle + + +def _dir_ns(tmp_path: Path) -> dict[str, Any]: + return {"namespace_impl": "dir", "namespace_properties": {"root": str(tmp_path)}} + + +def _double_score(batch: Any) -> Any: + import pyarrow as pa + import pyarrow.compute as pc + + return pa.RecordBatch.from_arrays([pc.multiply(batch["score"], 2)], ["doubled"]) + + +@pytest.mark.parametrize( + ("location", "expected"), + [ + ("file:///tmp/plain/t.lance", "/tmp/plain/t.lance"), + ("file:///tmp/root%20space/t.lance", "/tmp/root space/t.lance"), + ("file:///tmp/%E4%B8%AD%E6%96%87/t.lance", "/tmp/中文/t.lance"), + ("file:///tmp/100%25/t.lance", "/tmp/100%/t.lance"), + # Object-store URIs pass through untouched, encoding included. + ("s3://bucket/root%20space/t.lance", "s3://bucket/root%20space/t.lance"), + ], +) +def test_normalize_file_uri_percent_decodes_paths(location: str, expected: str) -> None: + assert namespace_mod._normalize_file_uri(location) == expected + + +@pytest.mark.parametrize("root_name", ["daft lance", "中文目录"]) +def test_namespace_roundtrip_with_encoded_location(tmp_path: Path, root_name: str) -> None: + """A namespace root/table whose location percent-encodes must still round-trip. + + The dir namespace vends ``file://.../daft%20lance/table%20space.lance``; writing + to the literal encoded form would create a directory that the later + ``describe_table`` never resolves to. + """ + root = tmp_path / root_name + root.mkdir() + ns = _dir_ns(root) + table_id = ["table space"] + + df = daft.from_pydict({"id": [1, 2], "label": ["a", "b"]}) + daft_lance.write_lance(df, table_id=table_id, mode="create", **ns).collect() + + assert daft_lance.read_lance(table_id=table_id, **ns).to_pydict() == {"id": [1, 2], "label": ["a", "b"]} + # The data landed under the decoded name, not a literal "%20" sibling. + assert (root / "table space.lance").is_dir() + assert not (root / "table%20space.lance").exists() + + +def test_namespace_write_read_append_roundtrip(tmp_path: Path) -> None: + ns = _dir_ns(tmp_path) + table_id = ["roundtrip"] + + df1 = daft.from_pydict({"id": [1, 2], "label": ["a", "b"]}) + df2 = daft.from_pydict({"id": [3], "label": ["c"]}) + + 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: Path) -> None: + 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: Path) -> None: + 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_overwrite_does_not_declare_on_ambiguous_error( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + class FakeNamespace: + declared = False + + def describe_table(self, request: Any) -> Any: + raise RuntimeError("permission denied: parent catalog does not exist") + + def declare_table(self, request: Any) -> Any: + 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_read_supports_pushdowns(tmp_path: Path) -> None: + 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() + + 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: Path) -> None: + 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_merge_columns_df(tmp_path: Path) -> None: + ns = _dir_ns(tmp_path) + table_id = ["merge_cols_df"] + + daft_lance.write_lance( + daft.from_pydict({"id": [1, 2, 3], "score": [1, 2, 3]}), + 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("tripled", df["score"] * 3) + daft_lance.merge_columns_df(df.select("fragment_id", "_rowaddr", "tripled"), table_id=table_id, **ns) + + result = daft_lance.read_lance(table_id=table_id, **ns).sort("id").to_pydict() + assert result["tripled"] == [3, 6, 9] + + +def test_namespace_merge_columns_transform(tmp_path: Path) -> None: + ns = _dir_ns(tmp_path) + table_id = ["merge_cols_transform"] + + daft_lance.write_lance( + daft.from_pydict({"id": [1, 2, 3], "score": [10, 20, 30]}), + table_id=table_id, + mode="create", + **ns, + ).collect() + + daft_lance.merge_columns(table_id=table_id, transform=_double_score, read_columns=["score"], **ns) + + result = daft_lance.read_lance(table_id=table_id, **ns).sort("id").to_pydict() + assert result["doubled"] == [20, 40, 60] + + +def test_namespace_merge_columns_df_slow_path(tmp_path: Path) -> None: + ns = _dir_ns(tmp_path) + table_id = ["merge_cols_slow"] + + daft_lance.write_lance( + daft.from_pydict({"id": [1, 2, 3], "score": [10, 20, 30]}), + table_id=table_id, + mode="create", + **ns, + ).collect() + + source = daft_lance.read_lance( + table_id=table_id, + default_scan_options={"with_row_address": True}, + include_fragment_id=True, + **ns, + ).limit(2) + source = source.with_column("partial_score", source["score"] * 10) + daft_lance.merge_columns_df( + source.select("fragment_id", "_rowaddr", "partial_score"), + table_id=table_id, + **ns, + ) + + result = daft_lance.read_lance(table_id=table_id, **ns).sort("id").to_pydict() + assert result["partial_score"] == [100, 200, None] + + +def test_namespace_create_scalar_index(tmp_path: Path) -> None: + ns = _dir_ns(tmp_path) + table_id = ["indexed"] + + 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) + + +@pytest.mark.parametrize("segmented", [False, True]) +def test_namespace_create_distributed_inverted_index(tmp_path: Path, segmented: bool) -> None: + ns = _dir_ns(tmp_path) + table_id = [f"inverted_{segmented}"] + + daft_lance.write_lance( + daft.from_pydict({"id": list(range(20)), "text": [f"document {i}" for i in range(20)]}), + table_id=table_id, + mode="create", + **ns, + ).collect() + + daft_lance.create_scalar_index( + table_id=table_id, + column="text", + index_type="INVERTED", + name="text_idx", + segmented=segmented, + **ns, + ) + + import lance + import lance_namespace as ln + from lance_namespace import DescribeTableRequest + + location = ln.connect("dir", {"root": str(tmp_path)}).describe_table(DescribeTableRequest(id=table_id)).location + indices = lance.dataset(location).list_indices() + assert any(idx["name"] == "text_idx" and idx["fields"] == ["text"] for idx in indices) + + +def test_namespace_compact_files(tmp_path: Path) -> None: + ns = _dir_ns(tmp_path) + table_id = ["compacted"] + + 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() + + 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]} + + +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 = 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: Path) -> None: + ns = _dir_ns(tmp_path) + schema = daft.from_pydict({"id": [1]}).schema() + + with pytest.raises(ValueError, match="blob_columns"): + 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(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_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"] + + namespace = ln.connect("dir", {"root": str(tmp_path)}) + namespace.declare_table(DeclareTableRequest(id=table_id, location=None)) + + 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_namespace_create_declares_without_describe(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + requests: list[Any] = [] + + class RecordingNamespace: + def describe_table(self, request: Any) -> Any: + raise AssertionError("create must not describe before declaring") + + 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()) + + resolved = 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_namespace_overwrite_uses_plain_describe(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + requests: list[Any] = [] + + class RecordingNamespace: + def describe_table(self, request: Any) -> Any: + requests.append(request) + return SimpleNamespace(location=str(tmp_path / "t.lance"), storage_options=None) + + def declare_table(self, request: Any) -> Any: + raise AssertionError("an existing overwrite target must not be declared") + + 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" + ) + + 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_namespace_response_preserves_managed_versioning(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + response = SimpleNamespace( + location=str(tmp_path / "managed.lance"), + storage_options={"token": "temporary"}, + managed_versioning=True, + ) + + resolved = namespace_mod._resolved_from_response(response) + + assert resolved.managed_versioning is True + + +def test_namespace_commit_kwargs_include_managed_versioning(monkeypatch: pytest.MonkeyPatch) -> None: + namespace_client = object() + monkeypatch.setattr( + namespace_mod, + "get_namespace_kwargs", + lambda *args: {"namespace_client": namespace_client, "table_id": ["catalog", "table"]}, + ) + + kwargs = namespace_mod.get_namespace_commit_kwargs("rest", {}, ["catalog", "table"], True) + + assert kwargs == { + "namespace_client": namespace_client, + "table_id": ["catalog", "table"], + "namespace_client_managed_versioning": True, + } + + +def test_namespace_with_mem_wal_is_rejected(tmp_path: Path) -> None: + schema = daft.from_pydict({"id": [1]}).schema() + with pytest.raises(ValueError, match="use_mem_wal=True is not supported"): + LanceDataSink( + uri=None, + schema=schema, + mode="create", + table_id=["memwal"], + use_mem_wal=True, + **_dir_ns(tmp_path), + ) + + +def test_construct_lance_dataset_empty_storage_options_falls_back_to_io_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import daft_lance.utils as utils_mod + from daft.io import IOConfig, S3Config + + captured = {} + + class FakeDataset: + # construct_lance_dataset_handle pins this into open_kwargs["version"]. + version = 1 + + def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **kwargs: Any) -> Any: + captured["storage_options"] = storage_options + 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")) + handle = utils_mod.construct_lance_dataset_handle("s3://bucket/t.lance", storage_options={}, io_config=io_config) + + merged = captured["storage_options"] + assert isinstance(handle.dataset, FakeDataset) + assert merged is not None + 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 = 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_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_remain_explicit_for_uri(tmp_path: Path) -> None: + """The URI sink historically treats storage_options={} as explicitly empty.""" + from daft.io import IOConfig, S3Config + + io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) + schema = daft.from_pydict({"id": [1]}).schema() + + 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 == {} + + +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 + + captured = {} + + class FakeDataset: + # construct_lance_dataset_handle pins this into open_kwargs["version"]. + version = 1 + + def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **kwargs: Any) -> Any: + captured["uri"] = uri + captured["storage_options"] = storage_options + return FakeDataset() + + monkeypatch.setattr("daft_lance.utils.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"}, + managed_versioning=True, + ), + ) + + io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) + handle = utils_mod.construct_lance_dataset_handle( + 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 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 + assert merged["secret_access_key"] == "io-secret" # io_config fills the gaps + assert captured["uri"] is None # namespace addressing passes uri=None + assert handle.storage_options == merged + assert handle.uri == "s3://bucket/t.lance" + assert handle.managed_versioning is True + assert not hasattr(handle.dataset, "_lance_open_kwargs") + + +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 + + captured = {} + + class FakeDataset: + # construct_lance_dataset_handle pins this into open_kwargs["version"]. + version = 1 + + def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **kwargs: Any) -> Any: + captured["storage_options"] = storage_options + return FakeDataset() + + monkeypatch.setattr("daft_lance.utils.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_handle( + None, + io_config=io_config, + namespace_impl="rest", + namespace_properties={"uri": "http://namespace.example"}, + table_id=["t"], + ) + + 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: Path) -> None: + 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 = 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 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: Path) -> None: + 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: Path) -> None: + with pytest.raises(ValueError, match="'table_id' must be provided"): + daft_lance.read_lance( + namespace_impl="dir", + namespace_properties={"root": str(tmp_path)}, + ) + + +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"]) + + +def test_namespace_requires_uri_or_namespace() -> None: + with pytest.raises(ValueError, match="Must provide either 'uri' OR"): + daft_lance.read_lance() + + +# --------------------------------------------------------------------------- +# Distributed worker reopen (DatasetOpenContext) +# +# LanceDataset.__reduce__ drops _namespace_client / _table_id / +# _namespace_client_managed_versioning, so a pickled dataset reaches a worker +# stripped of its namespace identity. Maintenance paths therefore ship a +# DatasetOpenContext and reopen. These tests lock that contract down. +# --------------------------------------------------------------------------- + + +def _ns_handle(tmp_path: Path, table: str = "ctx_tbl"): + daft_lance.write_lance( + daft.from_pydict({"score": [1, 2, 3, 4]}), + table_id=[table], + mode="create", + **_dir_ns(tmp_path), + ).collect() + return construct_lance_dataset_handle(None, table_id=[table], **_dir_ns(tmp_path)) + + +def test_worker_open_context_is_free_of_live_objects(tmp_path: Path) -> None: + """The context must survive pickling without dragging unpicklable state along.""" + import pickle + + import lance + from lance_namespace import LanceNamespace + + context = _ns_handle(tmp_path).worker_open_context() + restored = pickle.loads(pickle.dumps(context)) + + for value in vars(restored).values(): + assert not isinstance(value, lance.LanceDataset) + assert not isinstance(value, LanceNamespace) + assert not hasattr(restored, "serialized_manifest") + + assert restored.table_id == ["ctx_tbl"] + assert restored.namespace_impl == "dir" + assert restored.version == context.version + + +def test_worker_open_restores_namespace_identity(tmp_path: Path) -> None: + """Reopening on a worker must rebuild the wiring pickle would have dropped.""" + context = _ns_handle(tmp_path, "identity_tbl").worker_open_context() + + worker_ds = context.open_pinned() + + assert worker_ds._namespace_client is not None + assert worker_ds._table_id == ["identity_tbl"] + assert worker_ds.version == context.version + assert worker_ds.count_rows() == 4 + + +def test_worker_open_propagates_managed_versioning(tmp_path: Path) -> None: + """A catalog-managed table must not be reopened as an unmanaged uri table.""" + import dataclasses + + context = dataclasses.replace( + _ns_handle(tmp_path, "managed_tbl").worker_open_context(), + managed_versioning=True, + ) + + assert context.commit_kwargs["namespace_client_managed_versioning"] is True + assert context.open_pinned()._namespace_client_managed_versioning is True + + +def test_worker_open_does_not_describe_the_table_location(tmp_path: Path) -> None: + """Workers must not put a namespace round-trip on every task. + + ``lance.dataset(None, namespace_client=..., table_id=...)`` resolves the + location with a describe_table on every call; the context passes the uri the + driver already resolved, so the worker open costs zero namespace calls. + """ + import lance + + metered = { + "namespace_impl": "dir", + "namespace_properties": {"root": str(tmp_path), "ops_metrics_enabled": "true"}, + } + daft_lance.write_lance( + daft.from_pydict({"score": [1, 2, 3, 4]}), table_id=["metered_tbl"], mode="create", **metered + ).collect() + handle = construct_lance_dataset_handle(None, table_id=["metered_tbl"], **metered) + context = handle.worker_open_context() + client = namespace_mod.get_or_create_namespace(metered["namespace_impl"], metered["namespace_properties"]) + + client.reset_ops_metrics() + worker_ds = context.open_pinned() + worker_ds.count_rows() + assert client.retrieve_ops_metrics().get("describe_table", 0) == 0 + + # The high-level entry point is what the low-level open exists to avoid. + client.reset_ops_metrics() + lance.dataset(None, namespace_client=client, table_id=["metered_tbl"]) + assert client.retrieve_ops_metrics().get("describe_table", 0) == 1 + + +def test_open_latest_sees_versions_written_after_pinning(tmp_path: Path) -> None: + """Workers stay on the planned snapshot; only coordinator steps move forward.""" + handle = _ns_handle(tmp_path, "versioned_tbl") + context = handle.worker_open_context() + + daft_lance.write_lance( + daft.from_pydict({"score": [5, 6]}), table_id=["versioned_tbl"], mode="append", **_dir_ns(tmp_path) + ).collect() + + assert context.open_pinned().version == context.version + assert context.open_pinned().count_rows() == 4 + assert context.open_latest().version > context.version + assert context.open_latest().count_rows() == 6 + + +def test_maintenance_udfs_hold_a_context_not_a_dataset(tmp_path: Path) -> None: + """No maintenance UDF may capture the driver's live dataset.""" + import pickle + + import lance + + from daft_lance.lance_compaction import CompactionTaskUDF + from daft_lance.lance_merge_column import ( + FastPathFragmentWriter, + FragmentHandler, + GroupFragmentMergeUDF, + ) + from daft_lance.lance_scalar_index import ( + FragmentIndexHandler, + SegmentedFragmentIndexHandler, + ) + + context = _ns_handle(tmp_path, "udf_tbl").worker_open_context() + + plain = [ + CompactionTaskUDF(context), + FragmentIndexHandler(context, "score", "BTREE", "idx", "uuid", False), + SegmentedFragmentIndexHandler(context, "score", "BTREE", "idx"), + ] + # daft.cls wraps these, so reach through to the instance it actually holds. + wrapped = [ + FragmentHandler(context, {"doubled": "score * 2"}, ["score"]), + GroupFragmentMergeUDF(context), + FastPathFragmentWriter(context, ["doubled"]), + ] + + instances = plain + [udf._daft_get_instance() for udf in wrapped] + for udf in instances: + state = vars(udf) + assert not any(isinstance(value, lance.LanceDataset) for value in state.values()), ( + f"{type(udf).__name__} captured a live LanceDataset" + ) + assert state["open_context"] is context + + for udf in plain: + assert isinstance(pickle.loads(pickle.dumps(udf)), type(udf)) + + +def test_worker_udf_opens_the_dataset_once_per_instance(tmp_path: Path) -> None: + """The pinned-manifest read is per UDF instance, never per call.""" + from daft_lance.lance_scalar_index import SegmentedFragmentIndexHandler + + context = _ns_handle(tmp_path, "reopen_tbl").worker_open_context() + opens = [] + + class CountingContext: + uri = context.uri + + def open_pinned(self): + opens.append(1) + return context.open_pinned() + + handler = SegmentedFragmentIndexHandler(CountingContext(), "score", "BTREE", "idx") + handler._dataset() + handler._dataset() + + assert len(opens) == 1 + + +def test_scan_open_kwargs_pin_the_planned_version(tmp_path: Path) -> None: + """Workers must read the snapshot the driver planned against, not latest. + + ``open_kwargs`` crosses to scan workers, so a ``version=None`` there means + each task independently opens latest. A compaction landing in between makes + planned fragment ids vanish; an overwrite silently substitutes the data. + """ + import lance + import pyarrow as pa + + ds_path = str(tmp_path / "pinned.lance") + lance.write_dataset(pa.table({"a": [1, 2, 3]}), ds_path) + handle = construct_lance_dataset_handle(ds_path) + planned_version = handle.dataset.version + + assert handle.open_kwargs["version"] == planned_version + + lance.write_dataset(pa.table({"a": [99]}), ds_path, mode="overwrite") + + worker_ds = namespace_mod.open_dataset_from_open_kwargs(handle.uri, handle.open_kwargs) + assert worker_ds.version == planned_version + assert worker_ds.to_table().to_pydict() == {"a": [1, 2, 3]} + + +def test_scan_open_kwargs_resolve_asof_into_a_numeric_version(tmp_path: Path) -> None: + """A tag/asof request is resolved on the driver; workers get the number.""" + import lance + import pyarrow as pa + + ds_path = str(tmp_path / "asof.lance") + lance.write_dataset(pa.table({"a": [1]}), ds_path) + lance.write_dataset(pa.table({"a": [2]}), ds_path, mode="append") + + handle = construct_lance_dataset_handle(ds_path, asof="2030-01-01 00:00:00") + + assert "asof" not in handle.open_kwargs + assert handle.open_kwargs["version"] == handle.dataset.version + + +def test_namespace_overwrite_recovers_from_a_lost_declare_race(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Overwrite must survive another writer declaring the table first.""" + from lance_namespace.errors import TableAlreadyExistsError, TableNotFoundError + + class RacingNamespace: + def __init__(self) -> None: + self.describes = 0 + self.declares = 0 + + def describe_table(self, request: Any) -> Any: + self.describes += 1 + # First describe loses the race; by the second, the rival's table exists. + if self.describes == 1: + raise TableNotFoundError("not found") + return SimpleNamespace(location=str(tmp_path / "rival.lance"), storage_options=None) + + def declare_table(self, request: Any) -> Any: + self.declares += 1 + raise TableAlreadyExistsError("declared by another writer") + + namespace = RacingNamespace() + monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: namespace) + + resolved = namespace_mod.resolve_namespace_table( + namespace_impl="rest", + namespace_properties={"uri": "http://namespace.example"}, + table_id=["raced"], + mode="overwrite", + ) + + assert resolved.uri == str(tmp_path / "rival.lance") + assert (namespace.describes, namespace.declares) == (2, 1) + + +def test_namespace_create_still_fails_on_a_lost_declare_race(monkeypatch: pytest.MonkeyPatch) -> None: + """For create, losing the race is the answer -- it must not be swallowed.""" + from lance_namespace.errors import TableAlreadyExistsError + + class RacingNamespace: + def describe_table(self, request: Any) -> Any: + raise AssertionError("create must not describe") + + def declare_table(self, request: Any) -> Any: + raise TableAlreadyExistsError("declared by another writer") + + monkeypatch.setattr(namespace_mod, "get_or_create_namespace", lambda *args: RacingNamespace()) + + with pytest.raises(TableAlreadyExistsError): + namespace_mod.resolve_namespace_table( + namespace_impl="rest", + namespace_properties={"uri": "http://namespace.example"}, + table_id=["raced"], + mode="create", + ) 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..8dafb6a --- /dev/null +++ b/tests/io/lancedb/test_namespace_rest_e2e.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import os +import uuid +from typing import Any + +import pytest + +import daft +import daft_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(monkeypatch: pytest.MonkeyPatch) -> None: + import lance + import lance_namespace as ln + from lance_namespace import CreateNamespaceRequest, DescribeTableRequest, NamespaceExistsRequest + from lance_namespace.errors import TableAlreadyExistsError + + namespace_properties = {"uri": os.environ["DAFT_LANCE_REST_URI"]} + + # Ensure object-store access in this test cannot silently succeed through + # ambient AWS credentials instead of credentials vended by the catalog. + for name in ( + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_ROLE_ARN", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + "AWS_CONTAINER_AUTHORIZATION_TOKEN", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", os.devnull) + monkeypatch.setenv("AWS_CONFIG_FILE", os.devnull) + catalog = os.environ.get("DAFT_LANCE_REST_CATALOG", "lance_catalog") + schema = os.environ.get("DAFT_LANCE_REST_SCHEMA", "daft_ns_e2e") + table_id = [catalog, schema, f"orders_{uuid.uuid4().hex[:8]}"] + ns: dict[str, Any] = {"namespace_impl": "rest", "namespace_properties": namespace_properties} + + 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_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_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, vend_credentials=True)) + location = getattr(describe, "location", None) or getattr(describe, "table_uri", None) + assert location + assert describe.storage_options, "catalog must vend storage options for this integration test" + + result = daft_lance.read_lance(table_id=table_id, **ns).to_pydict() + assert result == { + "id": [1, 2, 3, 4, 5], + "label": ["a", "b", "c", "d", "e"], + "score": [10, 20, 30, 40, 50], + } + + 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 + 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], + }