From 539e8a72f1c0ec406ba3381be75bc053885d2899 Mon Sep 17 00:00:00 2001 From: fanng Date: Thu, 8 Oct 2026 18:29:47 +0900 Subject: [PATCH 1/9] feat: add transform-driven column updates --- README.md | 21 +- daft_lance/__init__.py | 2 + daft_lance/_lance.py | 113 +++++++- daft_lance/lance_update_column.py | 397 +++++++++++++++++++++++--- tests/io/lance/test_update_columns.py | 156 ++++++++++ 5 files changed, 650 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index ee6d1b5..90cd882 100644 --- a/README.md +++ b/README.md @@ -220,6 +220,25 @@ later re-run, so callers should filter the input first when they need idempotent ### Column Updates +Use a row-preserving transform to overwrite existing columns directly from the +current Lance snapshot. `where` is evaluated before the transform, and Daft +distributes work by fragment without grouping a user DataFrame: + +```python +from daft_lance import update_columns + +result = update_columns( + "s3://bucket/my_dataset", + transform={"score": "score * 2", "status": "'processed'"}, + where="date >= DATE '2026-07-01'", +) +print(result.version, result.rows_updated) +``` + +Python callables receive Arrow record batches after filtering and must return +exactly the existing columns named by `columns`, with the same row count and +order. Use `read_columns` to restrict their input. + Use a prepared Daft DataFrame to overwrite existing Lance columns while preserving row addresses and untouched column files: @@ -274,7 +293,7 @@ Address Lance tables through a [Lance Namespace](https://lancedb.github.io/lance (catalog) instead of a raw URI. Pass `namespace_impl` + `namespace_properties` + `table_id` in place of `uri` — the namespace resolves the table's storage location and vends any storage credentials. This works across `read_lance`, `write_lance`, `merge_columns_df`, -`update_columns_df`, `create_scalar_index`, and `compact_files`. +`update_columns`, `update_columns_df`, `create_scalar_index`, and `compact_files`. ```python import daft diff --git a/daft_lance/__init__.py b/daft_lance/__init__.py index 7397441..5819b62 100644 --- a/daft_lance/__init__.py +++ b/daft_lance/__init__.py @@ -12,6 +12,7 @@ merge_columns_df, optimize_indices, read_lance, + update_columns, update_columns_df, write_lance, ) @@ -28,6 +29,7 @@ "read_blobs", "read_lance", "take_blobs", + "update_columns", "update_columns_df", "write_lance", ] diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index c43e496..3ceaefd 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -18,7 +18,12 @@ from .lance_merge_column import merge_columns_from_df, merge_columns_internal from .lance_scalar_index import create_scalar_index_internal, optimize_indices_internal from .lance_scan import LanceScanOperator -from .lance_update_column import UpdateColumnsResult, update_columns_from_df, validate_update_arguments +from .lance_update_column import ( + UpdateColumnsResult, + update_columns_from_df, + update_columns_with_transform, + validate_update_arguments, +) from .lance_vector_index import create_vector_index_internal from .namespace import validate_uri_or_namespace from .utils import construct_lance_dataset_handle @@ -391,6 +396,112 @@ def merge_columns_df( ) +@PublicAPI +def update_columns( + uri: str | pathlib.Path | None = None, + io_config: IOConfig | None = None, + *, + transform: dict[str, str] | BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch], + where: str | None = None, + columns: Sequence[str] | None = None, + read_columns: Sequence[str] | None = None, + batch_size: int | None = None, + table_id: list[str] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, + storage_options: dict[str, Any] | None = None, + version: int | str | None = None, + asof: str | None = None, + block_size: int | None = None, + commit_lock: Any | None = None, + index_cache_size: int | None = None, + default_scan_options: dict[str, Any] | None = None, + metadata_cache_size_bytes: int | None = None, + max_concurrency: int | None = None, +) -> UpdateColumnsResult: + """Overwrite existing Lance columns with a row-preserving transform. + + The target snapshot is scanned fragment by fragment. Rows matching ``where`` + are passed through ``transform`` and joined back to their fragment by the + internal ``_rowaddr``. Daft distributes only fragment identifiers; callers do + not need to build or group an update DataFrame. + + Args: + uri: URI of the target Lance dataset. Mutually exclusive with namespace parameters. + io_config: Daft object-store configuration. + transform: Either a mapping from existing column names to Lance SQL + expressions, a Lance ``BatchUDF``, or a callable accepting and + returning an Arrow ``RecordBatch``. Callable output must preserve row + count and order and contain only the columns being updated. + where: Optional Lance SQL predicate evaluated against the pinned source + snapshot before the transform. ``None`` updates every live row. + columns: Existing top-level columns to overwrite. Inferred from dict keys + or a ``BatchUDF.output_schema``; required for other callables. + read_columns: Columns passed to callable transforms. Defaults to all table + columns. Ignored for SQL-expression transforms. + batch_size: Maximum scan batch size per fragment. + table_id: Namespace table identifier. + namespace_impl: Lance Namespace implementation. + namespace_properties: Properties used to connect to the namespace. + storage_options: Additional object-store options. + version: Dataset version or tag to read and commit against. Defaults to + the version current when this function starts. + asof: Find the latest version created on or before this timestamp. + block_size: Hint for the minimum object-store I/O request size. + commit_lock: Custom Lance commit lock. + index_cache_size: Index cache size. + default_scan_options: Default Lance scan options. + metadata_cache_size_bytes: Metadata cache size in bytes. + max_concurrency: Maximum number of concurrent fragment-update workers. + + Returns: + The committed dataset version and the number of live rows that matched + ``where``. If no rows match, no version is created. + + Note: + A fragment failure can leave unreferenced data files written by other + workers. No transaction is committed; Lance cleanup removes those files. + + Raises: + NotImplementedError: If the dataset uses stable row IDs, because the + installed pylance API cannot return matched fragment offsets. + """ + io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config + dataset_handle = construct_lance_dataset_handle( + uri, + storage_options=storage_options, + io_config=io_config, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, + version=version, + asof=asof, + block_size=block_size, + commit_lock=commit_lock, + index_cache_size=index_cache_size, + default_scan_options=default_scan_options, + metadata_cache_size_bytes=metadata_cache_size_bytes, + ) + if dataset_handle.dataset.has_stable_row_ids: + raise NotImplementedError( + "update_columns does not support datasets with stable row IDs: " + "pylance does not yet expose updated fragment offsets, so " + "_row_last_updated_at_version and CDF metadata cannot be updated correctly." + ) + + return update_columns_with_transform( + dataset_handle.dataset, + dataset_handle.worker_open_context(), + transform=transform, + columns=columns, + read_columns=read_columns, + where=where, + batch_size=batch_size, + commit_lock=commit_lock, + max_concurrency=max_concurrency, + ) + + @PublicAPI def update_columns_df( df: DataFrame, diff --git a/daft_lance/lance_update_column.py b/daft_lance/lance_update_column.py index 0b06409..b51e96f 100644 --- a/daft_lance/lance_update_column.py +++ b/daft_lance/lance_update_column.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast @@ -10,6 +10,7 @@ import daft.pickle from daft.datatype import DataType from daft.dependencies import pa +from daft.runners import get_or_create_runner from daft.udf import method from daft_lance._metadata import _is_lance_blob @@ -17,6 +18,9 @@ from daft_lance.namespace import DatasetOpenContext +UpdateTransform = dict[str, str] | lance.udf.BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch] + + _ROW_ADDRESS = "_rowaddr" _FRAGMENT_ID = "fragment_id" _METADATA_COLUMNS = {_ROW_ADDRESS, "_rowid", _FRAGMENT_ID} @@ -33,10 +37,10 @@ class UpdateColumnsResult: """Result of a distributed Lance column update. - ``rows_updated`` counts the rows the source submitted, not the rows Lance - matched. They differ when the source carries a ``_rowaddr`` that is not a - live row of the pinned snapshot, which is silently ignored, or the same - address more than once, which updates one row. + For ``update_columns``, ``rows_updated`` is the number of live rows matching + its predicate. For ``update_columns_df``, it counts the rows submitted by the + source, which can differ from the number Lance matches when addresses are + stale or repeated. """ version: int @@ -88,29 +92,35 @@ def validate_update_arguments(columns: Sequence[str], max_concurrency: int | Non return resolved_columns -def _validate_update_columns( - df: daft.DataFrame, +def _validate_update_target_columns( lance_ds: lance.LanceDataset, columns: Sequence[str], - max_concurrency: int | None = None, + *, + operation_name: str, ) -> list[str]: - resolved_columns = validate_update_arguments(columns, max_concurrency) - + resolved_columns = validate_update_arguments(columns, None) target_names = set(lance_ds.schema.names) for name in resolved_columns: if name not in target_names: raise ValueError( - f"Cannot update non-existent column {name!r}; update_columns_df only overwrites existing columns." + f"Cannot update non-existent column {name!r}; {operation_name} only overwrites existing columns." ) arrow_field = lance_ds.schema.field(name) if pa.types.is_struct(arrow_field.type): - # A source struct missing one of the target's fields casts cleanly - # with that field filled in as null, so a partial struct would drop - # data silently. Supporting structs needs an explicit field-set - # check first. - raise ValueError(f"Struct column {name!r} is not supported by update_columns_df.") + raise ValueError(f"Struct column {name!r} is not supported by {operation_name}.") if _is_lance_blob(arrow_field): - raise ValueError(f"Blob column {name!r} cannot be updated by update_columns_df.") + raise ValueError(f"Blob column {name!r} cannot be updated by {operation_name}.") + return resolved_columns + + +def _validate_update_columns( + df: daft.DataFrame, + lance_ds: lance.LanceDataset, + columns: Sequence[str], + max_concurrency: int | None = None, +) -> list[str]: + validate_update_arguments(columns, max_concurrency) + resolved_columns = _validate_update_target_columns(lance_ds, columns, operation_name="update_columns_df") source_names = df.column_names for required in [_ROW_ADDRESS, _FRAGMENT_ID, *resolved_columns]: @@ -224,29 +234,273 @@ def __call__(self, *series: Any) -> list[dict[str, Any]]: ] -def update_columns_from_df( - df: daft.DataFrame, +def _resolve_transform_columns( + transform: UpdateTransform, + columns: Sequence[str] | None, + max_concurrency: int | None, +) -> list[str]: + if max_concurrency is not None and max_concurrency <= 0: + raise ValueError("max_concurrency must be a positive integer.") + + inferred_columns: Sequence[str] | None = None + if isinstance(transform, dict): + for name, expression in transform.items(): + if not isinstance(name, str): + raise TypeError(f"Transform column names must be strings, got {type(name).__name__}.") + if not isinstance(expression, str): + raise TypeError(f"Transform expressions must be strings, got {type(expression).__name__}.") + inferred_columns = list(transform) + elif isinstance(transform, lance.udf.BatchUDF): + if transform.cache is not None: + raise ValueError("BatchUDF checkpoint files are not supported by update_columns fragment workers.") + if transform.output_schema is not None: + inferred_columns = transform.output_schema.names + elif not callable(transform): + raise TypeError("'transform' must be a dict of Lance SQL expressions, a BatchUDF, or a callable.") + + if columns is None: + if inferred_columns is None: + raise ValueError("'columns' is required for a callable transform without an output schema.") + return validate_update_arguments(inferred_columns, max_concurrency) + + resolved_columns = validate_update_arguments(columns, max_concurrency) + if inferred_columns is not None and set(resolved_columns) != set(inferred_columns): + raise ValueError("'columns' must name exactly the columns produced by 'transform'.") + return resolved_columns + + +def _validate_read_columns(lance_ds: lance.LanceDataset, read_columns: Sequence[str] | None) -> list[str] | None: + if read_columns is None: + return None + if isinstance(read_columns, str): + raise TypeError("'read_columns' must be a sequence of column names, not a bare string.") + + resolved = list(read_columns) + seen: set[str] = set() + valid_names = set(lance_ds.schema.names) | {_ROW_ADDRESS, "_rowid"} + for name in resolved: + if not isinstance(name, str): + raise TypeError(f"'read_columns' entries must be strings, got {type(name).__name__}.") + if name in seen: + raise ValueError(f"Duplicate column {name!r} in 'read_columns'.") + if name not in valid_names: + raise ValueError(f"Read column {name!r} does not exist in the target dataset.") + seen.add(name) + return resolved + + +def validate_transform_update_arguments( lance_ds: lance.LanceDataset, - open_context: DatasetOpenContext, + transform: UpdateTransform, *, - columns: Sequence[str], - commit_lock: Any | None = None, - max_concurrency: int | None = None, -) -> UpdateColumnsResult: - """Execute a distributed, DataFrame-driven RewriteColumns transaction.""" - resolved_columns = _validate_update_columns(df, lance_ds, columns, max_concurrency) - source = df.select(*resolved_columns, _ROW_ADDRESS, _FRAGMENT_ID) + columns: Sequence[str] | None, + read_columns: Sequence[str] | None, + where: str | None, + batch_size: int | None, + max_concurrency: int | None, +) -> tuple[list[str], list[str] | None, str | None]: + """Validate a transform-driven update before any fragment writes.""" + if batch_size is not None and batch_size <= 0: + raise ValueError("batch_size must be a positive integer.") + resolved_columns = _resolve_transform_columns(transform, columns, max_concurrency) + resolved_columns = _validate_update_target_columns(lance_ds, resolved_columns, operation_name="update_columns") + resolved_read_columns = _validate_read_columns(lance_ds, read_columns) + + if where is not None: + if not isinstance(where, str): + raise TypeError("'where' must be a Lance SQL predicate string or None.") + where = where.strip() + if not where: + raise ValueError("'where' must be a non-empty Lance SQL predicate when provided.") + + projection: dict[str, str] | list[str] + if isinstance(transform, dict): + projection = {name: transform[name] for name in resolved_columns} + else: + projection = [] + try: + lance_ds.scanner(columns=projection, filter=where, limit=1).explain_plan(True) + except Exception as exc: + raise ValueError(f"Transform or where predicate is not valid for the target Lance dataset: {exc}") from exc - handler_cls = daft.cls(_FragmentUpdateHandler, max_concurrency=max_concurrency) - handler = handler_cls(open_context, resolved_columns) - grouped = source.groupby(_FRAGMENT_ID).map_groups( - handler( - *(source[name] for name in resolved_columns), - source[_ROW_ADDRESS], - source[_FRAGMENT_ID], - ).alias("commit_message") + return resolved_columns, resolved_read_columns, where + + +def _cast_transform_output( + output: pa.RecordBatch, + row_addresses: pa.Array[Any], + *, + columns: list[str], + target_schema: pa.Schema, + expected_rows: int, +) -> pa.RecordBatch: + if not isinstance(output, pa.RecordBatch): + raise TypeError(f"Transform must return a pyarrow.RecordBatch, got {type(output).__name__}.") + if output.num_rows != expected_rows: + raise ValueError( + f"Transform changed the row count from {expected_rows} to {output.num_rows}; row-preserving output is required." + ) + if len(set(output.schema.names)) != len(output.schema.names): + raise ValueError("Transform output contains duplicate column names.") + if set(output.schema.names) != set(columns): + raise ValueError(f"Transform must return exactly the update columns {columns!r}, got {output.schema.names!r}.") + + arrays: list[pa.Array[Any]] = [] + try: + for field in target_schema: + array = output.column(output.schema.get_field_index(field.name)) + if not field.nullable and array.null_count: + raise ValueError(f"Update produced nulls for non-nullable column {field.name!r}.") + arrays.append(array.cast(field.type, safe=True)) + except (pa.ArrowInvalid, pa.ArrowNotImplementedError, pa.ArrowTypeError) as exc: + raise ValueError(f"Update columns cannot be safely cast to the target Lance schema: {exc}") from exc + + output_schema = pa.schema([*target_schema, pa.field(_ROW_ADDRESS, pa.uint64(), nullable=False)]) + return pa.RecordBatch.from_arrays([*arrays, row_addresses.cast(pa.uint64(), safe=True)], schema=output_schema) + + +def _transform_fragment( + lance_ds: lance.LanceDataset, + fragment_id: int, + *, + transform: UpdateTransform, + columns: list[str], + read_columns: list[str] | None, + where: str | None, + batch_size: int | None, +) -> dict[str, Any]: + fragment = lance_ds.get_fragment(fragment_id) + if fragment is None: + raise ValueError(f"Fragment {fragment_id} does not exist in target snapshot version {lance_ds.version}.") + + target_schema = pa.schema([lance_ds.schema.field(name) for name in columns]) + if isinstance(transform, dict): + source_batches = fragment.to_batches( + columns={name: transform[name] for name in columns}, + filter=where, + with_row_address=True, + batch_size=batch_size, + ) + + def transformed_batches() -> Iterator[pa.RecordBatch]: + for batch in source_batches: + row_addresses = batch.column(batch.schema.get_field_index(_ROW_ADDRESS)) + output = batch.select(columns) + yield _cast_transform_output( + output, + row_addresses, + columns=columns, + target_schema=target_schema, + expected_rows=batch.num_rows, + ) + + else: + requested = list(lance_ds.schema.names) if read_columns is None else read_columns + physical_columns = [name for name in requested if name not in {_ROW_ADDRESS, "_rowid"}] + source_batches = fragment.to_batches( + columns=physical_columns, + filter=where, + with_row_address=True, + with_row_id="_rowid" in requested, + batch_size=batch_size, + ) + + def transformed_batches() -> Iterator[pa.RecordBatch]: + for batch in source_batches: + row_addresses = batch.column(batch.schema.get_field_index(_ROW_ADDRESS)) + transform_input = batch.select(requested) + output = transform(transform_input) + yield _cast_transform_output( + output, + row_addresses, + columns=columns, + target_schema=target_schema, + expected_rows=batch.num_rows, + ) + + iterator = iter(transformed_batches()) + first = next(iterator, None) + if first is None: + return {"fragment_meta": None, "rows_updated": 0, "fields_modified": []} + + rows_updated = first.num_rows + + def counted_batches() -> Iterator[pa.RecordBatch]: + nonlocal rows_updated + yield first + for batch in iterator: + rows_updated += batch.num_rows + yield batch + + reader = pa.RecordBatchReader.from_batches(first.schema, counted_batches()) + fragment_meta, fields_modified = fragment.update_columns( + reader, + left_on=_ROW_ADDRESS, + right_on=_ROW_ADDRESS, ) - commit_messages = grouped.collect().to_pydict()["commit_message"] + if int(fragment_meta.id) != fragment_id: + raise ValueError(f"Fragment rewrite changed fragment id: expected {fragment_id}, got {fragment_meta.id}.") + return { + "fragment_meta": daft.pickle.dumps(fragment_meta), + "rows_updated": rows_updated, + "fields_modified": [int(field_id) for field_id in fields_modified], + } + + +class _FragmentTransformUpdateHandler: + """Apply a row-preserving transform to matching rows in pinned fragments.""" + + def __init__( + self, + open_context: DatasetOpenContext, + transform: UpdateTransform, + columns: list[str], + read_columns: list[str] | None, + where: str | None, + batch_size: int | None, + ) -> None: + self.open_context = open_context + self.transform = transform + self.columns = columns + self.read_columns = read_columns + self.where = where + self.batch_size = batch_size + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds + + @method.batch(return_dtype=_FRAGMENT_UPDATE_RESULT_DTYPE) + def __call__(self, fragment_ids: Any) -> list[dict[str, Any]]: + results = [] + for scalar in _to_arrow_array(fragment_ids): + if not scalar.is_valid: + raise ValueError("fragment_id cannot be null.") + fragment_id = int(scalar.cast(pa.int64(), safe=True).as_py()) + results.append( + _transform_fragment( + self._dataset(), + fragment_id, + transform=self.transform, + columns=self.columns, + read_columns=self.read_columns, + where=self.where, + batch_size=self.batch_size, + ) + ) + return results + + +def _commit_update_messages( + commit_messages: list[dict[str, Any]], + lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, + *, + commit_lock: Any | None, +) -> UpdateColumnsResult: + commit_messages = [message for message in commit_messages if message["fragment_meta"] is not None] if not commit_messages: return UpdateColumnsResult(version=lance_ds.version, rows_updated=0) @@ -260,7 +514,6 @@ def update_columns_from_df( if fragment_id in seen_fragment_ids: raise ValueError(f"Duplicate update result for fragment {fragment_id}.") seen_fragment_ids.add(fragment_id) - updated_fragments.append(fragment_meta) fields_modified.update(int(field_id) for field_id in message["fields_modified"]) rows_updated += int(message["rows_updated"]) @@ -280,3 +533,73 @@ def update_columns_from_df( **open_context.commit_kwargs, ) return UpdateColumnsResult(version=committed.version, rows_updated=rows_updated) + + +def update_columns_with_transform( + lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, + *, + transform: UpdateTransform, + columns: Sequence[str] | None, + read_columns: Sequence[str] | None, + where: str | None, + batch_size: int | None, + commit_lock: Any | None, + max_concurrency: int | None, +) -> UpdateColumnsResult: + resolved_columns, resolved_read_columns, resolved_where = validate_transform_update_arguments( + lance_ds, + transform, + columns=columns, + read_columns=read_columns, + where=where, + batch_size=batch_size, + max_concurrency=max_concurrency, + ) + fragment_ids = [fragment.fragment_id for fragment in lance_ds.get_fragments()] + if not fragment_ids: + return UpdateColumnsResult(version=lance_ds.version, rows_updated=0) + + # The control DataFrame contains only fragment ids. Split it without a + # hash shuffle so distributed runners can schedule fragments independently. + source = daft.from_pydict({_FRAGMENT_ID: fragment_ids}) + if len(fragment_ids) > 1 and get_or_create_runner().name != "native": + source = source.into_partitions(len(fragment_ids)) + handler_cls = daft.cls(_FragmentTransformUpdateHandler, max_concurrency=max_concurrency) + handler = handler_cls( + open_context, + transform, + resolved_columns, + resolved_read_columns, + resolved_where, + batch_size, + ) + results = source.with_column("commit_message", handler(source[_FRAGMENT_ID])) + commit_messages = results.collect().to_pydict()["commit_message"] + return _commit_update_messages(commit_messages, lance_ds, open_context, commit_lock=commit_lock) + + +def update_columns_from_df( + df: daft.DataFrame, + lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, + *, + columns: Sequence[str], + commit_lock: Any | None = None, + max_concurrency: int | None = None, +) -> UpdateColumnsResult: + """Execute a distributed, DataFrame-driven RewriteColumns transaction.""" + resolved_columns = _validate_update_columns(df, lance_ds, columns, max_concurrency) + source = df.select(*resolved_columns, _ROW_ADDRESS, _FRAGMENT_ID) + + handler_cls = daft.cls(_FragmentUpdateHandler, max_concurrency=max_concurrency) + handler = handler_cls(open_context, resolved_columns) + grouped = source.groupby(_FRAGMENT_ID).map_groups( + handler( + *(source[name] for name in resolved_columns), + source[_ROW_ADDRESS], + source[_FRAGMENT_ID], + ).alias("commit_message") + ) + commit_messages = grouped.collect().to_pydict()["commit_message"] + return _commit_update_messages(commit_messages, lance_ds, open_context, commit_lock=commit_lock) diff --git a/tests/io/lance/test_update_columns.py b/tests/io/lance/test_update_columns.py index 0aa5de7..51cf96e 100644 --- a/tests/io/lance/test_update_columns.py +++ b/tests/io/lance/test_update_columns.py @@ -31,6 +31,162 @@ def _read_update_source(path: str) -> daft.DataFrame: ) +def test_update_columns_sql_transform_with_where(tmp_path: Path) -> None: + path = str(tmp_path / "transform-where.lance") + daft.from_pydict( + {"id": list(range(6)), "value": [value * 10 for value in range(6)], "score": list(range(6))} + ).write_lance(path, max_rows_per_file=2) + before_version = lance.dataset(path).version + before_files = _fragment_files(path) + + result = daft_lance.update_columns( + path, + transform={"value": "value + 100", "score": "score * 2"}, + where="id >= 2 AND id < 4", + max_concurrency=2, + ) + + assert result == daft_lance.UpdateColumnsResult(version=before_version + 1, rows_updated=2) + assert lance.dataset(path).to_table().sort_by("id").to_pydict() == { + "id": list(range(6)), + "value": [0, 10, 120, 130, 40, 50], + "score": [0, 1, 4, 6, 4, 5], + } + after_files = _fragment_files(path) + assert {fragment_id for fragment_id in before_files if before_files[fragment_id] != after_files[fragment_id]} == {1} + + +def test_update_columns_callable_receives_only_filtered_rows(tmp_path: Path) -> None: + path = str(tmp_path / "callable-where.lance") + daft.from_pydict({"id": [1, 2, 3], "value": [10, 20, 30]}).write_lance(path) + + def transform(batch: pa.RecordBatch) -> pa.RecordBatch: + assert batch.schema.names == ["id", "value"] + assert all(value is not None and value >= 2 for value in batch.column("id").to_pylist()) + return pa.record_batch([pa.compute.multiply(batch.column("value"), 10)], names=["value"]) + + result = daft_lance.update_columns( + path, + transform=transform, + columns=["value"], + read_columns=["id", "value"], + where="id >= 2", + batch_size=1, + ) + + assert result.rows_updated == 2 + assert lance.dataset(path).to_table().sort_by("id").to_pydict() == { + "id": [1, 2, 3], + "value": [10, 200, 300], + } + + +def test_update_columns_batch_udf_infers_columns(tmp_path: Path) -> None: + path = str(tmp_path / "batch-udf.lance") + daft.from_pydict({"id": [1, 2], "value": [10, 20]}).write_lance(path) + + @lance.batch_udf( # type: ignore[no-untyped-call, untyped-decorator] + output_schema=pa.schema([pa.field("value", pa.int64())]) + ) + def transform(batch: pa.RecordBatch) -> pa.RecordBatch: + return pa.record_batch([pa.compute.add(batch.column("value"), 5)], names=["value"]) + + result = daft_lance.update_columns(path, transform=transform, read_columns=["value"]) + + assert result.rows_updated == 2 + assert lance.dataset(path).to_table().sort_by("id").to_pydict()["value"] == [15, 25] + + +def test_update_columns_no_matches_is_noop(tmp_path: Path) -> None: + path = str(tmp_path / "transform-noop.lance") + daft.from_pydict({"id": [1, 2], "value": [10, 20]}).write_lance(path, max_rows_per_file=1) + version = lance.dataset(path).version + + result = daft_lance.update_columns(path, transform={"value": "value + 1"}, where="id > 100") + + assert result == daft_lance.UpdateColumnsResult(version=version, rows_updated=0) + assert lance.dataset(path).version == version + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"transform": {"value": "value + 1"}, "where": " "}, "non-empty"), + ({"transform": {"value": "value ==== 1"}}, "not valid"), + ({"transform": lambda batch: batch}, "columns.*required"), + ({"transform": {"missing": "value + 1"}}, "non-existent"), + ({"transform": {"value": "value + 1"}, "batch_size": 0}, "batch_size"), + ({"transform": {"value": "value + 1"}, "max_concurrency": 0}, "max_concurrency"), + ], +) +def test_update_columns_validates_before_writing(tmp_path: Path, kwargs: dict[str, Any], message: str) -> None: + path = str(tmp_path / "invalid-transform.lance") + daft.from_pydict({"id": [1], "value": [10]}).write_lance(path) + version = lance.dataset(path).version + + with pytest.raises((TypeError, ValueError), match=message): + daft_lance.update_columns(path, **kwargs) + + assert lance.dataset(path).version == version + + +def test_update_columns_rejects_non_row_preserving_transform(tmp_path: Path) -> None: + path = str(tmp_path / "row-count.lance") + daft.from_pydict({"id": [1, 2], "value": [10, 20]}).write_lance(path) + version = lance.dataset(path).version + + def transform(batch: pa.RecordBatch) -> pa.RecordBatch: + return batch.select(["value"]).slice(0, max(0, batch.num_rows - 1)) + + with pytest.raises(Exception, match="row count"): + daft_lance.update_columns(path, transform=transform, columns=["value"], read_columns=["value"]) + + assert lance.dataset(path).version == version + + +def test_update_columns_rejects_stable_row_ids(tmp_path: Path) -> None: + path = str(tmp_path / "transform-stable-row-ids.lance") + lance.write_dataset(pa.table({"id": [1], "value": [10]}), path, enable_stable_row_ids=True) + version = lance.dataset(path).version + + with pytest.raises(NotImplementedError, match="does not support datasets with stable row IDs"): + daft_lance.update_columns(path, transform={"value": "value + 1"}) + + assert lance.dataset(path).version == version + + +def test_update_columns_preserves_pinned_version_for_occ(tmp_path: Path) -> None: + path = str(tmp_path / "transform-occ.lance") + daft.from_pydict({"id": [1, 2], "value": [10, 20]}).write_lance(path) + source_version = lance.dataset(path).version + lance.dataset(path).update({"value": "999"}, where="id = 1") + competing_version = lance.dataset(path).version + + with pytest.raises(Exception, match="conflict|Conflict"): + daft_lance.update_columns(path, version=source_version, transform={"value": "value + 100"}) + + assert lance.dataset(path).version == competing_version + assert lance.dataset(path).to_table().sort_by("id").to_pydict()["value"] == [999, 20] + + +def test_update_columns_namespace_roundtrip(tmp_path: Path) -> None: + namespace: dict[str, Any] = { + "namespace_impl": "dir", + "namespace_properties": {"root": str(tmp_path)}, + "table_id": ["transform_updates"], + } + daft_lance.write_lance( + daft.from_pydict({"id": [1, 2], "value": [10, 20]}), + mode="create", + **namespace, + ).collect() + + result = daft_lance.update_columns(transform={"value": "value * 10"}, where="id = 2", **namespace) + + assert result.rows_updated == 1 + assert daft_lance.read_lance(**namespace).sort("id").to_pydict()["value"] == [10, 200] + + def test_fragment_update_handler_reuses_pinned_dataset() -> None: dataset = object() From 2ebd1460c16b66e6fa6de316f20e6a31d866c1f8 Mon Sep 17 00:00:00 2001 From: fanng Date: Sat, 10 Oct 2026 08:33:50 +0900 Subject: [PATCH 2/9] fix: propagate update column offsets --- daft_lance/_lance.py | 22 ---------- daft_lance/lance_update_column.py | 39 ++++++++++++----- tests/io/lance/test_update_columns.py | 61 ++++++++++++++++++--------- 3 files changed, 71 insertions(+), 51 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 3ceaefd..3aa0f8e 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -462,9 +462,6 @@ def update_columns( A fragment failure can leave unreferenced data files written by other workers. No transaction is committed; Lance cleanup removes those files. - Raises: - NotImplementedError: If the dataset uses stable row IDs, because the - installed pylance API cannot return matched fragment offsets. """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config dataset_handle = construct_lance_dataset_handle( @@ -482,13 +479,6 @@ def update_columns( default_scan_options=default_scan_options, metadata_cache_size_bytes=metadata_cache_size_bytes, ) - if dataset_handle.dataset.has_stable_row_ids: - raise NotImplementedError( - "update_columns does not support datasets with stable row IDs: " - "pylance does not yet expose updated fragment offsets, so " - "_row_last_updated_at_version and CDF metadata cannot be updated correctly." - ) - return update_columns_with_transform( dataset_handle.dataset, dataset_handle.worker_open_context(), @@ -576,11 +566,6 @@ def update_columns_df( unreferenced files stay until Lance cleans them up (see ``LanceDataset.cleanup_old_versions``). - Raises: - NotImplementedError: If the target dataset uses stable row IDs. The - current pylance transaction binding cannot propagate the updated - fragment offsets required for correct CDF metadata. - Examples: >>> import daft >>> import daft_lance @@ -614,13 +599,6 @@ def update_columns_df( default_scan_options=default_scan_options, metadata_cache_size_bytes=metadata_cache_size_bytes, ) - if dataset_handle.dataset.has_stable_row_ids: - raise NotImplementedError( - "update_columns_df does not support datasets with stable row IDs: " - "pylance does not yet expose updated fragment offsets, so " - "_row_last_updated_at_version and CDF metadata cannot be updated correctly." - ) - return update_columns_from_df( df, dataset_handle.dataset, diff --git a/daft_lance/lance_update_column.py b/daft_lance/lance_update_column.py index b51e96f..6b323da 100644 --- a/daft_lance/lance_update_column.py +++ b/daft_lance/lance_update_column.py @@ -23,12 +23,13 @@ _ROW_ADDRESS = "_rowaddr" _FRAGMENT_ID = "fragment_id" -_METADATA_COLUMNS = {_ROW_ADDRESS, "_rowid", _FRAGMENT_ID} +_METADATA_COLUMNS = {_ROW_ADDRESS, "_rowid"} _FRAGMENT_UPDATE_RESULT_DTYPE = DataType.struct( { "fragment_meta": DataType.binary(), "rows_updated": DataType.int64(), "fields_modified": DataType.list(DataType.int64()), + "updated_offsets": DataType.binary(), } ) @@ -56,7 +57,12 @@ class _FragmentUpdateBatch: values: pa.Table -def validate_update_arguments(columns: Sequence[str], max_concurrency: int | None) -> list[str]: +def validate_update_arguments( + columns: Sequence[str], + max_concurrency: int | None, + *, + allow_fragment_id: bool = False, +) -> list[str]: """Check the arguments that do not depend on the target dataset. Called before the dataset is opened so a bad argument does not first cost a @@ -79,7 +85,7 @@ def validate_update_arguments(columns: Sequence[str], max_concurrency: int | Non if name in seen: raise ValueError(f"Duplicate column {name!r} in 'columns'.") seen.add(name) - if name == _FRAGMENT_ID: + if name == _FRAGMENT_ID and not allow_fragment_id: raise ValueError( f"Cannot update {name!r}; it is the grouping key update_columns_df injects, " "not a column of the target dataset." @@ -97,8 +103,9 @@ def _validate_update_target_columns( columns: Sequence[str], *, operation_name: str, + allow_fragment_id: bool = False, ) -> list[str]: - resolved_columns = validate_update_arguments(columns, None) + resolved_columns = validate_update_arguments(columns, None, allow_fragment_id=allow_fragment_id) target_names = set(lance_ds.schema.names) for name in resolved_columns: if name not in target_names: @@ -187,10 +194,11 @@ def _rewrite_fragment( raise ValueError(f"Update columns cannot be safely cast to the target Lance schema: {exc}") from exc update_table = values.append_column(_ROW_ADDRESS, batch.row_addresses) - fragment_meta, fields_modified = fragment.update_columns( + fragment_meta, fields_modified, updated_offsets = fragment.update_columns( update_table, left_on=_ROW_ADDRESS, right_on=_ROW_ADDRESS, + with_offsets=True, ) if int(fragment_meta.id) != batch.fragment_id: raise ValueError(f"Fragment rewrite changed fragment id: expected {batch.fragment_id}, got {fragment_meta.id}.") @@ -199,6 +207,7 @@ def _rewrite_fragment( "fragment_meta": daft.pickle.dumps(fragment_meta), "rows_updated": len(batch.row_addresses), "fields_modified": [int(field_id) for field_id in fields_modified], + "updated_offsets": updated_offsets, } @@ -261,9 +270,9 @@ def _resolve_transform_columns( if columns is None: if inferred_columns is None: raise ValueError("'columns' is required for a callable transform without an output schema.") - return validate_update_arguments(inferred_columns, max_concurrency) + return validate_update_arguments(inferred_columns, max_concurrency, allow_fragment_id=True) - resolved_columns = validate_update_arguments(columns, max_concurrency) + resolved_columns = validate_update_arguments(columns, max_concurrency, allow_fragment_id=True) if inferred_columns is not None and set(resolved_columns) != set(inferred_columns): raise ValueError("'columns' must name exactly the columns produced by 'transform'.") return resolved_columns @@ -303,7 +312,12 @@ def validate_transform_update_arguments( if batch_size is not None and batch_size <= 0: raise ValueError("batch_size must be a positive integer.") resolved_columns = _resolve_transform_columns(transform, columns, max_concurrency) - resolved_columns = _validate_update_target_columns(lance_ds, resolved_columns, operation_name="update_columns") + resolved_columns = _validate_update_target_columns( + lance_ds, + resolved_columns, + operation_name="update_columns", + allow_fragment_id=True, + ) resolved_read_columns = _validate_read_columns(lance_ds, read_columns) if where is not None: @@ -421,7 +435,7 @@ def transformed_batches() -> Iterator[pa.RecordBatch]: iterator = iter(transformed_batches()) first = next(iterator, None) if first is None: - return {"fragment_meta": None, "rows_updated": 0, "fields_modified": []} + return {"fragment_meta": None, "rows_updated": 0, "fields_modified": [], "updated_offsets": None} rows_updated = first.num_rows @@ -433,10 +447,11 @@ def counted_batches() -> Iterator[pa.RecordBatch]: yield batch reader = pa.RecordBatchReader.from_batches(first.schema, counted_batches()) - fragment_meta, fields_modified = fragment.update_columns( + fragment_meta, fields_modified, updated_offsets = fragment.update_columns( reader, left_on=_ROW_ADDRESS, right_on=_ROW_ADDRESS, + with_offsets=True, ) if int(fragment_meta.id) != fragment_id: raise ValueError(f"Fragment rewrite changed fragment id: expected {fragment_id}, got {fragment_meta.id}.") @@ -444,6 +459,7 @@ def counted_batches() -> Iterator[pa.RecordBatch]: "fragment_meta": daft.pickle.dumps(fragment_meta), "rows_updated": rows_updated, "fields_modified": [int(field_id) for field_id in fields_modified], + "updated_offsets": updated_offsets, } @@ -507,6 +523,7 @@ def _commit_update_messages( updated_fragments = [] seen_fragment_ids: set[int] = set() fields_modified: set[int] = set() + updated_fragment_offsets: dict[int, bytes] = {} rows_updated = 0 for message in commit_messages: fragment_meta = daft.pickle.loads(message["fragment_meta"]) @@ -516,6 +533,7 @@ def _commit_update_messages( seen_fragment_ids.add(fragment_id) updated_fragments.append(fragment_meta) fields_modified.update(int(field_id) for field_id in message["fields_modified"]) + updated_fragment_offsets[fragment_id] = message["updated_offsets"] rows_updated += int(message["rows_updated"]) operation = lance.LanceOperation.Update( @@ -523,6 +541,7 @@ def _commit_update_messages( fields_modified=sorted(fields_modified), fields_for_preserving_frag_bitmap=[], update_mode="rewrite_columns", + updated_fragment_offsets=updated_fragment_offsets, ) committed = lance.LanceDataset.commit( open_context.uri, diff --git a/tests/io/lance/test_update_columns.py b/tests/io/lance/test_update_columns.py index 51cf96e..e257c0a 100644 --- a/tests/io/lance/test_update_columns.py +++ b/tests/io/lance/test_update_columns.py @@ -85,9 +85,7 @@ def test_update_columns_batch_udf_infers_columns(tmp_path: Path) -> None: path = str(tmp_path / "batch-udf.lance") daft.from_pydict({"id": [1, 2], "value": [10, 20]}).write_lance(path) - @lance.batch_udf( # type: ignore[no-untyped-call, untyped-decorator] - output_schema=pa.schema([pa.field("value", pa.int64())]) - ) + @lance.batch_udf(output_schema=pa.schema([pa.field("value", pa.int64())])) def transform(batch: pa.RecordBatch) -> pa.RecordBatch: return pa.record_batch([pa.compute.add(batch.column("value"), 5)], names=["value"]) @@ -144,15 +142,42 @@ def transform(batch: pa.RecordBatch) -> pa.RecordBatch: assert lance.dataset(path).version == version -def test_update_columns_rejects_stable_row_ids(tmp_path: Path) -> None: +def test_update_columns_updates_stable_row_id_metadata(tmp_path: Path) -> None: path = str(tmp_path / "transform-stable-row-ids.lance") - lance.write_dataset(pa.table({"id": [1], "value": [10]}), path, enable_stable_row_ids=True) - version = lance.dataset(path).version + lance.write_dataset( + pa.table({"id": [1, 2, 3], "value": [10, 20, 30]}), + path, + enable_stable_row_ids=True, + ) + base_version = lance.dataset(path).version - with pytest.raises(NotImplementedError, match="does not support datasets with stable row IDs"): - daft_lance.update_columns(path, transform={"value": "value + 1"}) + result = daft_lance.update_columns( + path, + transform={"value": "value + 100"}, + where="id IN (1, 3)", + ) + + assert result.rows_updated == 2 + table = lance.dataset(path).to_table(columns=["id", "value", "_row_last_updated_at_version"]).sort_by("id") + assert table.column("value").to_pylist() == [110, 20, 130] + assert table.column("_row_last_updated_at_version").to_pylist() == [ + result.version, + base_version, + result.version, + ] - assert lance.dataset(path).version == version + +def test_update_columns_can_update_business_fragment_id_column(tmp_path: Path) -> None: + path = str(tmp_path / "business-fragment-id.lance") + lance.write_dataset(pa.table({"id": [1, 2], "fragment_id": [10, 20]}), path) + + result = daft_lance.update_columns( + path, + transform={"fragment_id": "fragment_id + 1"}, + ) + + assert result.rows_updated == 2 + assert lance.dataset(path).to_table().sort_by("id").column("fragment_id").to_pylist() == [11, 21] def test_update_columns_preserves_pinned_version_for_occ(tmp_path: Path) -> None: @@ -414,24 +439,22 @@ def test_update_columns_df_safe_casts_to_target_type(tmp_path: Path) -> None: assert table.column("value").to_pylist() == [123, 123] -def test_update_columns_df_rejects_stable_row_ids_before_writing(tmp_path: Path) -> None: +def test_update_columns_df_updates_stable_row_id_metadata(tmp_path: Path) -> None: path = str(tmp_path / "stable-row-ids.lance") lance.write_dataset( pa.table({"id": [1, 2], "value": [10, 20]}), path, enable_stable_row_ids=True, ) - source = _read_update_source(path).with_column("value", daft.col("value") + 1) - before = lance.dataset(path) - before_version = before.version - before_values = before.to_table().column("value").to_pylist() + source = _read_update_source(path).where("id = 1").with_column("value", daft.col("value") + 1) + base_version = lance.dataset(path).version - with pytest.raises(NotImplementedError, match="does not support datasets with stable row IDs"): - daft_lance.update_columns_df(source, path, columns=["value"]) + result = daft_lance.update_columns_df(source, path, columns=["value"]) - after = lance.dataset(path) - assert after.version == before_version - assert after.to_table().column("value").to_pylist() == before_values + assert result.rows_updated == 1 + table = lance.dataset(path).to_table(columns=["id", "value", "_row_last_updated_at_version"]).sort_by("id") + assert table.column("value").to_pylist() == [11, 20] + assert table.column("_row_last_updated_at_version").to_pylist() == [result.version, base_version] def test_update_columns_df_uses_commit_lock(tmp_path: Path) -> None: From fadbd3a841b824b1ebce1ad08e75c601bf1b92e6 Mon Sep 17 00:00:00 2001 From: fanng Date: Sat, 10 Oct 2026 09:39:12 +0900 Subject: [PATCH 3/9] refactor: inline update transform type --- daft_lance/lance_update_column.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/daft_lance/lance_update_column.py b/daft_lance/lance_update_column.py index 6b323da..7cb18cf 100644 --- a/daft_lance/lance_update_column.py +++ b/daft_lance/lance_update_column.py @@ -18,9 +18,6 @@ from daft_lance.namespace import DatasetOpenContext -UpdateTransform = dict[str, str] | lance.udf.BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch] - - _ROW_ADDRESS = "_rowaddr" _FRAGMENT_ID = "fragment_id" _METADATA_COLUMNS = {_ROW_ADDRESS, "_rowid"} @@ -244,7 +241,7 @@ def __call__(self, *series: Any) -> list[dict[str, Any]]: def _resolve_transform_columns( - transform: UpdateTransform, + transform: dict[str, str] | lance.udf.BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch], columns: Sequence[str] | None, max_concurrency: int | None, ) -> list[str]: @@ -300,7 +297,7 @@ def _validate_read_columns(lance_ds: lance.LanceDataset, read_columns: Sequence[ def validate_transform_update_arguments( lance_ds: lance.LanceDataset, - transform: UpdateTransform, + transform: dict[str, str] | lance.udf.BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch], *, columns: Sequence[str] | None, read_columns: Sequence[str] | None, @@ -377,7 +374,7 @@ def _transform_fragment( lance_ds: lance.LanceDataset, fragment_id: int, *, - transform: UpdateTransform, + transform: dict[str, str] | lance.udf.BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch], columns: list[str], read_columns: list[str] | None, where: str | None, @@ -469,7 +466,7 @@ class _FragmentTransformUpdateHandler: def __init__( self, open_context: DatasetOpenContext, - transform: UpdateTransform, + transform: dict[str, str] | lance.udf.BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch], columns: list[str], read_columns: list[str] | None, where: str | None, @@ -558,7 +555,7 @@ def update_columns_with_transform( lance_ds: lance.LanceDataset, open_context: DatasetOpenContext, *, - transform: UpdateTransform, + transform: dict[str, str] | lance.udf.BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch], columns: Sequence[str] | None, read_columns: Sequence[str] | None, where: str | None, From 088ca51e17f09ee87d38586bb2284e87e814c258 Mon Sep 17 00:00:00 2001 From: fanng Date: Sat, 10 Oct 2026 09:46:26 +0900 Subject: [PATCH 4/9] refactor: remove update snapshot parameters --- daft_lance/_lance.py | 21 ++------------------- tests/io/lance/test_update_columns.py | 20 +++++++++++++------- 2 files changed, 15 insertions(+), 26 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 3aa0f8e..ceabd7f 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -410,8 +410,6 @@ def update_columns( namespace_impl: str | None = None, namespace_properties: dict[str, str] | None = None, storage_options: dict[str, Any] | None = None, - version: int | str | None = None, - asof: str | None = None, block_size: int | None = None, commit_lock: Any | None = None, index_cache_size: int | None = None, @@ -444,9 +442,6 @@ def update_columns( namespace_impl: Lance Namespace implementation. namespace_properties: Properties used to connect to the namespace. storage_options: Additional object-store options. - version: Dataset version or tag to read and commit against. Defaults to - the version current when this function starts. - asof: Find the latest version created on or before this timestamp. block_size: Hint for the minimum object-store I/O request size. commit_lock: Custom Lance commit lock. index_cache_size: Index cache size. @@ -471,8 +466,6 @@ def update_columns( namespace_impl=namespace_impl, namespace_properties=namespace_properties, table_id=table_id, - version=version, - asof=asof, block_size=block_size, commit_lock=commit_lock, index_cache_size=index_cache_size, @@ -503,8 +496,6 @@ def update_columns_df( namespace_impl: str | None = None, namespace_properties: dict[str, str] | None = None, storage_options: dict[str, Any] | None = None, - version: int | str | None = None, - asof: str | None = None, block_size: int | None = None, commit_lock: Any | None = None, index_cache_size: int | None = None, @@ -536,12 +527,6 @@ def update_columns_df( namespace_impl: Lance Namespace implementation. namespace_properties: Properties used to connect to the namespace. storage_options: Additional object-store options. - version: Dataset version or tag to read and commit against. Defaults to - the current version. It must be the snapshot the source ``_rowaddr`` - values were read from. Lance rebases the commit over concurrent - transactions that touch other fragments; a concurrent write to a - fragment this update rewrites raises ``CommitConflictError``. - asof: If specified, find the latest version created on or earlier than the given argument value. block_size: Block size in bytes. Provide a hint for the size of the minimal I/O request. commit_lock: Custom Lance commit lock. index_cache_size: Index cache size. @@ -556,8 +541,8 @@ def update_columns_df( of the pinned snapshot updates nothing, and a ``_rowaddr`` repeated in the source updates its row once with an unspecified one of the submitted values. Both are silent and both are still counted. Keep the - source aligned with ``version``, and unique, to keep the count - meaningful. + source aligned with the current target snapshot, and unique, to keep + the count meaningful. Note: Fragments are rewritten in parallel, so a failure raised by one @@ -591,8 +576,6 @@ def update_columns_df( namespace_impl=namespace_impl, namespace_properties=namespace_properties, table_id=table_id, - version=version, - asof=asof, block_size=block_size, commit_lock=commit_lock, index_cache_size=index_cache_size, diff --git a/tests/io/lance/test_update_columns.py b/tests/io/lance/test_update_columns.py index e257c0a..3e7fb14 100644 --- a/tests/io/lance/test_update_columns.py +++ b/tests/io/lance/test_update_columns.py @@ -1,7 +1,8 @@ from __future__ import annotations import contextlib -from collections.abc import Iterator +import inspect +from collections.abc import Callable, Iterator from pathlib import Path from typing import Any, cast @@ -31,6 +32,13 @@ def _read_update_source(path: str) -> daft.DataFrame: ) +def test_update_columns_do_not_accept_historical_snapshot_parameters() -> None: + for operation in [daft_lance.update_columns, daft_lance.update_columns_df]: + parameters = inspect.signature(cast(Callable[..., Any], operation)).parameters + assert "version" not in parameters + assert "asof" not in parameters + + def test_update_columns_sql_transform_with_where(tmp_path: Path) -> None: path = str(tmp_path / "transform-where.lance") daft.from_pydict( @@ -180,18 +188,16 @@ def test_update_columns_can_update_business_fragment_id_column(tmp_path: Path) - assert lance.dataset(path).to_table().sort_by("id").column("fragment_id").to_pylist() == [11, 21] -def test_update_columns_preserves_pinned_version_for_occ(tmp_path: Path) -> None: +def test_update_columns_uses_current_snapshot(tmp_path: Path) -> None: path = str(tmp_path / "transform-occ.lance") daft.from_pydict({"id": [1, 2], "value": [10, 20]}).write_lance(path) - source_version = lance.dataset(path).version lance.dataset(path).update({"value": "999"}, where="id = 1") competing_version = lance.dataset(path).version - with pytest.raises(Exception, match="conflict|Conflict"): - daft_lance.update_columns(path, version=source_version, transform={"value": "value + 100"}) + result = daft_lance.update_columns(path, transform={"value": "value + 100"}) - assert lance.dataset(path).version == competing_version - assert lance.dataset(path).to_table().sort_by("id").to_pydict()["value"] == [999, 20] + assert result == daft_lance.UpdateColumnsResult(version=competing_version + 1, rows_updated=2) + assert lance.dataset(path).to_table().sort_by("id").to_pydict()["value"] == [1099, 120] def test_update_columns_namespace_roundtrip(tmp_path: Path) -> None: From b9ca2dd09a7858f5e9728441a737714ce6df85c8 Mon Sep 17 00:00:00 2001 From: fanng Date: Sat, 10 Oct 2026 09:51:02 +0900 Subject: [PATCH 5/9] docs: expand update transform examples --- README.md | 59 +++++++++++++++++++++++++++++++++++--------- daft_lance/_lance.py | 36 +++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 90cd882..3f7060e 100644 --- a/README.md +++ b/README.md @@ -235,9 +235,46 @@ result = update_columns( print(result.version, result.rows_updated) ``` -Python callables receive Arrow record batches after filtering and must return -exactly the existing columns named by `columns`, with the same row count and -order. Use `read_columns` to restrict their input. +`BatchUDF` transforms infer update columns from `output_schema`; use +`read_columns` to restrict the source columns supplied to the UDF: + +```python +import lance +import pyarrow as pa +from daft_lance import update_columns + +@lance.batch_udf(output_schema=pa.schema([pa.field("score", pa.int64())])) +def double_score(batch: pa.RecordBatch) -> pa.RecordBatch: + return pa.record_batch([pa.compute.multiply(batch.column("score"), 2)], names=["score"]) + +update_columns( + "s3://bucket/my_dataset", + transform=double_score, + where="status = 'new'", + read_columns=["score"], +) +``` + +Regular Python callables receive Arrow record batches after filtering. They +must declare `columns`, and return exactly those existing columns with the same +row count and order: + +```python +import pyarrow as pa +from daft_lance import update_columns + +def normalize_score(batch: pa.RecordBatch) -> pa.RecordBatch: + score = batch.column("score") + return pa.record_batch([pa.compute.divide(score, 100)], names=["score"]) + +update_columns( + "s3://bucket/my_dataset", + transform=normalize_score, + columns=["score"], + read_columns=["score"], + batch_size=8192, +) +``` Use a prepared Daft DataFrame to overwrite existing Lance columns while preserving row addresses and untouched column files: @@ -272,14 +309,14 @@ left-outer join on `_rowaddr` and both failures are expressible in it: an address matching no live row updates nothing, and a repeated address updates its row once with one of the submitted values, chosen by row order rather than by any rule. Both are silent, and both still count towards `rows_updated`. -Read the source from the snapshot you are updating (pass the same `version` if -you pin one) rather than replaying an address list produced against an older -snapshot, and make sure an upstream join cannot fan a row address out. The -update itself is committed atomically using Lance `RewriteColumns`. - -Stable-row-ID datasets are rejected before any fragments are written because -current pylance bindings cannot propagate the offsets required to advance -`_row_last_updated_at_version` and keep CDF metadata correct. +Read the source immediately before updating rather than replaying an address +list produced against an older snapshot, and make sure an upstream join cannot +fan a row address out. The update itself is committed atomically using Lance +`RewriteColumns`. + +Stable-row-ID datasets are supported: each fragment rewrite propagates its +matched physical offsets so Lance updates `_row_last_updated_at_version` and +CDF metadata only for rows that actually changed. Fragments are rewritten in parallel, so a per-fragment failure can surface after other fragments have already written their new column files. Nothing is diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index ceabd7f..f852094 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -457,6 +457,42 @@ def update_columns( A fragment failure can leave unreferenced data files written by other workers. No transaction is committed; Lance cleanup removes those files. + Examples: + Update columns with Lance SQL expressions. Dict keys infer ``columns``: + + >>> import daft_lance # doctest: +SKIP + >>> daft_lance.update_columns( # doctest: +SKIP + ... "/tmp/events.lance", + ... transform={"score": "score * 2", "status": "'processed'"}, + ... where="date >= DATE '2026-07-01'", + ... ) + + Use a ``BatchUDF`` when its output schema is known. Its output schema + infers ``columns`` and ``read_columns`` limits data provided to the UDF: + + >>> import lance # doctest: +SKIP + >>> import pyarrow as pa # doctest: +SKIP + >>> @lance.batch_udf(output_schema=pa.schema([pa.field("score", pa.int64())])) # doctest: +SKIP + ... def double_score(batch): + ... return pa.record_batch([pa.compute.multiply(batch.column("score"), 2)], names=["score"]) + >>> daft_lance.update_columns( # doctest: +SKIP + ... "/tmp/events.lance", transform=double_score, read_columns=["score"] + ... ) + + A regular callable must declare ``columns`` and return exactly those + columns, preserving the input row count and order: + + >>> def normalize_score(batch): # doctest: +SKIP + ... score = batch.column("score") + ... return pa.record_batch([pa.compute.divide(score, 100)], names=["score"]) + >>> daft_lance.update_columns( # doctest: +SKIP + ... "/tmp/events.lance", + ... transform=normalize_score, + ... columns=["score"], + ... read_columns=["score"], + ... batch_size=8192, + ... ) + """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config dataset_handle = construct_lance_dataset_handle( From baabd114db3f6f235baa60488585565df0f40aec Mon Sep 17 00:00:00 2001 From: fanng Date: Sat, 10 Oct 2026 10:01:29 +0900 Subject: [PATCH 6/9] docs: derive update values from source columns --- README.md | 23 ++++++++++++----------- daft_lance/_lance.py | 21 +++++++++++---------- 2 files changed, 23 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 3f7060e..005e409 100644 --- a/README.md +++ b/README.md @@ -229,7 +229,7 @@ from daft_lance import update_columns result = update_columns( "s3://bucket/my_dataset", - transform={"score": "score * 2", "status": "'processed'"}, + transform={"score": "raw_score * weight", "status": "'processed'"}, where="date >= DATE '2026-07-01'", ) print(result.version, result.rows_updated) @@ -244,14 +244,15 @@ import pyarrow as pa from daft_lance import update_columns @lance.batch_udf(output_schema=pa.schema([pa.field("score", pa.int64())])) -def double_score(batch: pa.RecordBatch) -> pa.RecordBatch: - return pa.record_batch([pa.compute.multiply(batch.column("score"), 2)], names=["score"]) +def compute_score(batch: pa.RecordBatch) -> pa.RecordBatch: + score = pa.compute.multiply(batch.column("raw_score"), batch.column("multiplier")) + return pa.record_batch([score], names=["score"]) update_columns( "s3://bucket/my_dataset", - transform=double_score, + transform=compute_score, where="status = 'new'", - read_columns=["score"], + read_columns=["raw_score", "multiplier"], ) ``` @@ -263,15 +264,15 @@ row count and order: import pyarrow as pa from daft_lance import update_columns -def normalize_score(batch: pa.RecordBatch) -> pa.RecordBatch: - score = batch.column("score") - return pa.record_batch([pa.compute.divide(score, 100)], names=["score"]) +def compute_average(batch: pa.RecordBatch) -> pa.RecordBatch: + average = pa.compute.divide(batch.column("total"), batch.column("count")) + return pa.record_batch([average], names=["average"]) update_columns( "s3://bucket/my_dataset", - transform=normalize_score, - columns=["score"], - read_columns=["score"], + transform=compute_average, + columns=["average"], + read_columns=["total", "count"], batch_size=8192, ) ``` diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index f852094..90b36c7 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -463,7 +463,7 @@ def update_columns( >>> import daft_lance # doctest: +SKIP >>> daft_lance.update_columns( # doctest: +SKIP ... "/tmp/events.lance", - ... transform={"score": "score * 2", "status": "'processed'"}, + ... transform={"score": "raw_score * weight", "status": "'processed'"}, ... where="date >= DATE '2026-07-01'", ... ) @@ -473,23 +473,24 @@ def update_columns( >>> import lance # doctest: +SKIP >>> import pyarrow as pa # doctest: +SKIP >>> @lance.batch_udf(output_schema=pa.schema([pa.field("score", pa.int64())])) # doctest: +SKIP - ... def double_score(batch): - ... return pa.record_batch([pa.compute.multiply(batch.column("score"), 2)], names=["score"]) + ... def compute_score(batch): + ... score = pa.compute.multiply(batch.column("raw_score"), batch.column("multiplier")) + ... return pa.record_batch([score], names=["score"]) >>> daft_lance.update_columns( # doctest: +SKIP - ... "/tmp/events.lance", transform=double_score, read_columns=["score"] + ... "/tmp/events.lance", transform=compute_score, read_columns=["raw_score", "multiplier"] ... ) A regular callable must declare ``columns`` and return exactly those columns, preserving the input row count and order: - >>> def normalize_score(batch): # doctest: +SKIP - ... score = batch.column("score") - ... return pa.record_batch([pa.compute.divide(score, 100)], names=["score"]) + >>> def compute_average(batch): # doctest: +SKIP + ... average = pa.compute.divide(batch.column("total"), batch.column("count")) + ... return pa.record_batch([average], names=["average"]) >>> daft_lance.update_columns( # doctest: +SKIP ... "/tmp/events.lance", - ... transform=normalize_score, - ... columns=["score"], - ... read_columns=["score"], + ... transform=compute_average, + ... columns=["average"], + ... read_columns=["total", "count"], ... batch_size=8192, ... ) From f01fad2ee06c2adfbab29efe8aea96073a4c0c70 Mon Sep 17 00:00:00 2001 From: fanng Date: Sat, 10 Oct 2026 11:25:13 +0900 Subject: [PATCH 7/9] refactor: remove update scan defaults --- daft_lance/_lance.py | 6 ------ tests/io/lance/test_update_columns.py | 1 + 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 90b36c7..c188859 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -413,7 +413,6 @@ def update_columns( block_size: int | None = None, commit_lock: Any | None = None, index_cache_size: int | None = None, - default_scan_options: dict[str, Any] | None = None, metadata_cache_size_bytes: int | None = None, max_concurrency: int | None = None, ) -> UpdateColumnsResult: @@ -445,7 +444,6 @@ def update_columns( block_size: Hint for the minimum object-store I/O request size. commit_lock: Custom Lance commit lock. index_cache_size: Index cache size. - default_scan_options: Default Lance scan options. metadata_cache_size_bytes: Metadata cache size in bytes. max_concurrency: Maximum number of concurrent fragment-update workers. @@ -506,7 +504,6 @@ def update_columns( block_size=block_size, commit_lock=commit_lock, index_cache_size=index_cache_size, - default_scan_options=default_scan_options, metadata_cache_size_bytes=metadata_cache_size_bytes, ) return update_columns_with_transform( @@ -536,7 +533,6 @@ def update_columns_df( block_size: int | None = None, commit_lock: Any | None = None, index_cache_size: int | None = None, - default_scan_options: dict[str, Any] | None = None, metadata_cache_size_bytes: int | None = None, max_concurrency: int | None = None, ) -> UpdateColumnsResult: @@ -567,7 +563,6 @@ def update_columns_df( block_size: Block size in bytes. Provide a hint for the size of the minimal I/O request. commit_lock: Custom Lance commit lock. index_cache_size: Index cache size. - default_scan_options: Default scan options. metadata_cache_size_bytes: Size of the metadata cache in bytes. max_concurrency: Maximum number of concurrent fragment-update workers. @@ -616,7 +611,6 @@ def update_columns_df( block_size=block_size, commit_lock=commit_lock, index_cache_size=index_cache_size, - default_scan_options=default_scan_options, metadata_cache_size_bytes=metadata_cache_size_bytes, ) return update_columns_from_df( diff --git a/tests/io/lance/test_update_columns.py b/tests/io/lance/test_update_columns.py index 3e7fb14..51b6c93 100644 --- a/tests/io/lance/test_update_columns.py +++ b/tests/io/lance/test_update_columns.py @@ -37,6 +37,7 @@ def test_update_columns_do_not_accept_historical_snapshot_parameters() -> None: parameters = inspect.signature(cast(Callable[..., Any], operation)).parameters assert "version" not in parameters assert "asof" not in parameters + assert "default_scan_options" not in parameters def test_update_columns_sql_transform_with_where(tmp_path: Path) -> None: From b345174663bc7858a42a565cd9493c3c07e36400 Mon Sep 17 00:00:00 2001 From: fanng Date: Sat, 10 Oct 2026 12:36:48 +0900 Subject: [PATCH 8/9] feat: add update transform resource options --- README.md | 23 ++++++++++++ daft_lance/_lance.py | 31 +++++++++++++++++ daft_lance/lance_update_column.py | 37 ++++++++++++++++++-- tests/io/lance/test_update_columns.py | 50 ++++++++++++++++++++++++++- 4 files changed, 138 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 005e409..61db88a 100644 --- a/README.md +++ b/README.md @@ -277,6 +277,29 @@ update_columns( ) ``` +For a GPU-backed version of the callable above, request resources for each +Daft transform worker. The actual worker count remains bounded by Ray's +available resources. Include custom `ray_options` only when the Ray cluster +advertises the named resource: + +```python +update_columns( + "s3://bucket/my_dataset", + transform=compute_average, + columns=["average"], + read_columns=["total", "count"], + gpus=1, + cpus=2, + use_process=True, + max_concurrency=8, + ray_options={"resources": {"gpu_type_a10": 0.001}}, +) +``` + +`cpus=None`, `gpus=0`, `use_process=None`, and `ray_options=None` preserve +Daft's defaults. These options apply only to transform workers; they are not +needed by `update_columns_df`, whose input values have already been computed. + Use a prepared Daft DataFrame to overwrite existing Lance columns while preserving row addresses and untouched column files: diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index c188859..45ea120 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -415,6 +415,10 @@ def update_columns( index_cache_size: int | None = None, metadata_cache_size_bytes: int | None = None, max_concurrency: int | None = None, + cpus: float | None = None, + gpus: float = 0, + use_process: bool | None = None, + ray_options: dict[str, Any] | None = None, ) -> UpdateColumnsResult: """Overwrite existing Lance columns with a row-preserving transform. @@ -446,6 +450,14 @@ def update_columns( index_cache_size: Index cache size. metadata_cache_size_bytes: Metadata cache size in bytes. max_concurrency: Maximum number of concurrent fragment-update workers. + cpus: CPU resources requested for each transform worker. ``None`` lets + Daft choose. + gpus: GPU resources requested for each transform worker. Defaults to + ``0`` and is passed to Daft unchanged. + use_process: Whether Daft runs each transform worker in a separate + process. ``None`` lets Daft choose. + ray_options: Additional Ray options for transform workers. Daft + validates these options and runner compatibility. Returns: The committed dataset version and the number of live rows that matched @@ -492,6 +504,21 @@ def update_columns( ... batch_size=8192, ... ) + A GPU-backed version of the callable above can request resources for + each Daft worker. The actual worker count is still bounded by available + Ray resources: + + >>> daft_lance.update_columns( # doctest: +SKIP + ... "/tmp/events.lance", + ... transform=compute_average, + ... columns=["average"], + ... read_columns=["total", "count"], + ... gpus=1, + ... cpus=2, + ... use_process=True, + ... max_concurrency=8, + ... ) + """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config dataset_handle = construct_lance_dataset_handle( @@ -516,6 +543,10 @@ def update_columns( batch_size=batch_size, commit_lock=commit_lock, max_concurrency=max_concurrency, + cpus=cpus, + gpus=gpus, + use_process=use_process, + ray_options=ray_options, ) diff --git a/daft_lance/lance_update_column.py b/daft_lance/lance_update_column.py index 7cb18cf..cc8f956 100644 --- a/daft_lance/lance_update_column.py +++ b/daft_lance/lance_update_column.py @@ -506,6 +506,31 @@ def __call__(self, fragment_ids: Any) -> list[dict[str, Any]]: return results +def _make_transform_update_handler( + open_context: DatasetOpenContext, + transform: dict[str, str] | lance.udf.BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch], + columns: list[str], + read_columns: list[str] | None, + where: str | None, + batch_size: int | None, + *, + cpus: float | None, + gpus: float, + use_process: bool | None, + max_concurrency: int | None, + ray_options: dict[str, Any] | None, +) -> Any: + handler_cls = daft.cls( + _FragmentTransformUpdateHandler, + cpus=cpus, + gpus=gpus, + use_process=use_process, + max_concurrency=max_concurrency, + ray_options=ray_options, + ) + return handler_cls(open_context, transform, columns, read_columns, where, batch_size) + + def _commit_update_messages( commit_messages: list[dict[str, Any]], lance_ds: lance.LanceDataset, @@ -562,6 +587,10 @@ def update_columns_with_transform( batch_size: int | None, commit_lock: Any | None, max_concurrency: int | None, + cpus: float | None, + gpus: float, + use_process: bool | None, + ray_options: dict[str, Any] | None, ) -> UpdateColumnsResult: resolved_columns, resolved_read_columns, resolved_where = validate_transform_update_arguments( lance_ds, @@ -581,14 +610,18 @@ def update_columns_with_transform( source = daft.from_pydict({_FRAGMENT_ID: fragment_ids}) if len(fragment_ids) > 1 and get_or_create_runner().name != "native": source = source.into_partitions(len(fragment_ids)) - handler_cls = daft.cls(_FragmentTransformUpdateHandler, max_concurrency=max_concurrency) - handler = handler_cls( + handler = _make_transform_update_handler( open_context, transform, resolved_columns, resolved_read_columns, resolved_where, batch_size, + cpus=cpus, + gpus=gpus, + use_process=use_process, + max_concurrency=max_concurrency, + ray_options=ray_options, ) results = source.with_column("commit_message", handler(source[_FRAGMENT_ID])) commit_messages = results.collect().to_pydict()["commit_message"] diff --git a/tests/io/lance/test_update_columns.py b/tests/io/lance/test_update_columns.py index 51b6c93..6b1293f 100644 --- a/tests/io/lance/test_update_columns.py +++ b/tests/io/lance/test_update_columns.py @@ -11,6 +11,7 @@ import daft import daft_lance +import daft_lance.lance_update_column as update_column from daft.dependencies import pa from daft_lance.lance_update_column import _FragmentUpdateHandler @@ -32,13 +33,60 @@ def _read_update_source(path: str) -> daft.DataFrame: ) -def test_update_columns_do_not_accept_historical_snapshot_parameters() -> None: +def test_update_columns_public_api_parameters() -> None: for operation in [daft_lance.update_columns, daft_lance.update_columns_df]: parameters = inspect.signature(cast(Callable[..., Any], operation)).parameters assert "version" not in parameters assert "asof" not in parameters assert "default_scan_options" not in parameters + transform_parameters = inspect.signature(daft_lance.update_columns).parameters + assert {"cpus", "gpus", "use_process", "ray_options"} <= set(transform_parameters) + assert transform_parameters["cpus"].default is None + assert transform_parameters["gpus"].default == 0 + assert transform_parameters["use_process"].default is None + assert transform_parameters["ray_options"].default is None + + dataframe_parameters = inspect.signature(daft_lance.update_columns_df).parameters + assert {"cpus", "gpus", "use_process", "ray_options"}.isdisjoint(dataframe_parameters) + + +def test_transform_handler_forwards_resource_options(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + open_context = object() + transform = {"value": "value + 1"} + + def fake_cls(class_: type, **kwargs: Any) -> Callable[..., dict[str, Any]]: + captured["class"] = class_ + captured.update(kwargs) + return lambda *args: {"args": args} + + monkeypatch.setattr(daft, "cls", fake_cls) + + handler = update_column._make_transform_update_handler( + cast(Any, open_context), + transform, + ["value"], + ["value"], + "id > 0", + 128, + cpus=2, + gpus=1, + use_process=True, + max_concurrency=4, + ray_options={"resources": {"gpu_type_a10": 0.001}}, + ) + + assert captured == { + "class": update_column._FragmentTransformUpdateHandler, + "cpus": 2, + "gpus": 1, + "use_process": True, + "max_concurrency": 4, + "ray_options": {"resources": {"gpu_type_a10": 0.001}}, + } + assert handler["args"] == (open_context, transform, ["value"], ["value"], "id > 0", 128) + def test_update_columns_sql_transform_with_where(tmp_path: Path) -> None: path = str(tmp_path / "transform-where.lance") From d88c18912b1a6ec963e4ea945bb3fc0196085b44 Mon Sep 17 00:00:00 2001 From: fanng Date: Sat, 10 Oct 2026 15:57:00 +0900 Subject: [PATCH 9/9] refactor: drop update transform dry run --- daft_lance/lance_update_column.py | 10 ---------- tests/io/lance/test_update_columns.py | 2 +- 2 files changed, 1 insertion(+), 11 deletions(-) diff --git a/daft_lance/lance_update_column.py b/daft_lance/lance_update_column.py index cc8f956..b684ef8 100644 --- a/daft_lance/lance_update_column.py +++ b/daft_lance/lance_update_column.py @@ -324,16 +324,6 @@ def validate_transform_update_arguments( if not where: raise ValueError("'where' must be a non-empty Lance SQL predicate when provided.") - projection: dict[str, str] | list[str] - if isinstance(transform, dict): - projection = {name: transform[name] for name in resolved_columns} - else: - projection = [] - try: - lance_ds.scanner(columns=projection, filter=where, limit=1).explain_plan(True) - except Exception as exc: - raise ValueError(f"Transform or where predicate is not valid for the target Lance dataset: {exc}") from exc - return resolved_columns, resolved_read_columns, where diff --git a/tests/io/lance/test_update_columns.py b/tests/io/lance/test_update_columns.py index 6b1293f..89a68e7 100644 --- a/tests/io/lance/test_update_columns.py +++ b/tests/io/lance/test_update_columns.py @@ -167,7 +167,7 @@ def test_update_columns_no_matches_is_noop(tmp_path: Path) -> None: ("kwargs", "message"), [ ({"transform": {"value": "value + 1"}, "where": " "}, "non-empty"), - ({"transform": {"value": "value ==== 1"}}, "not valid"), + ({"transform": {"value": "value ==== 1"}}, "sql parser error"), ({"transform": lambda batch: batch}, "columns.*required"), ({"transform": {"missing": "value + 1"}}, "non-existent"), ({"transform": {"value": "value + 1"}, "batch_size": 0}, "batch_size"),