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..4e2d321 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -2,22 +2,21 @@ 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.io._checkpoint import attach_checkpoint -from daft.io.object_store_options import io_config_to_storage_options from daft.logical.builder import LogicalPlanBuilder from .lance_compaction import compact_files_internal +from .lance_data_sink import LanceDataSink from .lance_merge_column import merge_columns_from_df, merge_columns_internal from .lance_scalar_index import create_scalar_index_internal from .lance_scan import LanceDBScanOperator +from .namespace import validate_uri_or_namespace from .utils import construct_lance_dataset if TYPE_CHECKING: @@ -28,20 +27,35 @@ from daft.dependencies import pa +def _effective_uri(lance_ds: LanceDataset, uri: str | pathlib.Path | None) -> str: + """The dataset location: the user-supplied URI, or the namespace-resolved one.""" + return str(uri) if uri is not None else str(lance_ds.uri) + + +def _effective_storage_options(lance_ds: LanceDataset, storage_options: dict[str, Any] | None) -> dict[str, Any] | None: + """Storage options actually used to open the dataset (includes namespace-vended ones).""" + open_kwargs = getattr(lance_ds, "_lance_open_kwargs", None) or {} + return open_kwargs.get("storage_options") or storage_options + + @PublicAPI def read_lance( - uri: str | pathlib.Path, + 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 +63,11 @@ def read_lance( uri: The URI of the Lance table to read from. Accepts a local path or an object-store URI like "s3://bucket/path". io_config: A custom IOConfig to use when accessing LanceDB data. Defaults to None. + table_id: Table identifier within the namespace, e.g. ["catalog", "schema", "table"]. + Mutually exclusive with ``uri``. + namespace_impl: Lance Namespace implementation, e.g. "dir" or "rest". + namespace_properties: Properties for connecting to the namespace, e.g. + {"root": "/data"} for "dir" or {"uri": "http://host:port"} for "rest". version : optional, int | str If specified, load a specific version of the Lance dataset. Else, loads the latest version. A version number (`int`) or a tag (`str`) can be provided. @@ -123,21 +142,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( 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, @@ -161,9 +184,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 +212,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 +253,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( 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, @@ -239,11 +273,11 @@ def merge_columns( return merge_columns_internal( lance_ds, - uri, + _effective_uri(lance_ds, uri), transform=transform, read_columns=read_columns, reader_schema=reader_schema, - storage_options=storage_options, + storage_options=_effective_storage_options(lance_ds, storage_options), daft_remote_args=daft_remote_args, concurrency=concurrency, ) @@ -252,9 +286,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 +317,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 +360,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( 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, @@ -342,10 +387,10 @@ def merge_columns_df( return merge_columns_from_df( df, lance_ds=lance_ds, - uri=uri, + uri=_effective_uri(lance_ds, uri), read_columns=read_columns, reader_schema=reader_schema, - storage_options=storage_options, + storage_options=_effective_storage_options(lance_ds, storage_options), daft_remote_args=daft_remote_args, concurrency=concurrency, left_on=left_on, @@ -356,9 +401,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 +434,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 +517,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( 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, @@ -480,12 +536,12 @@ def create_scalar_index( create_scalar_index_internal( lance_ds=lance_ds, - uri=uri, + uri=_effective_uri(lance_ds, uri), column=column, index_type=index_type, name=name, replace=replace, - storage_options=storage_options, + storage_options=_effective_storage_options(lance_ds, storage_options), fragment_group_size=fragment_group_size, num_partitions=num_partitions, max_concurrency=max_concurrency, @@ -496,9 +552,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 +578,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 +613,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( + lance_ds = construct_lance_dataset( uri, storage_options=storage_options, + io_config=io_config, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, version=version, asof=asof, block_size=block_size, @@ -571,3 +636,58 @@ def compact_files( 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: Any | 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). + + 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_data_sink.py b/daft_lance/lance_data_sink.py index 62c36bb..bda33d1 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,14 @@ detect_blob_v2_columns, resolve_storage_version, ) +from daft_lance.namespace import ( + ResolvedNamespaceTable, + 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 +48,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 +68,17 @@ def __init__( compact_after_write: bool = True, ) -> None: self._reject_unsupported_modes(mode, use_legacy_format) - if not isinstance(uri, (str, pathlib.Path)): + validate_uri_or_namespace(uri, namespace_impl, table_id, 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 +88,46 @@ def __init__( ) self._pyarrow_schema = self._normalize_schema(schema) self._init_blob_policy(blob_columns) + self._requested_storage_version = self._blob.apply_blob_v2_default(data_storage_version) self._use_mem_wal = use_mem_wal self._compact_after_write = compact_after_write self._mem_wal_total_rows: int = 0 self._mem_wal_total_bytes: int = 0 + # Resolved by start() on the driver, before the sink is serialized to workers. + # Construction must stay side-effect free: no namespace calls, no dataset opens. + self._table_uri: str | None = None + self._storage_options: dict[str, str] | None = None + self._data_storage_version: LanceStorageVersion | None = None + self._effective_pyarrow_schema: pa.Schema | None = None self._version: int = 0 self._table_schema: pa.Schema | None = None + + self._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) + 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 +139,41 @@ 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 _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, non-empty user-provided options replace the + io_config-derived ones entirely (falsy fallthrough, matching the read + entry points via construct_lance_dataset). + """ + io_derived = io_config_to_storage_options(self._io_config, resolved.uri) + if self._uri is not None: + base = self._user_storage_options or io_derived + return merge_storage_options(base, resolved.storage_options) + return merge_storage_options(io_derived, self._user_storage_options, resolved.storage_options) @staticmethod def _reject_unsupported_modes( @@ -167,7 +233,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 +247,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 +290,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 +303,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 +313,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 +325,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 +400,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_kwargs, ) stats = dataset.stats.dataset_stats() stats_dict = MicroPartition.from_pydict( @@ -348,7 +429,7 @@ def _finalize_cow(self, write_results: list[WriteResult[list[FragmentMetadata]]] return stats_dict def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadata]]]) -> MicroPartition: - dataset = lance.dataset(self._table_uri, storage_options=self._storage_options) + dataset = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) if self._compact_after_write: logger.info( @@ -359,7 +440,9 @@ def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadat from daft_lance.lance_compaction import compact_files_internal compact_files_internal(dataset) - dataset = lance.dataset(self._table_uri, storage_options=self._storage_options) + dataset = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + ) stats = dataset.stats.dataset_stats() return MicroPartition.from_pydict( diff --git a/daft_lance/lance_merge_column.py b/daft_lance/lance_merge_column.py index 8984120..df2d682 100644 --- a/daft_lance/lance_merge_column.py +++ b/daft_lance/lance_merge_column.py @@ -14,6 +14,8 @@ from daft.udf import cls as daft_cls from daft.udf import method +from .namespace import namespace_kwargs_for_dataset + if TYPE_CHECKING: import pathlib from collections.abc import Callable @@ -93,6 +95,7 @@ def merge_columns_internal( op, read_version=lance_ds.version, storage_options=storage_options, + **namespace_kwargs_for_dataset(lance_ds), ) @@ -473,6 +476,7 @@ def _merge_fast_path( op, read_version=lance_ds.version, storage_options=storage_options, + **namespace_kwargs_for_dataset(lance_ds), ) @@ -521,4 +525,5 @@ def _merge_slow_path( op, read_version=lance_ds.version, storage_options=storage_options, + **namespace_kwargs_for_dataset(lance_ds), ) diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index bd8e60c..2cd2fa2 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -14,6 +14,7 @@ import lance from daft.dependencies import pa +from daft_lance.namespace import namespace_kwargs_for_dataset from daft_lance.utils import distribute_fragments_balanced logger = logging.getLogger(__name__) @@ -386,7 +387,8 @@ def _create_segmented_index( # Reload dataset to pick up the latest version (segment files were written # by workers against the version that was current at their invocation time). - lance_ds = lance.LanceDataset(uri, storage_options=storage_options) + namespace_kwargs = namespace_kwargs_for_dataset(lance_ds) + lance_ds = lance.dataset(None if namespace_kwargs else uri, storage_options=storage_options, **namespace_kwargs) index_metas = _prepare_index_segments_for_commit(lance_ds, index_type, index_metas) logger.info( @@ -460,7 +462,8 @@ def _create_partitioned_index( df.collect() logger.info("Starting index metadata merging by reloading dataset to get latest state") - lance_ds = lance.LanceDataset(uri, storage_options=storage_options) + namespace_kwargs = namespace_kwargs_for_dataset(lance_ds) + lance_ds = lance.dataset(None if namespace_kwargs else uri, storage_options=storage_options, **namespace_kwargs) lance_ds.merge_index_metadata(index_id, index_type) logger.info("Starting atomic index creation and commit") @@ -506,6 +509,7 @@ def _create_partitioned_index( create_index_op, read_version=lance_ds.version, storage_options=storage_options, + **namespace_kwargs, ) logger.info("Index %s created successfully with ID %s", name, index_id) diff --git a/daft_lance/lance_scan.py b/daft_lance/lance_scan.py index e0ea8fb..e350479 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) diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py new file mode 100644 index 0000000..fd6b2f4 --- /dev/null +++ b/daft_lance/namespace.py @@ -0,0 +1,243 @@ +"""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. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from functools import lru_cache +from typing import Any +from urllib.parse import urlparse + +import lance + +_NAMESPACE_CACHE_SIZE = int(os.environ.get("DAFT_LANCE_NAMESPACE_CACHE_SIZE", "16")) + + +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. + """ + parsed = urlparse(location) + if parsed.scheme == "file": + return parsed.path + return location + + +@lru_cache(maxsize=_NAMESPACE_CACHE_SIZE) +def _get_cached_namespace(namespace_impl: str, namespace_properties_tuple: tuple[tuple[str, str], ...] | None) -> Any: + import lance_namespace as ln + + namespace_properties = dict(namespace_properties_tuple) if namespace_properties_tuple else {} + return ln.connect(namespace_impl, namespace_properties) + + +def get_or_create_namespace(namespace_impl: str | None, namespace_properties: dict[str, str] | None) -> Any | None: + """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} + + +# `lance.fragment.write_fragments` accepts the same namespace kwargs as `lance.dataset`. +get_write_fragments_kwargs = get_namespace_kwargs + + +def _storage_options(response: Any) -> dict[str, str] | None: + storage_options = getattr(response, "storage_options", None) + if storage_options is None: + return None + return dict(storage_options) + + +def _response_location(response: Any) -> str: + location = getattr(response, "location", None) or getattr(response, "table_uri", None) + if not location: + raise ValueError("Namespace response did not include a table location.") + return _normalize_file_uri(str(location)) + + +@dataclass(frozen=True) +class ResolvedNamespaceTable: + """A namespace table resolved to a physical location and storage options.""" + + uri: str + storage_options: dict[str, str] | None = None + + +def _resolved_from_response(response: Any) -> ResolvedNamespaceTable: + return ResolvedNamespaceTable( + uri=_response_location(response), + storage_options=_storage_options(response), + ) + + +def _describe_table(namespace: Any, table_id: list[str]) -> Any: + from lance_namespace import DescribeTableRequest + + # vend_credentials is explicit: when unset, whether the namespace returns + # storage credentials is implementation-defined. + return namespace.describe_table(DescribeTableRequest(id=table_id, vend_credentials=True)) + + +def _declare_table(namespace: Any, table_id: list[str]) -> ResolvedNamespaceTable: + from lance_namespace import DeclareTableRequest + + response = namespace.declare_table(DeclareTableRequest(id=table_id, location=None, vend_credentials=True)) + return _resolved_from_response(response) + + +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; + 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 TableNotFoundError + + try: + return _resolved_from_response(_describe_table(namespace, table_id)) + except TableNotFoundError: + if mode == "overwrite": + return _declare_table(namespace, table_id) + raise + + +def merge_storage_options(*layers: dict[str, Any] | None) -> dict[str, Any] | None: + """Merge storage-option layers; later layers take precedence. + + Callers order layers from lowest to highest priority, conventionally: + io_config-derived < user-provided ``storage_options`` < namespace-vended. + """ + merged: dict[str, Any] = {} + 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 a ``_lance_open_kwargs`` dict.""" + return ( + open_kwargs.pop("namespace_impl", None), + open_kwargs.pop("namespace_properties", None), + open_kwargs.pop("table_id", None), + ) + + +def namespace_kwargs_for_dataset(ds: Any) -> dict[str, Any]: + """Namespace kwargs for commit-style calls, recovered from ``ds._lance_open_kwargs``. + + ``lance.LanceDataset.commit`` is a static method, so a dataset opened through a + namespace does not carry its client into commits; re-derive it from the open + kwargs stashed by ``construct_lance_dataset``. + """ + open_kwargs = getattr(ds, "_lance_open_kwargs", None) or {} + return get_namespace_kwargs( + open_kwargs.get("namespace_impl"), + open_kwargs.get("namespace_properties"), + open_kwargs.get("table_id"), + ) + + +def open_dataset_from_open_kwargs(ds_uri: str | None, open_kwargs: dict[str, Any] | None) -> lance.LanceDataset: + """Re-open a dataset on a worker from serialized ``_lance_open_kwargs``. + + 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..fb4cebc 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -6,11 +6,21 @@ import lance from daft.dependencies import pa +from daft.io.object_store_options import io_config_to_storage_options from daft.logical.schema import Schema as DaftSchema +from daft_lance.namespace import ( + get_namespace_kwargs, + 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__) @@ -82,12 +92,47 @@ def distribute_fragments_balanced(fragments: list[Any], fragment_group_size: int def construct_lance_dataset( - uri: str | pathlib.Path, + uri: str | pathlib.Path | None, version: int | str | None = None, storage_options: dict[str, Any] | None = None, + io_config: IOConfig | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, + table_id: list[str] | None = None, **kwargs: Any, ) -> lance.LanceDataset: - """Construct a Lance dataset with common options.""" + """Construct a Lance dataset with common options. + + Storage options are layered from lowest to highest priority: + io_config-derived < user-provided ``storage_options`` < namespace-vended. + For a plain ``uri``, user-provided ``storage_options`` replace the + io_config-derived ones entirely (historical behavior). + """ + validate_uri_or_namespace(uri, namespace_impl, table_id, namespace_properties) + resolved_uri = str(uri) if uri is not None else None + namespace_storage_options = None + 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 + 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,11 +143,21 @@ def construct_lance_dataset( # Non-dict defaults are forwarded as-is. kwargs["default_scan_options"] = original_default_scan_options - ds = lance.dataset(uri, storage_options=storage_options, version=version, **kwargs) + dataset_uri = None if has_namespace_params(namespace_impl, table_id) else resolved_uri + ds = lance.dataset( + dataset_uri, + storage_options=merged_storage_options, + version=version, + **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), + **kwargs, + ) effective_kwargs = { - "storage_options": storage_options, + "storage_options": merged_storage_options, "version": version, + "namespace_impl": namespace_impl, + "namespace_properties": namespace_properties, + "table_id": table_id, } effective_kwargs.update(kwargs or {}) try: diff --git a/tests/io/lancedb/test_lance_batching.py b/tests/io/lancedb/test_lance_batching.py index ea27359..c805115 100644 --- a/tests/io/lancedb/test_lance_batching.py +++ b/tests/io/lancedb/test_lance_batching.py @@ -55,6 +55,7 @@ def test_accumulate_small_micropartitions(schema, tmp_path): with patch("daft_lance.lance_data_sink.lance", fake): sink = LanceDataSink(uri=str(tmp_path / "tbl"), schema=schema, mode="create", max_rows_per_file=25) + sink.start() mps = [_make_mp(10), _make_mp(20), _make_mp(30)] results = list(sink.write(iter(mps))) @@ -69,6 +70,7 @@ def test_flush_remaining_at_end(schema, tmp_path): with patch("daft_lance.lance_data_sink.lance", fake): sink = LanceDataSink(uri=str(tmp_path / "tbl"), schema=schema, mode="create", max_rows_per_file=25) + sink.start() mps = [_make_mp(10), _make_mp(5)] results = list(sink.write(iter(mps))) @@ -83,6 +85,7 @@ def test_large_micropartition_writes_directly(schema, tmp_path): with patch("daft_lance.lance_data_sink.lance", fake): sink = LanceDataSink(uri=str(tmp_path / "tbl"), schema=schema, mode="create", max_rows_per_file=25) + sink.start() mps = [_make_mp(30)] results = list(sink.write(iter(mps))) @@ -102,6 +105,7 @@ def test_accumulation_default_batches_across_micropartitions(schema, tmp_path): with patch("daft_lance.lance_data_sink.lance", fake): sink = LanceDataSink(uri=str(tmp_path / "tbl"), schema=schema, mode="create") + sink.start() mps = [_make_mp(10), _make_mp(10), _make_mp(10)] results = list(sink.write(iter(mps))) diff --git a/tests/io/lancedb/test_lance_data_sink_internals.py b/tests/io/lancedb/test_lance_data_sink_internals.py index 47f3107..e6983c6 100644 --- a/tests/io/lancedb/test_lance_data_sink_internals.py +++ b/tests/io/lancedb/test_lance_data_sink_internals.py @@ -17,14 +17,16 @@ def test_load_existing_dataset_missing_raises_for_append(tmp_path): schema = pa.schema([("a", pa.int64())]) + sink = LanceDataSink(uri=str(tmp_path / f"missing-{uuid.uuid4()}"), schema=schema, mode="append") with pytest.raises(ValueError, match="Cannot append to non-existent Lance dataset"): - LanceDataSink(uri=str(tmp_path / f"missing-{uuid.uuid4()}"), schema=schema, mode="append") + sink.start() def test_load_existing_dataset_missing_returns_none_for_create(tmp_path): schema = pa.schema([("a", pa.int64())]) sink = LanceDataSink(uri=str(tmp_path / f"missing-{uuid.uuid4()}"), schema=schema, mode="create") - # Construction succeeds; the dataset will be created on the first write. + sink.start() + # Start succeeds; the dataset will be created on the first write. assert sink is not None @@ -143,6 +145,7 @@ def __getattr__(self, name): with patch("daft_lance.lance_data_sink.lance", fake): sink = LanceDataSink(uri=str(tmp_path / "tbl"), schema=schema, mode="create") + sink.start() write_results = list(sink.write(iter(mps))) # Default ``max_rows_per_file`` is 1_048_576. 20 partitions x 100K = 2M rows @@ -161,6 +164,7 @@ def test_prepare_arrow_table_missing_column_rejected(tmp_path): lance.write_dataset(initial, uri) # Try to append with only one column sink = LanceDataSink(uri=uri, schema=schema, mode="append") + sink.start() missing_col = pa.table({"a": pa.array([2], type=pa.int64())}) with pytest.raises(Exception): # any error is fine; the redesign currently relies on the cast to fail sink._prepare_arrow_table(missing_col) diff --git a/tests/io/lancedb/test_lance_data_sink_storage_versions.py b/tests/io/lancedb/test_lance_data_sink_storage_versions.py index 28496bc..fcf5f91 100644 --- a/tests/io/lancedb/test_lance_data_sink_storage_versions.py +++ b/tests/io/lancedb/test_lance_data_sink_storage_versions.py @@ -47,6 +47,7 @@ def test_append_inherits_storage_version(tmp_path): # inherit from the existing dataset. schema = pa.schema([("a", pa.int64())]) sink = LanceDataSink(uri=uri, schema=schema, mode="append") + sink.start() assert sink._data_storage_version == initial_version df2 = _make_df(5) @@ -64,8 +65,9 @@ def test_storage_version_mismatch_on_append_rejected(tmp_path): df1.write_lance(uri, mode="create", data_storage_version="2.1") schema = pa.schema([("a", pa.int64())]) + sink = LanceDataSink(uri=uri, schema=schema, mode="append", data_storage_version="2.2") with pytest.raises(ValueError, match="does not match existing dataset version"): - LanceDataSink(uri=uri, schema=schema, mode="append", data_storage_version="2.2") + sink.start() def test_use_legacy_format_emits_deprecation_warning(tmp_path): diff --git a/tests/io/lancedb/test_mem_wal_writes.py b/tests/io/lancedb/test_mem_wal_writes.py index dd51b80..5ac5715 100644 --- a/tests/io/lancedb/test_mem_wal_writes.py +++ b/tests/io/lancedb/test_mem_wal_writes.py @@ -205,6 +205,7 @@ def test_ensure_creates_dataset_if_missing(self, lance_dataset_path): target = os.path.join(lance_dataset_path, "new_ds") schema = pa.schema([("a", pa.int64())]) sink = LanceDataSink(uri=target, schema=schema, mode="create", use_mem_wal=True) + sink.start() ds = sink._ensure_mem_wal_dataset() assert ds is not None assert ds.mem_wal_index_details() is not None @@ -214,6 +215,7 @@ def test_ensure_reuses_existing_dataset(self, lance_dataset_path): lance.write_dataset(pa.table({"a": [1]}, schema=schema), lance_dataset_path) sink = LanceDataSink(uri=lance_dataset_path, schema=schema, mode="append", use_mem_wal=True) + sink.start() ds = sink._ensure_mem_wal_dataset() assert ds is not None assert ds.mem_wal_index_details() is not None @@ -226,6 +228,7 @@ def test_ensure_idempotent_initialization(self, lance_dataset_path): ds.initialize_mem_wal(unsharded=True) sink = LanceDataSink(uri=lance_dataset_path, schema=schema, mode="append", use_mem_wal=True) + sink.start() ds2 = sink._ensure_mem_wal_dataset() assert ds2.mem_wal_index_details() is not None @@ -236,6 +239,7 @@ def test_write_result_has_empty_fragments(self, lance_dataset_path): schema = pa.schema([("a", pa.int64())]) sink = LanceDataSink(uri=lance_dataset_path, schema=schema, mode="create", use_mem_wal=True) + sink.start() mp = MicroPartition.from_pydict({"a": [1, 2, 3]}) results = list(sink.write(iter([mp]))) assert len(results) == 1 @@ -254,6 +258,7 @@ def test_finalize_mem_wal_returns_stats(self, lance_dataset_path): use_mem_wal=True, compact_after_write=False, ) + sink.start() mp = MicroPartition.from_pydict({"a": [1, 2, 3]}) results = list(sink.write(iter([mp]))) stats_mp = sink.finalize(results) diff --git a/tests/io/lancedb/test_namespace.py b/tests/io/lancedb/test_namespace.py new file mode 100644 index 0000000..11c0109 --- /dev/null +++ b/tests/io/lancedb/test_namespace.py @@ -0,0 +1,506 @@ +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 + + +def _dir_ns(tmp_path: Path) -> dict[str, Any]: + return {"namespace_impl": "dir", "namespace_properties": {"root": str(tmp_path)}} + + +def test_namespace_write_read_append_roundtrip(tmp_path: 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_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) + + +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_construct_lance_dataset_empty_storage_options_falls_back_to_io_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import daft_lance.utils as utils_mod + from daft.io import IOConfig, S3Config + + captured = {} + + class FakeDataset: + pass + + def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **kwargs: Any) -> Any: + captured["storage_options"] = storage_options + return FakeDataset() + + monkeypatch.setattr("daft_lance.utils.lance.dataset", fake_dataset) + + io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) + utils_mod.construct_lance_dataset("s3://bucket/t.lance", storage_options={}, io_config=io_config) + + merged = captured["storage_options"] + assert merged is not None + assert merged["access_key_id"] == "io-key" + + +def test_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_falls_back_to_io_config(tmp_path: Path) -> None: + """storage_options={} must behave like the read entry points: fall through to io_config.""" + from daft.io import IOConfig, S3Config + + io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) + schema = daft.from_pydict({"id": [1]}).schema() + + sink = LanceDataSink("s3://bucket/t.lance", schema, "create", io_config, storage_options={}) + merged = sink._merged_storage_options(namespace_mod.ResolvedNamespaceTable(uri="s3://bucket/t.lance")) + assert merged is not None + assert merged["access_key_id"] == "io-key" + + +def test_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: + pass + + 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"}, + ), + ) + + io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) + utils_mod.construct_lance_dataset( + None, + io_config=io_config, + storage_options={"access_key_id": "user-key", "user_option": "kept"}, + namespace_impl="rest", + namespace_properties={"uri": "http://namespace.example"}, + table_id=["t"], + ) + + merged = captured["storage_options"] + assert merged 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 + + +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: + pass + + def fake_dataset(uri: Any, storage_options: Any = None, version: Any = None, **kwargs: Any) -> Any: + captured["storage_options"] = storage_options + return FakeDataset() + + monkeypatch.setattr("daft_lance.utils.lance.dataset", fake_dataset) + monkeypatch.setattr(utils_mod, "get_namespace_kwargs", lambda *args: {}) + monkeypatch.setattr( + utils_mod, + "resolve_namespace_table", + lambda **kwargs: namespace_mod.ResolvedNamespaceTable(uri="s3://bucket/t.lance", storage_options=None), + ) + + io_config = IOConfig(s3=S3Config(key_id="io-key", access_key="io-secret", region_name="us-east-1")) + utils_mod.construct_lance_dataset( + None, + io_config=io_config, + namespace_impl="rest", + namespace_properties={"uri": "http://namespace.example"}, + table_id=["t"], + ) + + merged = captured["storage_options"] + assert merged 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() 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..2445cc3 --- /dev/null +++ b/tests/io/lancedb/test_namespace_rest_e2e.py @@ -0,0 +1,100 @@ +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() -> None: + import lance + import lance_namespace as ln + from lance_namespace import CreateNamespaceRequest, DescribeTableRequest, NamespaceExistsRequest + from lance_namespace.errors import TableAlreadyExistsError + + namespace_properties = {"uri": os.environ["DAFT_LANCE_REST_URI"]} + catalog = os.environ.get("DAFT_LANCE_REST_CATALOG", "lance_catalog") + 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)) + location = getattr(describe, "location", None) or getattr(describe, "table_uri", None) + assert location + + 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], + }