From 36f382d6698da4cc3fc390062206ae3398ad0228 Mon Sep 17 00:00:00 2001 From: fanng Date: Tue, 25 Aug 2026 22:08:40 +0900 Subject: [PATCH 1/2] feat: support conditional overwrite writes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds mode="overwrite_where" to daft_lance.write_lance: one Lance commit deletes the rows matching a SQL predicate and adds the new data, so readers see either the whole replacement or none of it. The delete runs distributed, one task per fragment, reusing the DatasetOpenContext pattern the merge and compaction paths already use. When a scalar index covers the predicate, a row-address lookup prunes the fragment list first; without one, every fragment goes to the workers rather than paying for a serialized scan on the driver. Input rows are checked against the predicate by default: overwrite_where appends the input verbatim, so a row outside the predicate would survive the next run of the same write instead of being replaced. validate_predicate=False opts out. The new mode writes exactly like an append and is normalized to it in the constructor, so the existing schema, blob-column and storage-version checks cover it -- adding a fourth mode string to those checks would have skipped them silently. Known gap, documented in write_lance and the README: Lance does not treat a concurrent append as conflicting with the Update this commits, so rows another writer adds during the overwrite survive it. test_concurrent_append_survives _the_overwrite pins that behavior. Co-authored-by: fanng <“fanng@apache.org”> Claude-Session: https://claude.ai/code/session_017b8TVnEmyN8wbXFTHJwYip --- README.md | 25 ++ daft_lance/_lance.py | 36 ++- daft_lance/lance_data_sink.py | 199 ++++++++++++--- daft_lance/lance_overwrite_where.py | 178 ++++++++++++++ tests/io/lancedb/test_overwrite_where.py | 296 +++++++++++++++++++++++ 5 files changed, 699 insertions(+), 35 deletions(-) create mode 100644 daft_lance/lance_overwrite_where.py create mode 100644 tests/io/lancedb/test_overwrite_where.py diff --git a/README.md b/README.md index 106d6bd..8b99786 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,31 @@ 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). + +> **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..2cc5640 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,20 @@ 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. 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 +677,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 +707,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..3551a39 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,43 @@ 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: + 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 + class LanceDataSink(DataSink[list[FragmentMetadata]]): """WriteSink for writing data to a Lance dataset.""" @@ -51,9 +87,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 +108,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 +180,37 @@ 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 bad. + + 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. + """ + 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: + _compile_predicate(self._predicate) + @property def _namespace_kwargs(self) -> dict[str, Any]: return get_namespace_kwargs(self._namespace_impl, self._namespace_properties, self._table_id) @@ -164,7 +231,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 +254,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 +272,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 +357,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 +369,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 +407,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 +476,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 +493,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 +540,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 +560,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, + ) + 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, ) - return stats_dict + 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 +628,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..7337475 --- /dev/null +++ b/daft_lance/lance_overwrite_where.py @@ -0,0 +1,178 @@ +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.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()) + + +class FragmentDeleteHandler: + """Applies one delete predicate to a fragment and reports what changed. + + Runs on workers: 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 would cost + the same full pass the delete step already pays, except serialized on the + driver -- so we return None and let every fragment go to the workers. + """ + 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], + concurrency: int | None, +) -> tuple[list[FragmentMetadata], list[int]]: + """Run the per-fragment delete across the cluster; return (updated, removed).""" + if not fragment_ids: + return [], [] + + df = from_pylist([{"fragment_id": fragment_id} for fragment_id in fragment_ids]) + handler_cls = daft_cls(FragmentDeleteHandler, max_concurrency=concurrency) + handler = handler_cls(open_context, predicate) + df = df.with_column("delete_result", handler(df["fragment_id"])) + + 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], + concurrency: int | None = None, +) -> 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, concurrency) + + 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..4273ebc --- /dev/null +++ b/tests/io/lancedb/test_overwrite_where.py @@ -0,0 +1,296 @@ +"""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) + + _overwrite(uri, ["d2", "d2"], [100, 101], "dt = 'd2'") + first = _rows(uri) + _overwrite(uri, ["d2", "d2"], [100, 101], "dt = 'd2'") + + assert _rows(uri) == first + + +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_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 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) + + assert ("d2", 900) in _rows(uri) From bfcb52cb12fae70faa6a230e5c6aed2372ec92c6 Mon Sep 17 00:00:00 2001 From: fanng Date: Tue, 25 Aug 2026 23:17:44 +0900 Subject: [PATCH 2/2] fix: make the overwrite_where input check trustworthy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review of the previous commit found two ways the input-row check could hurt rather than help, both confirmed against real datasets. The check evaluates the predicate with Daft while the delete runs with Lance/DataFusion, and the two engines disagree on a decimal literal compared against a narrow float column: Daft widens 0.1f32 to 0.10000000149 and calls "score > 0.1" true, Lance narrows the literal and calls it false. A row like that passed validation and was then never deleted, so re-running the same write piled up duplicates -- exactly what the check exists to prevent. A predicate reading a float32/float16 column is now refused up front. The driver-side check also only parsed the Daft expression, never evaluated it, so a Lance-valid predicate Daft types differently ("t >= TIMESTAMP '...'" against a naive timestamp column, the shape the docstring advertises) failed mid-write with a raw DaftTypeError and no pointer to validate_predicate=False. It is now evaluated against a zero-row input typed like the table. Also: from_pylist produces a single partition, so the per-fragment delete ran as one task on a distributed runner. It is now repartitioned by fragment id, except on the native runner where repartition is a no-op that only warns. The dead concurrency parameter is gone and the docstrings that claimed the work spread across the cluster now match what it does. Tests: the concurrency and idempotence tests asserted so little that they passed with the overwrite stubbed out to a no-op, and no test drove the index-pruning path end to end. Both now assert whole row sets, and a new test overwrites twice through an index on the predicate column -- the second run only replaces the row the first one appended if pruning still finds that fragment. Stubbing the overwrite to a no-op now fails 13 of 26 tests instead of 7 of 22, and a pruning miss fails 2. Co-authored-by: fanng <“fanng@apache.org”> Claude-Session: https://claude.ai/code/session_017b8TVnEmyN8wbXFTHJwYip --- README.md | 6 +- daft_lance/_lance.py | 7 +- daft_lance/lance_data_sink.py | 76 +++++++++++++- daft_lance/lance_overwrite_where.py | 37 ++++--- tests/io/lancedb/test_overwrite_where.py | 123 ++++++++++++++++++++++- 5 files changed, 227 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index 8b99786..79b3ec8 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,11 @@ daft_lance.write_lance( 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). +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 diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 2cc5640..c668f18 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -661,7 +661,12 @@ def write_lance( 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. + 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. diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index 3551a39..7303577 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -70,6 +70,12 @@ def _dataset_stats(dataset: lance.LanceDataset) -> MicroPartition: 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: @@ -80,6 +86,28 @@ def _compile_predicate(predicate: str) -> Expression: ) 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.""" @@ -196,11 +224,12 @@ def start(self) -> None: 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 bad. + """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. + 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: @@ -209,7 +238,48 @@ def _validate_predicate_against_table(self, dataset: lance.LanceDataset) -> None raise ValueError(f"predicate={self._predicate!r} is not a valid Lance filter for this table: {e}") from e if self._validate_predicate: - _compile_predicate(self._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]: diff --git a/daft_lance/lance_overwrite_where.py b/daft_lance/lance_overwrite_where.py index 7337475..3dcaafb 100644 --- a/daft_lance/lance_overwrite_where.py +++ b/daft_lance/lance_overwrite_where.py @@ -10,6 +10,7 @@ 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 @@ -36,14 +37,19 @@ # 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 on workers: 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. + 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: @@ -89,9 +95,9 @@ def _candidate_fragment_ids(dataset: lance.LanceDataset, predicate: str) -> set[ """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 would cost - the same full pass the delete step already pays, except serialized on the - driver -- so we return None and let every fragment go to the workers. + 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): @@ -110,16 +116,20 @@ def _delete_matching_rows( open_context: DatasetOpenContext, predicate: str, fragment_ids: list[int], - concurrency: int | None, ) -> tuple[list[FragmentMetadata], list[int]]: - """Run the per-fragment delete across the cluster; return (updated, removed).""" + """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]) - handler_cls = daft_cls(FragmentDeleteHandler, max_concurrency=concurrency) - handler = handler_cls(open_context, predicate) - df = df.with_column("delete_result", handler(df["fragment_id"])) + 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] = [] @@ -136,7 +146,6 @@ def apply_conditional_overwrite( open_context: DatasetOpenContext, predicate: str, new_fragments: list[FragmentMetadata], - concurrency: int | None = None, ) -> lance.LanceDataset | None: """Delete the rows matching ``predicate`` and add ``new_fragments`` in one commit. @@ -156,7 +165,7 @@ def apply_conditional_overwrite( 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, concurrency) + 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 diff --git a/tests/io/lancedb/test_overwrite_where.py b/tests/io/lancedb/test_overwrite_where.py index 4273ebc..c64a901 100644 --- a/tests/io/lancedb/test_overwrite_where.py +++ b/tests/io/lancedb/test_overwrite_where.py @@ -56,12 +56,15 @@ def test_replaces_only_matching_rows_in_one_version(tmp_path: Path) -> None: 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'") - first = _rows(uri) + 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) == first + assert _rows(uri) == expected def test_predicate_matching_nothing_only_appends(tmp_path: Path) -> None: @@ -210,6 +213,34 @@ def test_pruning_uses_a_scalar_index_when_one_covers_the_predicate(tmp_path: Pat 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 @@ -249,6 +280,89 @@ def test_indexed_table_stays_queryable_after_overwrite(tmp_path: Path) -> None: 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"] @@ -293,4 +407,7 @@ def test_concurrent_append_survives_the_overwrite(tmp_path: Path) -> None: results = list(sink.write(iter([MicroPartition.from_pydict({"dt": ["d2"], "id": [100]})]))) sink.finalize(results) - assert ("d2", 900) in _rows(uri) + # 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)]