diff --git a/README.md b/README.md index 106d6bd..79b3ec8 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,35 @@ from daft_lance import merge_columns_df merge_columns_df(df, "s3://bucket/my_dataset") ``` +### Conditional Overwrite + +Replace just the rows matching a predicate. One Lance commit deletes them from the existing +table and adds the new data, so readers see either the whole replacement or none of it. + +```python +import daft_lance + +daft_lance.write_lance( + df, + "s3://bucket/events", + mode="overwrite_where", + predicate="dt = DATE '2026-08-25'", +).collect() +``` + +The table must already exist, and every input row must satisfy `predicate` — pass +`validate_predicate=False` to append rows outside it anyway (which makes re-running the same +write duplicate them instead of replacing them). That check evaluates the predicate with Daft, +so with it on the filter has to mean the same thing to both engines; a predicate Daft types +differently (a bare `TIMESTAMP` literal against a naive timestamp column) or evaluates +differently (a decimal literal against a `float32` column) is rejected up front, before any +data is written, and needs `validate_predicate=False`. + +> **Warning:** Lance does not treat a concurrent append or update as conflicting with this +> commit, so rows another writer adds while the overwrite runs survive it even when they match +> `predicate`, without any error. Make sure no other writer touches the table during a +> conditional overwrite. + ### Namespace Tables Address Lance tables through a [Lance Namespace](https://lancedb.github.io/lance-namespace/) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 0ad2231..c668f18 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -2,7 +2,7 @@ import pathlib from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any from daft import context from daft.api_annotations import PublicAPI @@ -14,7 +14,7 @@ from daft.schema import Schema from .lance_compaction import compact_files_internal -from .lance_data_sink import LanceDataSink +from .lance_data_sink import LanceDataSink, LanceWriteMode 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 @@ -632,10 +632,12 @@ def compact_files( def write_lance( df: DataFrame, uri: str | pathlib.Path | None = None, - mode: Literal["create", "append", "overwrite"] = "create", + mode: LanceWriteMode = "create", io_config: IOConfig | None = None, schema: Schema | pa.Schema | None = None, *, + predicate: str | None = None, + validate_predicate: bool = True, table_id: list[str] | None = None, namespace_impl: str | None = None, namespace_properties: dict[str, str] | None = None, @@ -646,9 +648,25 @@ def write_lance( Args: df: The DataFrame to write. uri: The URI of the Lance table. Mutually exclusive with the namespace parameters. - mode: One of "create", "append", or "overwrite". + mode: One of "create", "append", "overwrite", or "overwrite_where". + ``"overwrite_where"`` replaces just the rows matching ``predicate``: one Lance + commit deletes them from the existing table and adds this DataFrame's data, so + readers see either the whole replacement or none of it. It requires an existing + table and is not supported with ``use_mem_wal=True``. io_config: A custom IOConfig to use when accessing Lance data. schema: Desired schema to enforce during write; defaults to the DataFrame schema. + predicate: SQL predicate selecting the rows to replace. Required by, and only valid + with, ``mode="overwrite_where"``. Uses Lance's SQL filter dialect, e.g. + ``"dt = DATE '2026-08-25'"``. + validate_predicate: For ``mode="overwrite_where"``, check that every input row + satisfies ``predicate`` and fail the write otherwise (default True). Rows outside + the predicate are still appended when this is False, which makes re-running the + same write duplicate them instead of replacing them. The check evaluates + ``predicate`` with Daft, so leaving it on also requires Daft to read the filter the + same way Lance does; the write fails up front, before any data is written, when it + cannot (a bare ``TIMESTAMP`` literal against a naive timestamp column) or when the + two engines disagree (a decimal literal compared against a float32 column). Pass + False in those cases. 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. @@ -664,12 +682,25 @@ def write_lance( plan is constructed. This includes missing/duplicate targets, append schema compatibility, and storage-version conflicts. + Warning: + ``mode="overwrite_where"`` commits against the table version the write started + from, but Lance does not treat a concurrent append or update as conflicting with + it. Rows another writer adds during the overwrite therefore survive it, even when + they match ``predicate``, and the commit still succeeds. Make sure no other writer + touches the table while a conditional overwrite is running. + 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 + + Replace one day's rows and add this batch in a single commit: + + >>> daft_lance.write_lance( + ... df, "/tmp/events", mode="overwrite_where", predicate="dt = DATE '2026-08-25'" + ... ).collect() # doctest: +SKIP """ validate_uri_or_namespace(uri, namespace_impl, table_id, namespace_properties) @@ -681,6 +712,8 @@ def write_lance( schema, mode, io_config, + predicate=predicate, + validate_predicate=validate_predicate, table_id=table_id, namespace_impl=namespace_impl, namespace_properties=namespace_properties, diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index 3feef7e..7303577 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -10,9 +10,11 @@ import lance from lance.fragment import FragmentMetadata +import daft from daft.context import get_context from daft.datatype import DataType from daft.dependencies import pa +from daft.expressions import ExpressionsProjection from daft.io import DataSink from daft.io.object_store_options import io_config_to_storage_options from daft.io.sink import WriteResult @@ -40,9 +42,71 @@ from collections.abc import Iterator from daft.daft import IOConfig + from daft.expressions import Expression logger = logging.getLogger(__name__) +# What the caller asks for, and what the write physically does. ``overwrite_where`` +# writes exactly like an append -- it only differs at commit time -- so it is +# normalized to "append" once in the constructor. Every mode check outside the +# commit path reads the normalized value, because a check that forgets the new +# mode fails silently (see resolve_storage_version, which would skip the +# storage-version compatibility check entirely). +LanceWriteMode = Literal["create", "append", "overwrite", "overwrite_where"] +LancePhysicalWriteMode = Literal["create", "append", "overwrite"] + + +def _dataset_stats(dataset: lance.LanceDataset) -> MicroPartition: + """The single-row write result: dataset stats plus the version just produced.""" + stats = dataset.stats.dataset_stats() + return MicroPartition.from_pydict( + { + "num_fragments": pa.array([stats["num_fragments"]], type=pa.int64()), + "num_deleted_rows": pa.array([stats["num_deleted_rows"]], type=pa.int64()), + "num_small_files": pa.array([stats["num_small_files"]], type=pa.int64()), + "version": pa.array([dataset.version], type=pa.int64()), + } + ) + + +def _compile_predicate(predicate: str) -> Expression: + """Compile the Lance filter into the Daft expression the input check runs. + + The check deliberately runs a second SQL engine over the input, which is only + trustworthy where the two agree; ``_reject_untrusted_predicate`` rules out the + case where they do not. + """ + try: + return daft.sql_expr(predicate) + except Exception as e: + raise ValueError( + f"predicate={predicate!r} could not be parsed by Daft, so input rows cannot be " + "checked against it. Daft's SQL dialect does not cover every Lance filter; pass " + "validate_predicate=False to skip the check and write the input as-is." + ) from e + + +def _evaluates_against(expr: Expression, schema: pa.Schema) -> bool: + """Whether ``expr`` resolves and type-checks against a zero-row input.""" + try: + MicroPartition.from_arrow(schema.empty_table()).filter(ExpressionsProjection([expr])) + except Exception: + return False + return True + + +def _predicate_columns(expr: Expression, schema: pa.Schema) -> set[str]: + """The columns ``expr`` reads. + + Daft exposes no accessor for an expression's inputs, so this drops one column + at a time from a zero-row input and records which removals stop it resolving. + """ + return { + field.name + for field in schema + if not _evaluates_against(expr, pa.schema([f for f in schema if f.name != field.name])) + } + class LanceDataSink(DataSink[list[FragmentMetadata]]): """WriteSink for writing data to a Lance dataset.""" @@ -51,9 +115,11 @@ def __init__( self, uri: str | pathlib.Path | None, schema: Schema | pa.Schema, - mode: Literal["create", "append", "overwrite"] = "create", + mode: LanceWriteMode = "create", io_config: IOConfig | None = None, *, + predicate: str | None = None, + validate_predicate: bool = True, table_id: list[str] | None = None, namespace_impl: str | None = None, namespace_properties: dict[str, str] | None = None, @@ -70,11 +136,20 @@ def __init__( ) -> None: self._reject_unsupported_modes(mode, use_legacy_format) self._reject_namespace_mem_wal(namespace_impl, table_id, use_mem_wal) + self._validate_overwrite_where(mode, predicate, 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._mode = mode + self._is_overwrite_where = mode == "overwrite_where" + self._write_mode: LancePhysicalWriteMode = "append" if mode == "overwrite_where" else mode + self._predicate = predicate.strip() if predicate is not None else None + # Only meaningful for overwrite_where; other modes append nothing to filter. + self._validate_predicate = validate_predicate and self._is_overwrite_where + # Compiled lazily on whichever process evaluates it, so no daft Expression + # ever has to survive the pickling of this sink. + self._predicate_expr: Expression | None = None self._uri = uri self._namespace_impl = namespace_impl self._namespace_properties = namespace_properties @@ -133,17 +208,79 @@ def start(self) -> None: self._data_storage_version = resolve_storage_version( self._requested_storage_version, existing_version, - self._mode, + self._write_mode, ) # Auto-pick up any existing lance.blob.v2 columns when appending so the # write path wraps the matching daft binary columns. - if self._mode == "append" and self._table_schema is not None: + if self._write_mode == "append" and self._table_schema is not None: self._blob.add_columns(detect_blob_v2_columns(self._table_schema)) + if self._is_overwrite_where: + assert existing is not None, "overwrite_where requires an existing dataset" + self._validate_predicate_against_table(existing) + # 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) + def _validate_predicate_against_table(self, dataset: lance.LanceDataset) -> None: + """Fail on the driver, before any data is written, if the predicate is unusable. + + Planning a scan is enough to surface parse errors and unknown columns; + without this the write only fails at commit time, after the whole input + has been written to storage. When the input check is on, the predicate + must also hold up under Daft, which is checked here for the same reason. + """ + assert self._predicate is not None + try: + dataset.scanner(columns=[], filter=self._predicate, limit=1).explain_plan(True) + except Exception as e: + raise ValueError(f"predicate={self._predicate!r} is not a valid Lance filter for this table: {e}") from e + + if self._validate_predicate: + self._reject_untrusted_predicate(_compile_predicate(self._predicate), dataset.schema) + + def _reject_untrusted_predicate(self, expr: Expression, table_schema: pa.Schema) -> None: + """Refuse to run the input check when Daft would answer differently than Lance. + + Both problems below are silent at write time: the first surfaces as a raw + Daft type error from inside a worker, the second as input rows that pass + the check and are then never covered by Lance's delete. + """ + # Exactly the schema _prepare_arrow_table casts the input to, so this + # sees the types the check will actually evaluate against. + target_schema = self._blob.cast_target_schema(table_schema) + + if not _evaluates_against(expr, target_schema): + raise ValueError( + f"predicate={self._predicate!r} is a valid Lance filter, but Daft cannot evaluate it " + "against this table's schema, so the input rows cannot be checked against it (Daft " + "reads a bare TIMESTAMP literal as UTC-aware, for example, which will not compare " + "against a naive timestamp column). Pass validate_predicate=False to write without " + "the check." + ) + + # Daft widens a narrow float column to f64 before comparing it to a + # decimal literal (0.1f32 -> 0.10000000149...), where Lance narrows the + # literal to the column's type instead. "score > 0.1" therefore selects + # different rows in the two engines, and a row Daft accepts can be one + # Lance never deletes -- the duplication this check exists to prevent. + narrow_floats = sorted( + name + for name in _predicate_columns(expr, target_schema) + if pa.types.is_float32(target_schema.field(name).type) + or pa.types.is_float16(target_schema.field(name).type) + ) + if narrow_floats: + raise ValueError( + f"predicate={self._predicate!r} reads {', '.join(narrow_floats)}, which Lance stores " + "as a narrow float. Daft and Lance compare a decimal literal against such a column " + "differently, so a row that passes the input check may not be one Lance deletes, and " + "it would survive a re-run of this write. Compare against an exactly representable " + "value (0.5, 0.25), or pass validate_predicate=False and make sure the input really " + "is inside the predicate." + ) + @property def _namespace_kwargs(self) -> dict[str, Any]: return get_namespace_kwargs(self._namespace_impl, self._namespace_properties, self._table_id) @@ -164,7 +301,7 @@ def _dataset_uri_arg(self) -> str | None: 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" + mode = self._write_mode if self._write_mode in ("create", "overwrite") else "read" resolved = resolve_namespace_table( namespace_impl=self._namespace_impl, namespace_properties=self._namespace_properties, @@ -187,9 +324,7 @@ def _merged_storage_options(self, resolved: ResolvedNamespaceTable) -> dict[str, return merge_storage_options(io_derived, self._user_storage_options, resolved.storage_options) @staticmethod - def _reject_unsupported_modes( - mode: Literal["create", "append", "overwrite"], use_legacy_format: bool | None - ) -> None: + def _reject_unsupported_modes(mode: LanceWriteMode, use_legacy_format: bool | None) -> None: # This mode was never functional and customers must use merge_columns_df. if mode == "merge": # type: ignore[comparison-overlap] raise ValueError( @@ -207,6 +342,24 @@ def _reject_unsupported_modes( stacklevel=3, ) + @staticmethod + def _validate_overwrite_where(mode: LanceWriteMode, predicate: str | None, use_mem_wal: bool) -> None: + """Conditional overwrite needs a predicate, and only works copy-on-write.""" + if mode != "overwrite_where": + if predicate is not None: + raise ValueError(f'predicate is only supported with mode="overwrite_where", got mode="{mode}".') + return + if predicate is None or not predicate.strip(): + raise ValueError( + 'mode="overwrite_where" requires a non-empty SQL predicate, e.g. predicate="dt = \'2026-08-25\'".' + ) + if use_mem_wal: + raise ValueError( + 'mode="overwrite_where" is not supported with use_mem_wal=True. The conditional ' + "overwrite commits deletions against a pinned dataset version, which the mem-WAL " + "write path does not go through." + ) + @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. @@ -274,9 +427,9 @@ def _absorb_existing_dataset(self) -> lance.LanceDataset | None: raise 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 and self._table_uri is not None: + if self._write_mode == "append": + raise ValueError(f"Cannot {self._mode} to non-existent Lance dataset.") + if self._write_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.") @@ -286,13 +439,13 @@ def _absorb_existing_dataset(self) -> lance.LanceDataset | None: self._table_schema = table_schema self._version = dataset.latest_version - if self._mode == "create": + if self._write_mode == "create": 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( + if self._write_mode == "append" and not _pyarrow_schema_castable( blob_aware_schema_for_validation(self._pyarrow_schema, table_schema), blob_aware_schema_for_validation(table_schema, table_schema), ): @@ -324,7 +477,7 @@ def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetada fragments = lance.fragment.write_fragments( wrapped, dataset_uri=self._table_uri, - mode=self._mode, + mode=self._write_mode, storage_options=self._storage_options, max_rows_per_file=self._max_rows_per_file, max_rows_per_group=self._max_rows_per_group, @@ -393,6 +546,8 @@ def _write_cow(self, micropartitions: Iterator[MicroPartition]) -> Iterator[Writ for micropartition in micropartitions: arrow_table = self._prepare_arrow_table(micropartition.to_arrow()) + if self._validate_predicate: + self._assert_rows_match_predicate(arrow_table) # Oversized inputs flush whatever we already have, then write directly # so Lance can split internally. @@ -408,6 +563,32 @@ def _write_cow(self, micropartitions: Iterator[MicroPartition]) -> Iterator[Writ if buffer.has_rows(): yield self._write_arrow_table(buffer.drain()) + def _predicate_expression(self) -> Expression: + if self._predicate_expr is None: + assert self._predicate is not None + self._predicate_expr = _compile_predicate(self._predicate) + return self._predicate_expr + + def _assert_rows_match_predicate(self, table: pa.Table) -> None: + """Reject input rows the predicate does not select. + + overwrite_where deletes by predicate but appends the input verbatim, so a + row outside the predicate is not covered by the next run of the same + write: re-running it duplicates that row instead of replacing it. + """ + if table.num_rows == 0: + return + # Checked after the cast to the table schema, so the comparison sees the + # same types Lance will evaluate the predicate against. + matched = len(MicroPartition.from_arrow(table).filter(ExpressionsProjection([self._predicate_expression()]))) + if matched != table.num_rows: + raise ValueError( + f"{table.num_rows - matched} of {table.num_rows} input rows do not satisfy " + f'predicate={self._predicate!r}. mode="overwrite_where" appends the input as-is, so ' + "those rows would not be replaced by a re-run of this write. Filter the input, widen " + "the predicate, or pass validate_predicate=False to write them anyway." + ) + def _write_mem_wal( self, micropartitions: Iterator[MicroPartition] ) -> Iterator[WriteResult[list[FragmentMetadata]]]: @@ -429,6 +610,9 @@ 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)) + if self._is_overwrite_where: + return self._finalize_overwrite_where(fragments) + 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": @@ -446,16 +630,47 @@ def _finalize_cow(self, write_results: list[WriteResult[list[FragmentMetadata]]] storage_options=self._storage_options, **self._namespace_commit_kwargs, ) - stats = dataset.stats.dataset_stats() - stats_dict = MicroPartition.from_pydict( - { - "num_fragments": pa.array([stats["num_fragments"]], type=pa.int64()), - "num_deleted_rows": pa.array([stats["num_deleted_rows"]], type=pa.int64()), - "num_small_files": pa.array([stats["num_small_files"]], type=pa.int64()), - "version": pa.array([dataset.version], type=pa.int64()), - } + return _dataset_stats(dataset) + + def _finalize_overwrite_where(self, fragments: list[FragmentMetadata]) -> MicroPartition: + """Delete the predicate's rows and add this batch's fragments in one commit.""" + from daft_lance.lance_overwrite_where import apply_conditional_overwrite + from daft_lance.namespace import DatasetOpenContext + + assert self._table_uri is not None, "LanceDataSink.start() must run before finalize" + assert self._predicate is not None + + # Pinned to the version start() read: the deletions describe that snapshot, + # and the commit declares it as its read version. + pinned = lance.dataset( + self._dataset_uri_arg, + version=self._version, + storage_options=self._storage_options, + **self._namespace_kwargs, ) - return stats_dict + open_context = DatasetOpenContext.from_dataset( + pinned, + self._table_uri, + storage_options=self._storage_options, + namespace_impl=self._namespace_impl, + namespace_properties=self._namespace_properties, + table_id=self._table_id, + managed_versioning=self._managed_versioning, + ) + dataset = apply_conditional_overwrite( + open_context=open_context, + predicate=self._predicate, + new_fragments=fragments, + ) + if dataset is None: + logger.info( + "overwrite_where matched no rows and wrote no data for predicate %r; no version created", + self._predicate, + ) + dataset = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + ) + return _dataset_stats(dataset) def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadata]]]) -> MicroPartition: dataset = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) @@ -483,15 +698,7 @@ def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadat self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs ) - stats = dataset.stats.dataset_stats() - return MicroPartition.from_pydict( - { - "num_fragments": pa.array([stats["num_fragments"]], type=pa.int64()), - "num_deleted_rows": pa.array([stats["num_deleted_rows"]], type=pa.int64()), - "num_small_files": pa.array([stats["num_small_files"]], type=pa.int64()), - "version": pa.array([dataset.version], type=pa.int64()), - } - ) + return _dataset_stats(dataset) class _LanceFragmentBuffer: diff --git a/daft_lance/lance_overwrite_where.py b/daft_lance/lance_overwrite_where.py new file mode 100644 index 0000000..3dcaafb --- /dev/null +++ b/daft_lance/lance_overwrite_where.py @@ -0,0 +1,187 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, cast + +import lance +import pyarrow as pa +import pyarrow.compute as pc + +import daft.pickle +from daft import from_pylist +from daft.datatype import DataType +from daft.runners import get_or_create_runner +from daft.udf import cls as daft_cls +from daft.udf import method + +if TYPE_CHECKING: + from lance.fragment import FragmentMetadata + + from daft_lance.namespace import DatasetOpenContext + +logger = logging.getLogger(__name__) + +_FRAGMENT_DELETE_RETURN_DTYPE = DataType.struct( + { + "fragment_id": DataType.int64(), + "fragment_meta": DataType.binary(), + "removed": DataType.bool(), + } +) + +# Pinned to the plan node Lance emits when a scalar index answers the filter. A +# rename would silently cost us fragment pruning, so +# test_pruning_uses_a_scalar_index_when_one_covers_the_predicate asserts on it. +_SCALAR_INDEX_PLAN_MARKER = "ScalarIndexQuery" + +# A Lance row address packs the fragment id into its high 32 bits. +_FRAGMENT_ID_SHIFT = pa.scalar(32, type=pa.uint64()) + +# Each partition builds its own handler and reopens the pinned snapshot once, so +# this bounds the manifest reads a wide table pays for the extra parallelism. +_MAX_DELETE_PARTITIONS = 64 + + +@daft_cls +class FragmentDeleteHandler: + """Applies one delete predicate to a fragment and reports what changed. + + Runs as a Daft UDF: the driver ships fragment ids, each task reopens the + pinned snapshot and writes a deletion file for the rows the predicate + matches. Data files are never rewritten, so row addresses -- and every index + built on them -- stay valid. + """ + + def __init__(self, open_context: DatasetOpenContext, predicate: str) -> None: + self.open_context = open_context + self.predicate = predicate + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + # Opened once per instance, not per fragment: the reopen costs a pinned + # manifest read and 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 + + @method.batch(return_dtype=_FRAGMENT_DELETE_RETURN_DTYPE) + def __call__(self, fragment_ids: Any) -> list[dict[str, Any]]: + lance_ds = self._dataset() + results: list[dict[str, Any]] = [] + for fragment_id in fragment_ids: + fragment = lance_ds.get_fragment(fragment_id) + if fragment is None: + raise ValueError(f"Fragment {fragment_id} not found in dataset") + deletions_before = fragment.metadata.num_deletions + updated = fragment.delete(self.predicate) + if updated is None: + # Every row matched: the fragment leaves the dataset entirely. + results.append({"fragment_id": int(fragment_id), "fragment_meta": None, "removed": True}) + continue + # A fragment the predicate missed comes back unchanged; committing it + # as "updated" would only add noise to the transaction. + changed = updated.num_deletions != deletions_before + results.append( + { + "fragment_id": int(fragment_id), + "fragment_meta": daft.pickle.dumps(updated) if changed else None, + "removed": False, + } + ) + return results + + +def _candidate_fragment_ids(dataset: lance.LanceDataset, predicate: str) -> set[int] | None: + """Fragment ids that hold rows matching ``predicate``, or None when unknown. + + Only worth doing when a scalar index can answer the filter: then this is an + index lookup that skips most fragments. Without an index the scan costs the + same full pass the delete step already pays, so we return None and let the + delete visit every fragment rather than paying for both. + """ + scanner = dataset.scanner(columns=[], filter=predicate, with_row_address=True) + if _SCALAR_INDEX_PLAN_MARKER not in scanner.explain_plan(True): + return None + + fragment_ids: set[int] = set() + # Streamed, not to_table(): one overwritten partition can be hundreds of + # millions of row addresses, and we only need the ids they live in. + for batch in scanner.to_batches(): + batch_ids = pc.unique(pc.shift_right(batch.column("_rowaddr"), _FRAGMENT_ID_SHIFT)) + fragment_ids.update(cast("list[int]", batch_ids.to_pylist())) + return fragment_ids + + +def _delete_matching_rows( + open_context: DatasetOpenContext, + predicate: str, + fragment_ids: list[int], +) -> tuple[list[FragmentMetadata], list[int]]: + """Run the per-fragment delete as a Daft job; return (updated, removed).""" + if not fragment_ids: + return [], [] + + df = from_pylist([{"fragment_id": fragment_id} for fragment_id in fragment_ids]) + partitions = min(len(fragment_ids), _MAX_DELETE_PARTITIONS) + # from_pylist lands everything in one partition, which would pin the whole + # delete to a single task on a distributed runner. The native runner has no + # partitions to spread -- repartition there is a no-op that only warns. + if partitions > 1 and get_or_create_runner().name != "native": + df = df.repartition(partitions, "fragment_id") + handler = FragmentDeleteHandler(open_context, predicate) + df = df.with_column("delete_result", handler(df["fragment_id"])) # type: ignore[arg-type] + + updated_fragments: list[FragmentMetadata] = [] + removed_fragment_ids: list[int] = [] + for result in df.collect().to_pydict()["delete_result"]: + if result["removed"]: + removed_fragment_ids.append(int(result["fragment_id"])) + elif result["fragment_meta"] is not None: + updated_fragments.append(daft.pickle.loads(result["fragment_meta"])) + return updated_fragments, removed_fragment_ids + + +def apply_conditional_overwrite( + *, + open_context: DatasetOpenContext, + predicate: str, + new_fragments: list[FragmentMetadata], +) -> lance.LanceDataset | None: + """Delete the rows matching ``predicate`` and add ``new_fragments`` in one commit. + + ``open_context`` must be pinned to the version the write started from; that + version is what the commit declares as its read version, so the deletions + describe the snapshot they were computed against. + + Returns the committed dataset, or None when there was nothing to do (no + matching rows and no new data), in which case no version is created. + """ + pinned = open_context.open_pinned() + candidates = _candidate_fragment_ids(pinned, predicate) + if candidates is None: + fragment_ids = [fragment.fragment_id for fragment in pinned.get_fragments()] + logger.info("No scalar index covers %r; running delete over all %d fragments", predicate, len(fragment_ids)) + else: + fragment_ids = sorted(candidates) + logger.info("Scalar index pruned delete for %r down to %d fragments", predicate, len(fragment_ids)) + + updated_fragments, removed_fragment_ids = _delete_matching_rows(open_context, predicate, fragment_ids) + + if not updated_fragments and not removed_fragment_ids and not new_fragments: + return None + + operation = lance.LanceOperation.Update( + removed_fragment_ids=removed_fragment_ids, + updated_fragments=updated_fragments, + new_fragments=list(new_fragments), + # Deletions do not change any field's values, so no index needs to be + # dropped from the fragments that survive. + fields_modified=[], + ) + return lance.LanceDataset.commit( + open_context.uri, + operation, + read_version=open_context.version, + storage_options=open_context.storage_options, + **open_context.commit_kwargs, + ) diff --git a/tests/io/lancedb/test_overwrite_where.py b/tests/io/lancedb/test_overwrite_where.py new file mode 100644 index 0000000..c64a901 --- /dev/null +++ b/tests/io/lancedb/test_overwrite_where.py @@ -0,0 +1,413 @@ +"""Conditional overwrite: ``mode="overwrite_where"`` replaces a predicate's rows in one commit.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import lance +import pyarrow as pa +import pytest + +import daft +import daft_lance +from daft.recordbatch import MicroPartition +from daft_lance.lance_data_sink import LanceDataSink, _compile_predicate +from daft_lance.lance_overwrite_where import _SCALAR_INDEX_PLAN_MARKER, _candidate_fragment_ids + + +def _seed(uri: str) -> None: + """Three fragments; the first and second each mix two ``dt`` values.""" + daft_lance.write_lance( + daft.from_pydict({"dt": ["d1", "d2"], "id": [1, 2]}), uri, mode="create", max_rows_per_file=2 + ).collect() + daft_lance.write_lance( + daft.from_pydict({"dt": ["d2", "d3"], "id": [3, 4]}), uri, mode="append", max_rows_per_file=2 + ).collect() + daft_lance.write_lance( + daft.from_pydict({"dt": ["d2", "d2"], "id": [5, 6]}), uri, mode="append", max_rows_per_file=2 + ).collect() + + +def _rows(uri: str) -> list[tuple[str, int]]: + table = lance.dataset(uri).to_table().to_pydict() + return sorted(zip(table["dt"], table["id"])) + + +def _overwrite(uri: str, dts: list[str | None], ids: list[int], predicate: str, **kwargs: Any) -> dict[str, list[Any]]: + return daft_lance.write_lance( + daft.from_pydict({"dt": dts, "id": ids}), uri, mode="overwrite_where", predicate=predicate, **kwargs + ).to_pydict() + + +def test_replaces_only_matching_rows_in_one_version(tmp_path: Path) -> None: + uri = str(tmp_path / "tbl") + _seed(uri) + before = lance.dataset(uri).version + + stats = _overwrite(uri, ["d2", "d2"], [100, 101], "dt = 'd2'") + + assert _rows(uri) == [("d1", 1), ("d2", 100), ("d2", 101), ("d3", 4)] + # One commit, not a delete followed by an append: readers never see the gap. + assert lance.dataset(uri).version == before + 1 + assert stats["version"] == [before + 1] + + +def test_rerunning_the_same_overwrite_is_idempotent(tmp_path: Path) -> None: + uri = str(tmp_path / "tbl") + _seed(uri) + expected = [("d1", 1), ("d2", 100), ("d2", 101), ("d3", 4)] + + _overwrite(uri, ["d2", "d2"], [100, 101], "dt = 'd2'") + assert _rows(uri) == expected + # The second run has to delete the rows the first one appended, which sit in + # a fragment that did not exist when the first run planned its delete. + _overwrite(uri, ["d2", "d2"], [100, 101], "dt = 'd2'") + + assert _rows(uri) == expected + + +def test_predicate_matching_nothing_only_appends(tmp_path: Path) -> None: + uri = str(tmp_path / "tbl") + _seed(uri) + before = lance.dataset(uri).version + + _overwrite(uri, ["d9"], [42], "dt = 'd9'") + + assert _rows(uri) == [("d1", 1), ("d2", 2), ("d2", 3), ("d2", 5), ("d2", 6), ("d3", 4), ("d9", 42)] + assert lance.dataset(uri).version == before + 1 + + +def test_empty_input_deletes_the_matched_rows(tmp_path: Path) -> None: + uri = str(tmp_path / "tbl") + _seed(uri) + before = lance.dataset(uri).version + empty = daft.from_pydict({"dt": ["d2"], "id": [1]}).limit(0) + + daft_lance.write_lance(empty, uri, mode="overwrite_where", predicate="dt = 'd2'").collect() + + assert _rows(uri) == [("d1", 1), ("d3", 4)] + assert lance.dataset(uri).version == before + 1 + + +def test_fully_matched_fragment_is_removed_not_just_emptied(tmp_path: Path) -> None: + uri = str(tmp_path / "tbl") + _seed(uri) + # The third seeded fragment is all "d2", so the overwrite drops it entirely + # while the mixed fragments only gain deletion files. + _overwrite(uri, ["d2"], [100], "dt = 'd2'") + + fragments = lance.dataset(uri).get_fragments() + assert sum(fragment.count_rows() for fragment in fragments) == 3 + assert all(fragment.count_rows() > 0 for fragment in fragments) + + +def test_rows_outside_the_predicate_are_rejected(tmp_path: Path) -> None: + uri = str(tmp_path / "tbl") + _seed(uri) + before = _rows(uri) + + with pytest.raises(Exception, match="do not satisfy"): + _overwrite(uri, ["d2", "d9"], [100, 101], "dt = 'd2'") + + assert _rows(uri) == before + + +def test_null_rows_count_as_not_satisfying_the_predicate(tmp_path: Path) -> None: + uri = str(tmp_path / "tbl") + _seed(uri) + + # ``dt = 'd2'`` is NULL, not False, for a NULL dt: a naive "keep rows where + # NOT predicate" check would let this row through. + with pytest.raises(Exception, match="do not satisfy"): + _overwrite(uri, ["d2", None], [100, 101], "dt = 'd2'") + + +def test_validate_predicate_false_appends_rows_outside_the_predicate(tmp_path: Path) -> None: + uri = str(tmp_path / "tbl") + _seed(uri) + + _overwrite(uri, ["d2", "d9"], [100, 101], "dt = 'd2'", validate_predicate=False) + + assert _rows(uri) == [("d1", 1), ("d2", 100), ("d3", 4), ("d9", 101)] + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"mode": "overwrite_where"}, "requires a non-empty SQL predicate"), + ({"mode": "overwrite_where", "predicate": " "}, "requires a non-empty SQL predicate"), + ({"mode": "append", "predicate": "dt = 'd2'"}, 'only supported with mode="overwrite_where"'), + ({"mode": "overwrite_where", "predicate": "dt = 'd2'", "use_mem_wal": True}, "not supported with use_mem_wal"), + ], +) +def test_argument_validation(tmp_path: Path, kwargs: dict[str, Any], match: str) -> None: + with pytest.raises(ValueError, match=match): + daft_lance.write_lance(daft.from_pydict({"dt": ["d2"], "id": [1]}), str(tmp_path / "tbl"), **kwargs) + + +def test_requires_an_existing_table(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="Cannot overwrite_where to non-existent Lance dataset"): + _overwrite(str(tmp_path / "missing"), ["d2"], [1], "dt = 'd2'") + + +def test_schema_must_match_like_append(tmp_path: Path) -> None: + uri = str(tmp_path / "tbl") + _seed(uri) + + with pytest.raises(ValueError, match="Schema of data does not match table schema"): + daft_lance.write_lance( + daft.from_pydict({"dt": ["d2"]}), uri, mode="overwrite_where", predicate="dt = 'd2'" + ).collect() + + +def test_storage_version_conflict_is_detected_like_append(tmp_path: Path) -> None: + """Regression guard for the mode normalization. + + ``resolve_storage_version`` only checks the "append" mode; before + ``overwrite_where`` was normalized to it, a conflicting version was accepted + silently. + """ + uri = str(tmp_path / "tbl") + daft_lance.write_lance( + daft.from_pydict({"dt": ["d2"], "id": [1]}), uri, mode="create", data_storage_version="2.1" + ).collect() + + with pytest.raises(ValueError, match="does not match existing dataset version"): + _overwrite(uri, ["d2"], [2], "dt = 'd2'", data_storage_version="2.0") + + +@pytest.mark.parametrize("predicate", ["nosuchcol = 1", "dt ==== 'x'"]) +def test_predicate_must_be_a_valid_lance_filter(tmp_path: Path, predicate: str) -> None: + """Bad predicates fail on the driver, before any data is written.""" + uri = str(tmp_path / "tbl") + _seed(uri) + + with pytest.raises(ValueError, match="is not a valid Lance filter"): + _overwrite(uri, ["d2"], [100], predicate) + + +def test_unparseable_predicate_points_at_the_escape_hatch() -> None: + with pytest.raises(ValueError, match="validate_predicate=False"): + _compile_predicate("dt ==== 'x'") + + +def test_pruning_uses_a_scalar_index_when_one_covers_the_predicate(tmp_path: Path) -> None: + uri = str(tmp_path / "tbl") + _seed(uri) + dataset = lance.dataset(uri) + + # No index: every fragment has to be visited, and the planner says so. + assert _candidate_fragment_ids(dataset, "id = 4") is None + + dataset.create_scalar_index("id", "BTREE") + dataset = lance.dataset(uri) + + # Pinned plan-node name: pruning silently stops working if Lance renames it. + plan = dataset.scanner(columns=[], filter="id = 4", with_row_address=True).explain_plan(True) + assert _SCALAR_INDEX_PLAN_MARKER in plan + + # id 4 lives only in the second seeded fragment. + assert _candidate_fragment_ids(dataset, "id = 4") == {1} + assert _candidate_fragment_ids(dataset, "id > 2") == {1, 2} + assert _candidate_fragment_ids(dataset, "id = 999") == set() + + +def test_overwrite_through_a_scalar_index_on_the_predicate_column(tmp_path: Path) -> None: + """The pruning path, end to end -- a miss here silently leaves rows behind.""" + uri = str(tmp_path / "tbl") + daft_lance.write_lance( + daft.from_pydict({"day": [1, 1, 2, 2, 3], "id": [1, 2, 3, 4, 5]}), uri, mode="create", max_rows_per_file=2 + ).collect() + lance.dataset(uri).create_scalar_index("day", "BTREE") + + def rows() -> list[tuple[int, int]]: + table = lance.dataset(uri).to_table().to_pydict() + return sorted(zip(table["day"], table["id"])) + + def overwrite(new_id: int) -> None: + daft_lance.write_lance( + daft.from_pydict({"day": [2], "id": [new_id]}), uri, mode="overwrite_where", predicate="day = 2" + ).collect() + + assert _candidate_fragment_ids(lance.dataset(uri), "day = 2") is not None, "expected the pruning path" + overwrite(100) + assert rows() == [(1, 1), (1, 2), (2, 100), (3, 5)] + + # The index does not cover the fragment the first overwrite appended, so this + # run only replaces its row if pruning still finds that fragment. + assert _candidate_fragment_ids(lance.dataset(uri), "day = 2") is not None, "expected the pruning path" + overwrite(200) + assert rows() == [(1, 1), (1, 2), (2, 200), (3, 5)] + + +def test_indexed_table_stays_queryable_after_overwrite(tmp_path: Path) -> None: + uri = str(tmp_path / "tbl") + n = 300 + vector_type = pa.list_(pa.float32(), 2) + seed = pa.table( + { + "id": pa.array(range(n), pa.int64()), + "dt": pa.array(["d1" if i % 2 else "d2" for i in range(n)]), + "vector": pa.array([[float(i % 3), 0.0] for i in range(n)], type=vector_type), + } + ) + lance.write_dataset(seed, uri, max_rows_per_file=100) + dataset = lance.dataset(uri) + dataset.create_scalar_index("id", "BTREE") + try: + dataset.create_index("vector", "IVF_PQ", num_partitions=2, num_sub_vectors=1) + except Exception: + pytest.skip("Could not create vector index (lance version or dataset size issue)") + + new_rows = pa.table( + { + "id": pa.array([1000, 1001], pa.int64()), + "dt": pa.array(["d2", "d2"]), + "vector": pa.array([[7.0, 7.0], [8.0, 8.0]], type=vector_type), + } + ) + daft_lance.write_lance(daft.from_arrow(new_rows), uri, mode="overwrite_where", predicate="dt = 'd2'").collect() + + dataset = lance.dataset(uri) + # Deleted rows are invisible through the scalar index that still covers them. + assert dataset.to_table(filter="id = 0").num_rows == 0 + assert dataset.to_table(filter="id = 1").num_rows == 1 + assert sorted(dataset.to_table(filter="dt = 'd2'").to_pydict()["id"]) == [1000, 1001] + + # New fragments are not in the index; the search must still find them. + nearest = {"column": "vector", "q": pa.array([8.0, 8.0], type=pa.float32()), "k": 1, "use_index": True} + assert daft.read_lance(uri, default_scan_options={"nearest": nearest}).select("id").to_pydict()["id"] == [1001] + + +def _float_table(tmp_path: Path) -> str: + uri = str(tmp_path / "floats") + lance.write_dataset( + pa.table( + { + "score": pa.array([0.1, 0.2], pa.float32()), + "dt": pa.array(["d1", "d2"], pa.large_string()), + "id": pa.array([1, 2], pa.int64()), + } + ), + uri, + ) + return uri + + +def _float_rows(uri: str, dt: str, id_: int, score: float = 0.1) -> daft.DataFrame: + return daft.from_arrow( + pa.table( + { + "score": pa.array([score], pa.float32()), + "dt": pa.array([dt], pa.large_string()), + "id": pa.array([id_], pa.int64()), + } + ) + ) + + +def test_narrow_float_predicate_refuses_the_input_check(tmp_path: Path) -> None: + """Daft widens float32 before comparing; Lance narrows the literal instead. + + "score > 0.1" is false in Lance for a 0.1f row but true in Daft, so trusting + the check would let a row through that the delete never covers -- it would + pile up on every re-run. + """ + uri = _float_table(tmp_path) + + with pytest.raises(ValueError, match="narrow float"): + daft_lance.write_lance( + _float_rows(uri, "d2", 99), uri, mode="overwrite_where", predicate="score > 0.1" + ).collect() + + # The opt-out still writes: the caller takes responsibility for the input. + daft_lance.write_lance( + _float_rows(uri, "d2", 99), uri, mode="overwrite_where", predicate="score > 0.1", validate_predicate=False + ).collect() + assert sorted(lance.dataset(uri).to_table().to_pydict()["id"]) == [1, 99] + + +def test_float_column_outside_the_predicate_still_validates(tmp_path: Path) -> None: + """Only the columns the predicate reads matter, not every float in the table.""" + uri = _float_table(tmp_path) + + daft_lance.write_lance( + _float_rows(uri, "d2", 99, score=0.7), uri, mode="overwrite_where", predicate="dt = 'd2'" + ).collect() + + assert sorted(lance.dataset(uri).to_table().to_pydict()["id"]) == [1, 99] + with pytest.raises(Exception, match="do not satisfy"): + daft_lance.write_lance( + _float_rows(uri, "d9", 100), uri, mode="overwrite_where", predicate="dt = 'd2'" + ).collect() + + +def test_predicate_daft_cannot_evaluate_fails_before_writing(tmp_path: Path) -> None: + """A Lance-valid filter Daft types differently must not fail mid-write.""" + uri = str(tmp_path / "events") + timestamps = pa.array([1787529600000000, 1787616000000000], pa.timestamp("us")) + lance.write_dataset(pa.table({"t": timestamps, "id": pa.array([1, 2], pa.int64())}), uri) + predicate = "t >= TIMESTAMP '2026-08-25 00:00:00' AND t < TIMESTAMP '2026-08-26 00:00:00'" + new_row = daft.from_arrow( + pa.table({"t": pa.array([1787529600000000], pa.timestamp("us")), "id": pa.array([50], pa.int64())}) + ) + + with pytest.raises(ValueError, match="validate_predicate=False"): + daft_lance.write_lance(new_row, uri, mode="overwrite_where", predicate=predicate).collect() + assert sorted(lance.dataset(uri).to_table().to_pydict()["id"]) == [1, 2] + + daft_lance.write_lance( + new_row, uri, mode="overwrite_where", predicate=predicate, validate_predicate=False + ).collect() + assert 50 in lance.dataset(uri).to_table().to_pydict()["id"] + + +def test_namespace_addressed_table(tmp_path: Path) -> None: + ns: dict[str, Any] = {"namespace_impl": "dir", "namespace_properties": {"root": str(tmp_path)}} + table_id = ["events"] + + daft_lance.write_lance( + daft.from_pydict({"dt": ["d1", "d2"], "id": [1, 2]}), table_id=table_id, mode="create", **ns + ).collect() + daft_lance.write_lance( + daft.from_pydict({"dt": ["d2"], "id": [100]}), + table_id=table_id, + mode="overwrite_where", + predicate="dt = 'd2'", + **ns, + ).collect() + + result = daft_lance.read_lance(table_id=table_id, **ns).to_pydict() + assert sorted(zip(result["dt"], result["id"])) == [("d1", 1), ("d2", 100)] + + +def test_concurrent_append_survives_the_overwrite(tmp_path: Path) -> None: + """Documents a known gap, so a future Lance change surfaces here. + + Lance does not treat a concurrent append as conflicting with the Update this + mode commits, so rows another writer adds while the overwrite runs stay in + the table even when they match the predicate -- and nothing raises. The + docstring on ``write_lance`` warns about it; this test pins the behavior. + """ + uri = str(tmp_path / "tbl") + _seed(uri) + + sink = LanceDataSink( + uri=uri, + schema=daft.from_pydict({"dt": ["d2"], "id": [1]}).schema(), + mode="overwrite_where", + predicate="dt = 'd2'", + ) + sink.start() # pins the read version + + concurrent = pa.table({"dt": pa.array(["d2"], pa.large_string()), "id": pa.array([900], pa.int64())}) + lance.write_dataset(concurrent, uri, mode="append") + + results = list(sink.write(iter([MicroPartition.from_pydict({"dt": ["d2"], "id": [100]})]))) + sink.finalize(results) + + # The overwrite replaced every "d2" row it knew about, and the concurrent one + # it could not see survived: asserting the whole table keeps this test honest + # if the overwrite ever silently turns into a no-op. + assert _rows(uri) == [("d1", 1), ("d2", 100), ("d2", 900), ("d3", 4)]