From 3ec287d0dfab7cc82e4ab45796205556891c49de Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Mon, 27 Jul 2026 11:12:44 +0900 Subject: [PATCH 01/11] feat: add distributed update columns API --- lance_ray/__init__.py | 8 + lance_ray/datasource.py | 73 +-- lance_ray/io.py | 749 +++++++++++++++++++++++- tests/test_update_columns.py | 1051 ++++++++++++++++++++++++++++++++++ 4 files changed, 1844 insertions(+), 37 deletions(-) create mode 100644 tests/test_update_columns.py diff --git a/lance_ray/__init__.py b/lance_ray/__init__.py index cea756a3..a4a77721 100644 --- a/lance_ray/__init__.py +++ b/lance_ray/__init__.py @@ -17,10 +17,14 @@ from .fragment import LanceFragmentWriter from .index import create_index, create_scalar_index, optimize_indices from .io import ( + CommitOutcomeUnknown, + UpdateColumnsResult, + UpdateColumnsTransform, add_columns, add_columns_from, merge_columns_from, read_lance, + update_columns, write_lance, ) from .pool import clear_global_pool, get_global_pool, init_global_pool, set_global_pool @@ -37,6 +41,10 @@ "add_columns", "add_columns_from", "merge_columns_from", + "update_columns", + "UpdateColumnsResult", + "UpdateColumnsTransform", + "CommitOutcomeUnknown", "create_scalar_index", "create_index", "optimize_indices", diff --git a/lance_ray/datasource.py b/lance_ray/datasource.py index 436aad89..efba9ce9 100644 --- a/lance_ray/datasource.py +++ b/lance_ray/datasource.py @@ -288,6 +288,43 @@ def _read_fragments_with_retry( ) +def blob_field_kind(f: pa.Field) -> Optional[str]: + """Detect Lance blob columns. + + Returns: + "v2" for blob v2 extension columns, + "legacy" for legacy metadata-based blob columns, + or None if the field is not a blob. + """ + field_type = f.type + + # Blob v2: extension type `lance.blob.v2` + if isinstance(field_type, pa.ExtensionType): + ext_name = getattr(field_type, "extension_name", None) + if ext_name == "lance.blob.v2": + return "v2" + + # Legacy: LargeBinary with field metadata {"lance-encoding:blob": "true"} + try: + is_large_bin = field_type == pa.large_binary() + except Exception: + is_large_bin = False + if not is_large_bin: + return None + + meta = f.metadata + if meta is None: + return None + + # pyarrow may store metadata keys/values as str + if (meta.get("lance-encoding:blob") == "true") or ( + meta.get(b"lance-encoding:blob") == b"true" + ): + return "legacy" + + return None + + def _read_fragments( fragment_ids: list[int], lance_ds: "lance.LanceDataset", @@ -319,41 +356,7 @@ def _read_fragments( # Map column name -> blob kind ("legacy" or "v2") blob_columns: dict[str, str] = {} - def _is_blob_field(f: pa.Field) -> Optional[str]: - """Detect Lance blob columns. - - Returns: - "v2" for blob v2 extension columns, - "legacy" for legacy metadata-based blob columns, - or None if the field is not a blob. - """ - field_type = f.type - - # Blob v2: extension type `lance.blob.v2` - if isinstance(field_type, pa.ExtensionType): - ext_name = getattr(field_type, "extension_name", None) - if ext_name == "lance.blob.v2": - return "v2" - - # Legacy: LargeBinary with field metadata {"lance-encoding:blob": "true"} - try: - is_large_bin = field_type == pa.large_binary() - except Exception: - is_large_bin = False - if not is_large_bin: - return None - - meta = f.metadata - if meta is None: - return None - - # pyarrow may store metadata keys/values as str - if (meta.get("lance-encoding:blob") == "true") or ( - meta.get(b"lance-encoding:blob") == b"true" - ): - return "legacy" - - return None + _is_blob_field = blob_field_kind # Build list of blob columns to reconstruct, honoring column projection ds_field_names = ds_schema.names diff --git a/lance_ray/io.py b/lance_ray/io.py index df10171a..65818843 100644 --- a/lance_ray/io.py +++ b/lance_ray/io.py @@ -2,8 +2,11 @@ I/O operations for Lance-Ray integration. """ +import logging import pickle -from collections.abc import Callable +import uuid as uuid_module +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, Optional import pyarrow as pa @@ -15,7 +18,7 @@ from ray.util.multiprocessing import Pool from .datasink import LanceDatasink -from .datasource import LanceDatasource +from .datasource import LanceDatasource, blob_field_kind from .fragment import prepare_fragment_write_options from .utils import ( get_namespace_kwargs, @@ -36,6 +39,69 @@ | Callable[[pa.RecordBatch], pa.RecordBatch] ) +logger = logging.getLogger(__name__) + +#: Transform contract for :func:`update_columns`. +#: +#: Receives a ``pa.Table`` holding only the requested user columns (metadata +#: columns are hidden) and must return a ``pa.Table`` whose column set is +#: exactly ``output_schema.names``, with the same number of rows **in the same +#: order**. Row reordering cannot be detected and silently corrupts data. +UpdateColumnsTransform = Callable[[pa.Table], pa.Table] + +_METADATA_COLUMNS = frozenset({"_rowaddr", "_fragid", "_rowid"}) + +# Lance's built-in commit performs a conflict-checked rebase on every attempt. +# Bounded here only to limit tail latency (0 would still rebase once, it just +# would not retry). +_UPDATE_COMMIT_MAX_RETRIES = 5 + + +@dataclass(frozen=True) +class UpdateColumnsResult: + """Outcome of a distributed :func:`update_columns` run. + + The transaction UUID is intentionally absent: after a successful commit + ``version`` is the better handle, and ``dataset.read_transaction(version)`` + recovers the full transaction (including its UUID). The UUID only matters + when it is unknown whether a version was produced at all, so it is carried + by :class:`CommitOutcomeUnknown` instead. + """ + + read_version: int + version: int + columns: tuple[str, ...] + rows_updated: int + fragments_rewritten: int + + +class CommitOutcomeUnknown(RuntimeError): + """Raised when a commit neither clearly succeeded nor clearly failed. + + The transaction file is written *before* the manifest, so its presence does + not prove the commit landed. To confirm, enumerate versions after + ``read_version`` and match ``dataset.read_transaction(version).uuid`` + against :attr:`transaction_uuid`. Do not re-run the backfill until the + outcome is established. + """ + + def __init__( + self, + message: str, + *, + transaction_uuid: str, + read_version: int, + columns: tuple[str, ...], + field_ids: tuple[int, ...], + fragment_ids: tuple[int, ...], + ): + super().__init__(message) + self.transaction_uuid = transaction_uuid + self.read_version = read_version + self.columns = columns + self.field_ids = field_ids + self.fragment_ids = fragment_ids + def read_lance( uri: Optional[str] = None, @@ -144,6 +210,42 @@ def read_lance( ) +def _read_resolved_lance( + uri: str, + *, + table_id: Optional[list[str]], + columns: Optional[list[str]], + filter: Optional[str], + storage_options: Optional[dict[str, Any]], + base_store_params: Optional[dict[str, dict[str, Any]]], + dataset_options: Optional[dict[str, Any]], + namespace_impl: Optional[str], + namespace_properties: Optional[dict[str, str]], + with_metadata: bool, +) -> Dataset: + """Read an already-resolved URI while retaining namespace credentials. + + ``read_lance`` deliberately rejects an explicit URI combined with namespace + arguments. ``update_columns`` is different: the driver resolves a + namespace target once to pin its location and version, but worker dataset + opens still need the namespace client for credential refresh. Keep this + narrow internal path so the public validation contract remains unchanged. + """ + datasource = LanceDatasource( + uri=uri, + table_id=table_id, + columns=columns, + filter=filter, + storage_options=storage_options, + base_store_params=base_store_params, + dataset_options=dataset_options, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + with_metadata=with_metadata, + ) + return read_datasource(datasource=datasource, ray_remote_args={}) + + def write_lance( ds: Dataset, uri: Optional[str] = None, @@ -1141,6 +1243,649 @@ def _null_udf(in_batch: pa.RecordBatch) -> pa.RecordBatch: ) +def _blob_column_names(schema: pa.Schema) -> set[str]: + """Names of top-level blob columns (legacy or v2) in a Lance schema.""" + return {f.name for f in schema if blob_field_kind(f) is not None} + + +def _leaf_field_ids(lance_field: Any) -> list[int]: + """Leaf Lance field ids beneath (or of) ``lance_field``. + + A rewritten data file declares the *leaf* field ids it carries, which is + what ``LanceFragment.update_columns`` reports back as ``fields_modified``. + For a scalar column that is just the column's own id, but a ``list`` + column has a child item field and reports the child's id instead. + """ + children = lance_field.children() + if not children: + return [lance_field.id()] + ids: list[int] = [] + for child in children: + ids.extend(_leaf_field_ids(child)) + return ids + + +def _resolve_update_targets( + lance_ds: LanceDataset, + output_schema: pa.Schema, +) -> tuple[tuple[str, ...], tuple[int, ...]]: + """Validate ``output_schema`` against the target dataset. + + Everything here runs on the driver, before any Ray task starts, so schema + errors never surface after part of the fragments have been rewritten. + + Returns the target column names (in ``output_schema`` order) and the *leaf* + Lance field ids they cover, which the driver later cross-checks against the + ``fields_modified`` reported by every worker. + """ + if len(output_schema) == 0: + raise ValueError( + "'output_schema' must declare at least one column to update. " + "update_columns only overwrites existing columns; use " + "add_columns_from() to add new ones." + ) + + target_schema = lance_ds.schema + lance_schema = lance_ds.lance_schema + target_names = set(target_schema.names) + blob_columns = _blob_column_names(target_schema) + + names: list[str] = [] + field_ids: list[int] = [] + seen: set[str] = set() + + for field in output_schema: + name = field.name + if name in seen: + raise ValueError(f"Duplicate column '{name}' in 'output_schema'.") + seen.add(name) + + if name in _METADATA_COLUMNS: + raise ValueError( + f"Cannot update metadata column '{name}'. Metadata columns are " + "managed by lance-ray and must not appear in 'output_schema'." + ) + if "." in name: + raise ValueError( + f"Nested field path '{name}' is not supported. Only top-level " + "columns can be updated." + ) + if name not in target_names: + raise ValueError( + f"Cannot update non-existent column '{name}'. update_columns " + "only overwrites columns that already exist in the target " + "dataset; use add_columns_from() to add new ones." + ) + if name in blob_columns: + raise ValueError( + f"Cannot write blob column '{name}'. Blob columns may be read " + "via 'read_columns' but are not valid update targets." + ) + + target_field = target_schema.field(name) + if pa.types.is_struct(target_field.type): + raise ValueError( + f"Column '{name}' has a nested (struct) type, which is not " + "supported by update_columns yet." + ) + if target_field.type != field.type: + raise ValueError( + f"Type mismatch for column '{name}': target dataset has " + f"{target_field.type}, 'output_schema' declares {field.type}." + ) + if target_field.nullable != field.nullable: + raise ValueError( + f"Nullability mismatch for column '{name}': target dataset has " + f"nullable={target_field.nullable}, 'output_schema' declares " + f"nullable={field.nullable}. Arrow cast does not check " + "nullability, so this must match exactly." + ) + + lance_field = lance_schema.field(name) + if lance_field is None: + raise ValueError( + f"Column '{name}' has no Lance field id; cannot update it." + ) + names.append(name) + field_ids.extend(_leaf_field_ids(lance_field)) + + return tuple(names), tuple(field_ids) + + +def _resolve_read_columns( + lance_ds: LanceDataset, + read_columns: Optional[Sequence[str]], +) -> list[str]: + """Resolve the projection handed to the user transform. + + ``None`` expands to every top-level **non-blob** user column. Blob columns + are reconstructed into raw bytes on read, so pulling them in by default + would materialize every image/video in the table. + """ + target_schema = lance_ds.schema + blob_columns = _blob_column_names(target_schema) + + if read_columns is None: + return [name for name in target_schema.names if name not in blob_columns] + + resolved = list(read_columns) + metadata_requested = [c for c in resolved if c in _METADATA_COLUMNS] + if metadata_requested: + raise ValueError( + f"Metadata columns {metadata_requested} cannot be requested in " + "'read_columns'. lance-ray reads and re-attaches them internally; " + "the transform never sees them." + ) + missing = [c for c in resolved if c not in target_schema.names] + if missing: + raise ValueError( + f"'read_columns' references columns that do not exist in the " + f"target dataset: {missing}" + ) + return resolved + + +def _make_update_transform( + transform: "UpdateColumnsTransform", + output_schema: pa.Schema, + columns: tuple[str, ...], +) -> Callable[[pa.Table], pa.Table]: + """Wrap the user transform: hide metadata columns, validate, re-attach.""" + expected = set(columns) + column_list = list(columns) + non_nullable = [f.name for f in output_schema if not f.nullable] + + def _wrapped(batch: pa.Table) -> pa.Table: + rowaddr = batch.column("_rowaddr") + fragid = batch.column("_fragid") + user_batch = batch.drop_columns( + [c for c in _METADATA_COLUMNS if c in batch.column_names] + ) + + result = transform(user_batch) + + if isinstance(result, pa.RecordBatch): + raise TypeError( + "transform must return pa.Table, got pa.RecordBatch. Use " + "pa.Table.from_batches([rb]) to convert it." + ) + if not isinstance(result, pa.Table): + raise TypeError( + f"transform must return pa.Table, got {type(result).__name__}." + ) + if result.num_rows != batch.num_rows: + raise ValueError( + f"transform changed the row count: got {result.num_rows} rows " + f"for an input batch of {batch.num_rows}. The transform must be " + "a row-order-preserving batch mapping; filtering, sorting, " + "exploding or aggregating inside the transform is not allowed." + ) + actual = set(result.column_names) + if actual != expected: + unexpected = sorted(actual - expected) + missing = sorted(expected - actual) + raise ValueError( + "transform output columns must match 'output_schema' exactly. " + f"Unexpected: {unexpected}; missing: {missing}." + ) + + # Column order and types come from output_schema, not from the user's + # return value; cast is safe=True so precision loss raises. + out = result.select(column_list).cast(output_schema) + + for name in non_nullable: + if out.column(name).null_count: + raise ValueError( + f"transform produced nulls for non-nullable column '{name}'." + ) + + return out.append_column("_rowaddr", rowaddr).append_column("_fragid", fragid) + + return _wrapped + + +_UPDATE_RESULT_SCHEMA = pa.schema( + [ + pa.field("frag_id", pa.int64()), + pa.field("fragment_meta", pa.binary()), + pa.field("fields_modified", pa.binary()), + pa.field("rows_updated", pa.int64()), + ] +) + + +def _is_commit_conflict(exc: BaseException) -> bool: + """Whether a commit failure is a definite conflict (nothing was written). + + pylance maps every Lance error onto a builtin exception type, so the only + signal available is the message text produced by ``lance-core``'s + ``CommitConflict`` / ``RetryableCommitConflict`` / ``IncompatibleTransaction`` + variants. Anything we cannot positively identify is treated as an unknown + outcome, which is the conservative direction. + """ + message = str(exc).lower() + return "commit conflict" in message or "incompatible transaction" in message + + +def update_columns( + uri: Optional[str] = None, + *, + transform: "UpdateColumnsTransform", + output_schema: pa.Schema, + filter: Optional[str] = None, + read_columns: Optional[Sequence[str]] = None, + transform_batch_size: Optional[int] = None, + transform_ray_remote_args: Optional[Mapping[str, Any]] = None, + fragment_ray_remote_args: Optional[Mapping[str, Any]] = None, + fragment_concurrency: Optional[int] = None, + storage_options: Optional[Mapping[str, Any]] = None, + base_store_params: Optional[Mapping[str, Mapping[str, Any]]] = None, + namespace_impl: Optional[str] = None, + namespace_properties: Optional[Mapping[str, str]] = None, + table_id: Optional[Sequence[str]] = None, +) -> UpdateColumnsResult: + """Overwrite existing columns of a Lance dataset with Ray. + + This is the column-level counterpart of :func:`add_columns_from`. It pins a + snapshot, computes new values for existing columns with a distributed + transform, and rewrites only those columns via Lance's + ``RewriteColumns`` update mode: rows do not move, ``_rowaddr`` is preserved, + and untouched columns keep their original data files. It fits "most rows, + few columns" backfills; it is a poor fit for updating a handful of rows + scattered across many fragments, because a fragment's target column is + rewritten in full even for a single matching row. + + Examples: + >>> import lance_ray as lr + >>> import pyarrow as pa + >>> import pyarrow.compute as pc + >>> def bump_price(batch: pa.Table) -> pa.Table: + ... return pa.table({"price": pc.multiply(batch["price"], 1.1)}) + >>> lr.update_columns( # doctest: +SKIP + ... "/tmp/products.lance", + ... transform=bump_price, + ... output_schema=pa.schema([pa.field("price", pa.float64())]), + ... filter="status = 'active'", + ... read_columns=["price"], + ... ) + + Args: + uri: The path to the target Lance dataset. If omitted, provide + ``namespace_impl`` and ``table_id`` to resolve it from a namespace. + transform: A callable taking a ``pa.Table`` of the requested columns and + returning a ``pa.Table`` whose columns are exactly + ``output_schema.names``. It **must preserve row count and row + order**; reordering cannot be detected and silently writes values to + the wrong rows. Do not filter, sort, join, deduplicate, aggregate or + explode inside the transform. + output_schema: Schema of the transform output. Required, so that column + existence, Arrow type and nullability are checked before any Ray + task starts. Only names, types and nullability are compared; Lance + field ids and field metadata are resolved from the target dataset. + filter: A Lance filter expression. Only matching rows get new values; + unmatched rows keep their old values. Note that the physical rewrite + is still fragment-wide, so a very sparse filter is a poor fit. + read_columns: Columns handed to the transform. ``None`` expands to all + top-level non-blob columns; blob columns must be requested + explicitly because reading them materializes their raw bytes. + transform_batch_size: Arrow batch size for the transform stage, i.e. the + size of a single UDF/model-inference input. ``None`` uses Ray's + default. + transform_ray_remote_args: ``ray.remote`` args for the transform tasks + only, e.g. ``{"num_gpus": 1}``. + fragment_ray_remote_args: ``ray.remote`` args for the per-fragment + rewrite tasks only; these are CPU/memory/IO bound. + fragment_concurrency: Maximum number of fragments rewritten + concurrently. This is the knob that bounds fragment-worker peak + memory; it does not bound data accumulated by the transform stage. + storage_options: Storage options for the dataset. + base_store_params: Runtime object-store parameters keyed by base URI. + Required to read blob v2 columns backed by an external base. + namespace_impl: The namespace implementation type (e.g. "rest", "dir"). + namespace_properties: Properties for connecting to the namespace. + table_id: The table identifier as a list of strings. + + Returns: + An :class:`UpdateColumnsResult`. When the filter matches nothing the run + is a successful no-op: no transaction is created and ``version`` equals + ``read_version``. + + Raises: + CommitOutcomeUnknown: The commit neither clearly succeeded nor clearly + failed (timeout, dropped connection). Confirm before re-running. + """ + validate_uri_or_namespace(uri, namespace_impl, list(table_id) if table_id else None) + + table_id_list = list(table_id) if table_id is not None else None + namespace_properties_dict = ( + dict(namespace_properties) if namespace_properties is not None else None + ) + uri, resolved_storage_options = resolve_namespace_table( + uri, + dict(storage_options) if storage_options is not None else None, + namespace_impl, + namespace_properties_dict, + table_id_list, + ) + namespace_kwargs = get_namespace_kwargs( + namespace_impl, namespace_properties_dict, table_id_list + ) + base_store_params_dict = ( + {k: dict(v) for k, v in base_store_params.items()} + if base_store_params is not None + else None + ) + + lance_ds = LanceDataset( + uri=uri, + storage_options=resolved_storage_options, + base_store_params=base_store_params_dict, + **namespace_kwargs, + ) + resolved_read_version = lance_ds.version + + if lance_ds.has_stable_row_ids: + raise NotImplementedError( + "Distributed update_columns does not yet support datasets with " + "stable row IDs: pylance does not expose the matched row offsets " + "that Lance needs to advance _row_last_updated_at_version, so the " + "change would be invisible to CDF consumers " + "(lance-format/lance#6734)." + ) + + columns, field_ids = _resolve_update_targets(lance_ds, output_schema) + projection = _resolve_read_columns(lance_ds, read_columns) + + fragments_in_lance = {f.metadata.id for f in lance_ds.get_fragments()} + + ray_ds = _read_resolved_lance( + uri, + columns=projection, + filter=filter, + dataset_options={"version": resolved_read_version}, + storage_options=dict(resolved_storage_options), + base_store_params=base_store_params_dict, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties_dict, + table_id=table_id_list, + with_metadata=True, + ) + + map_batches_kwargs: dict[str, Any] = {} + if transform_batch_size is not None: + map_batches_kwargs["batch_size"] = transform_batch_size + if transform_ray_remote_args: + # map_batches ends in **ray_remote_args, so the options are splatted + # into the call rather than nested under a 'ray_remote_args' key. + map_batches_kwargs.update(transform_ray_remote_args) + + ray_ds = ray_ds.map_batches( + _make_update_transform(transform, output_schema, columns), + batch_format="pyarrow", + **map_batches_kwargs, + ) + + # Capture closure variables for worker tasks. + _uri = uri + _storage_options = dict(resolved_storage_options) + _base_store_params = base_store_params_dict + _namespace_impl = namespace_impl + _namespace_properties = namespace_properties_dict + _table_id = table_id_list + _read_version = resolved_read_version + _columns = list(columns) + + def _update_one_fragment(group: pa.Table) -> pa.Table: + if group.num_rows == 0: + return _UPDATE_RESULT_SCHEMA.empty_table() + + frag_ids = pc.unique(group.column("_fragid")) + if len(frag_ids) != 1: + raise ValueError( + f"map_groups received {len(frag_ids)} fragment ids in one " + f"group: {frag_ids.to_pylist()}. Expected exactly one." + ) + frag_id = int(frag_ids[0].as_py()) + + rowaddr = group.column("_rowaddr") + if pc.count_distinct(rowaddr).as_py() != group.num_rows: + raise ValueError( + f"Duplicate _rowaddr values for fragment {frag_id}. Each row " + "address must appear at most once; update_columns will not " + "silently pick one of several candidate values." + ) + derived = pc.cast(pc.shift_right(rowaddr, 32), pa.uint64()) + if pc.any(pc.not_equal(derived, frag_id)).as_py(): + raise ValueError( + f"Some _rowaddr values in the group for fragment {frag_id} do " + "not belong to that fragment." + ) + + # No sort: the fragment update is a hash join on _rowaddr, so input + # order carries no meaning, and duplicates were already rejected above. + update_table = group.select(["_rowaddr", *_columns]).combine_chunks() + + local_ns_kwargs = get_namespace_kwargs( + _namespace_impl, _namespace_properties, _table_id + ) + local_ds = LanceDataset( + uri=_uri, + storage_options=_storage_options, + base_store_params=_base_store_params, + version=_read_version, + **local_ns_kwargs, + ) + fragment = local_ds.get_fragment(frag_id) + if fragment is None: + raise ValueError(f"Fragment {frag_id} not found in Lance dataset at {_uri}") + + fragment_meta, fields_modified = fragment.update_columns( + update_table, + left_on="_rowaddr", + right_on="_rowaddr", + ) + + return pa.table( + { + "frag_id": pa.array([frag_id], type=pa.int64()), + "fragment_meta": pa.array( + [pickle.dumps(fragment_meta)], type=pa.binary() + ), + "fields_modified": pa.array( + [pickle.dumps(list(fields_modified))], type=pa.binary() + ), + "rows_updated": pa.array([group.num_rows], type=pa.int64()), + }, + schema=_UPDATE_RESULT_SCHEMA, + ) + + map_groups_kwargs: dict[str, Any] = {} + if fragment_ray_remote_args: + # Same as above: map_groups ends in **ray_remote_args. + map_groups_kwargs.update(fragment_ray_remote_args) + if fragment_concurrency is not None: + map_groups_kwargs["concurrency"] = fragment_concurrency + + result_ds = ray_ds.groupby("_fragid").map_groups( + _update_one_fragment, + batch_format="pyarrow", + **map_groups_kwargs, + ) + + rows = result_ds.take_all() + if not rows: + # The filter matched nothing: no files were written, no transaction. + return UpdateColumnsResult( + read_version=resolved_read_version, + version=resolved_read_version, + columns=columns, + rows_updated=0, + fragments_rewritten=0, + ) + + updated_fragments = [] + seen_frag_ids: set[int] = set() + rows_updated = 0 + observed_field_ids: Optional[list[int]] = None + + for row in rows: + frag_id = int(row["frag_id"]) + if frag_id not in fragments_in_lance: + raise ValueError( + f"Fragment {frag_id} is not part of the pinned snapshot " + f"(version {resolved_read_version}) of {uri}" + ) + if frag_id in seen_frag_ids: + raise ValueError(f"Duplicate fragment {frag_id} in map_groups output") + seen_frag_ids.add(frag_id) + + fragment_meta = pickle.loads(row["fragment_meta"]) + if fragment_meta.id != frag_id: + raise ValueError( + f"Fragment rewrite changed the fragment id: expected {frag_id}, " + f"got {fragment_meta.id}" + ) + updated_fragments.append(fragment_meta) + rows_updated += int(row["rows_updated"]) + + worker_field_ids = sorted(pickle.loads(row["fields_modified"])) + if observed_field_ids is None: + observed_field_ids = worker_field_ids + elif observed_field_ids != worker_field_ids: + raise ValueError( + "Workers disagree on the modified field ids: " + f"{observed_field_ids} vs {worker_field_ids}" + ) + + if observed_field_ids != sorted(field_ids): + raise ValueError( + f"Modified field ids {observed_field_ids} do not match the leaf " + f"field ids {sorted(field_ids)} of the target columns " + f"{list(columns)}." + ) + + op = LanceOperation.Update( + updated_fragments=updated_fragments, + # Lance itself reported these while rewriting; the driver-derived leaf + # ids above are only a cross-check. + fields_modified=observed_field_ids, + fields_for_preserving_frag_bitmap=[], + update_mode="rewrite_columns", + ) + # Unmodified fragments are intentionally not resubmitted: Lance's Update + # transaction merges by fragment id and carries the rest through untouched. + + transaction_uuid = str(uuid_module.uuid4()) + logger.info( + "Committing update_columns: columns=%s read_version=%s fragments=%s " + "transaction_uuid=%s", + list(columns), + resolved_read_version, + len(updated_fragments), + transaction_uuid, + ) + + committed = _commit_update( + uri=uri, + op=op, + read_version=resolved_read_version, + transaction_uuid=transaction_uuid, + storage_options=_storage_options, + namespace_kwargs=namespace_kwargs, + columns=columns, + field_ids=field_ids, + fragment_ids=tuple(sorted(seen_frag_ids)), + ) + + return UpdateColumnsResult( + read_version=resolved_read_version, + version=committed, + columns=columns, + rows_updated=rows_updated, + fragments_rewritten=len(updated_fragments), + ) + + +_COMMIT_SLOW_WARNING_S = 300.0 + + +def _commit_update( + *, + uri: str, + op: "LanceOperation.Update", + read_version: int, + transaction_uuid: str, + storage_options: dict[str, Any], + namespace_kwargs: dict[str, Any], + columns: tuple[str, ...], + field_ids: tuple[int, ...], + fragment_ids: tuple[int, ...], +) -> int: + """Commit the RewriteColumns transaction exactly once. + + Retries are delegated to Lance, whose commit loop re-runs + ``TransactionRebase.check_txn`` against every concurrent transaction and + only rebases when that check passes. lance-ray's own + ``_commit_with_retry`` must not be used here: it compares fragment id sets + and then advances ``read_version`` on its own, which would let an Update + built from stale fragment metadata overwrite a concurrent writer. + """ + import threading + + from lance.dataset import Transaction + + txn = Transaction( + read_version=read_version, + operation=op, + uuid=transaction_uuid, + ) + + timer = threading.Timer( + _COMMIT_SLOW_WARNING_S, + logger.warning, + args=( + "update_columns commit still pending after %.0fs " + "(transaction_uuid=%s, read_version=%s). If it never returns, the " + "outcome is unknown and must be confirmed against version history " + "before re-running.", + _COMMIT_SLOW_WARNING_S, + transaction_uuid, + read_version, + ), + ) + timer.daemon = True + timer.start() + try: + committed_ds = LanceDataset.commit( + uri, + txn, + max_retries=_UPDATE_COMMIT_MAX_RETRIES, + storage_options=storage_options, + **namespace_kwargs, + ) + except Exception as exc: + if _is_commit_conflict(exc): + # Nothing was committed. The caller must recompute from a fresh + # snapshot; retrying with this stale metadata is never safe. + raise + raise CommitOutcomeUnknown( + "update_columns could not determine whether its commit succeeded. " + "The transaction file is written before the manifest, so its " + "presence proves nothing. Enumerate versions after " + f"{read_version} and match read_transaction(version).uuid against " + f"{transaction_uuid!r} before re-running this backfill.", + transaction_uuid=transaction_uuid, + read_version=read_version, + columns=columns, + field_ids=field_ids, + fragment_ids=fragment_ids, + ) from exc + finally: + timer.cancel() + + return committed_ds.version + + def _validate_write_args( uri: Optional[str], namespace_impl: Optional[str], diff --git a/tests/test_update_columns.py b/tests/test_update_columns.py new file mode 100644 index 00000000..5b54a1de --- /dev/null +++ b/tests/test_update_columns.py @@ -0,0 +1,1051 @@ +"""Test cases for the distributed ``update_columns`` API.""" + +import tempfile +from pathlib import Path + +import lance +import lance_ray as lr +import pyarrow as pa +import pyarrow.compute as pc +import pytest +import ray +from lance_ray.io import _is_commit_conflict + + +@pytest.fixture +def temp_dir(): + with tempfile.TemporaryDirectory() as temp_dir: + yield temp_dir + + +PRICE_SCHEMA = pa.schema([pa.field("price", pa.float64())]) + + +def _write_products(path, rows=6, max_rows_per_file=2): + """A small multi-fragment dataset: ids 1..rows, alternating status.""" + table = pa.table( + { + "id": pa.array(list(range(1, rows + 1)), pa.int32()), + "price": pa.array( + [float(i * 10) for i in range(1, rows + 1)], pa.float64() + ), + "status": pa.array(["a" if i % 2 else "b" for i in range(1, rows + 1)]), + } + ) + lance.write_dataset(table, str(path), max_rows_per_file=max_rows_per_file) + return table + + +def _double_price(batch: pa.Table) -> pa.Table: + return pa.table({"price": pc.multiply(batch["price"], 2.0)}) + + +def _dataset_fingerprint(path): + """Version plus every data file, to prove nothing was written.""" + ds = lance.dataset(str(path)) + files = sorted( + data_file.path + for fragment in ds.get_fragments() + for data_file in fragment.data_files() + ) + return ds.version, files + + +class TestBasicBehavior: + def test_namespace_only_updates_columns(self, temp_dir): + """The driver resolves the namespace once, while workers retain it. + + This guards the internal read path used after update_columns has pinned + the namespace table location and snapshot. Public read_lance still + rejects URI-plus-namespace arguments. + """ + table_id = ["update_columns_namespace_only"] + table = pa.table( + { + "id": pa.array([1, 2, 3], pa.int32()), + "price": pa.array([10.0, 20.0, 30.0], pa.float64()), + } + ) + lr.write_lance( + ray.data.from_arrow(table), + namespace_impl="dir", + namespace_properties={"root": temp_dir}, + table_id=table_id, + min_rows_per_file=1, + max_rows_per_file=2, + ) + + result = lr.update_columns( + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + namespace_impl="dir", + namespace_properties={"root": temp_dir}, + table_id=table_id, + ) + + assert result.version == result.read_version + 1 + assert result.rows_updated == 3 + got = lr.read_lance( + namespace_impl="dir", + namespace_properties={"root": temp_dir}, + table_id=table_id, + ).take_all() + assert sorted(row["price"] for row in got) == [20.0, 40.0, 60.0] + + def test_updates_all_rows_across_fragments(self, temp_dir): + path = Path(temp_dir) / "all_rows.lance" + _write_products(path) + + result = lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + ) + + assert result.read_version == 1 + assert result.version == 2 + assert result.columns == ("price",) + assert result.rows_updated == 6 + assert result.fragments_rewritten == 3 + + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["price"] == [20.0, 40.0, 60.0, 80.0, 100.0, 120.0] + assert got["id"] == [1, 2, 3, 4, 5, 6] + + def test_filter_updates_only_matching_rows(self, temp_dir): + path = Path(temp_dir) / "filtered.lance" + _write_products(path) + + result = lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + filter="status = 'a'", + read_columns=["price"], + ) + + assert result.rows_updated == 3 + got = lance.dataset(str(path)).to_table().to_pydict() + # Only the odd ids (status 'a') doubled; the rest keep their old value. + assert got["price"] == [20.0, 20.0, 60.0, 40.0, 100.0, 60.0] + + def test_partial_fragment_coverage_is_allowed(self, temp_dir): + path = Path(temp_dir) / "partial.lance" + _write_products(path) + + # Only touches rows in the first fragment. + result = lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + filter="id <= 2", + read_columns=["price"], + ) + + assert result.rows_updated == 2 + assert result.fragments_rewritten == 1 + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["price"] == [20.0, 40.0, 30.0, 40.0, 50.0, 60.0] + + def test_updates_multiple_columns(self, temp_dir): + path = Path(temp_dir) / "multi_col.lance" + table = pa.table( + { + "id": pa.array([1, 2, 3, 4], pa.int32()), + "price": pa.array([1.0, 2.0, 3.0, 4.0], pa.float64()), + "label": pa.array(["w", "x", "y", "z"]), + } + ) + lance.write_dataset(table, str(path), max_rows_per_file=2) + + def bump_both(batch: pa.Table) -> pa.Table: + return pa.table( + { + "price": pc.multiply(batch["price"], 10.0), + "label": pc.binary_join_element_wise(batch["label"], "!", ""), + } + ) + + result = lr.update_columns( + str(path), + transform=bump_both, + output_schema=pa.schema( + [pa.field("price", pa.float64()), pa.field("label", pa.string())] + ), + read_columns=["price", "label"], + ) + + assert result.columns == ("price", "label") + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["price"] == [10.0, 20.0, 30.0, 40.0] + assert got["label"] == ["w!", "x!", "y!", "z!"] + + def test_transform_may_read_columns_it_does_not_update(self, temp_dir): + path = Path(temp_dir) / "aux_read.lance" + _write_products(path, rows=4) + + def price_from_id(batch: pa.Table) -> pa.Table: + return pa.table({"price": pc.cast(batch["id"], pa.float64())}) + + lr.update_columns( + str(path), + transform=price_from_id, + output_schema=PRICE_SCHEMA, + read_columns=["id"], + ) + + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["price"] == [1.0, 2.0, 3.0, 4.0] + + def test_no_op_when_filter_matches_nothing(self, temp_dir): + path = Path(temp_dir) / "noop.lance" + _write_products(path) + before = lance.dataset(str(path)).version + + result = lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + filter="id > 1000", + read_columns=["price"], + ) + + assert result.rows_updated == 0 + assert result.fragments_rewritten == 0 + assert result.version == result.read_version == before + assert lance.dataset(str(path)).version == before + + def test_transaction_recoverable_from_result_version(self, temp_dir): + path = Path(temp_dir) / "txn_lookup.lance" + _write_products(path, rows=2, max_rows_per_file=2) + + result = lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + ) + + # The result deliberately omits the UUID: `version` is enough to + # recover the whole transaction, including its uuid. + assert not hasattr(result, "transaction_uuid") + txn = lance.dataset(str(path)).read_transaction(result.version) + assert txn is not None + assert txn.uuid + + +class TestTransformContract: + def test_transform_does_not_see_metadata_columns(self, temp_dir): + path = Path(temp_dir) / "hidden_meta.lance" + _write_products(path, rows=4) + + def assert_no_metadata(batch: pa.Table) -> pa.Table: + for hidden in ("_rowaddr", "_fragid", "_rowid"): + assert hidden not in batch.column_names + return pa.table({"price": pc.multiply(batch["price"], 3.0)}) + + lr.update_columns( + str(path), + transform=assert_no_metadata, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + ) + + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["price"] == [30.0, 60.0, 90.0, 120.0] + + @pytest.mark.parametrize( + "bad_transform, match", + [ + ( + lambda b: pa.table( + {"price": pc.multiply(b["price"], 2.0), "extra": b["price"]} + ), + "Unexpected: \\['extra'\\]", + ), + (lambda b: pa.table({"nope": b["price"]}), "missing: \\['price'\\]"), + ( + # Ray may hand the transform single-row batches, so drop a row + # by doubling instead of slicing: that always changes the count. + lambda b: pa.concat_tables([b.select(["price"])] * 2), + "changed the row count", + ), + ( + lambda b: pa.RecordBatch.from_pydict({"price": [1.0] * b.num_rows}), + "must return pa.Table, got pa.RecordBatch", + ), + ( + lambda b: {"price": [1.0] * b.num_rows}, + "must return pa.Table, got dict", + ), + ], + ) + def test_rejects_bad_transform_output(self, temp_dir, bad_transform, match): + path = Path(temp_dir) / "bad_output.lance" + _write_products(path, rows=2, max_rows_per_file=2) + + with pytest.raises(Exception, match=match): + lr.update_columns( + str(path), + transform=bad_transform, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + ) + + def test_rejects_non_existent_output_column(self, temp_dir): + path = Path(temp_dir) / "missing_col.lance" + _write_products(path, rows=2) + + with pytest.raises(ValueError, match="non-existent column 'brand_new'"): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=pa.schema([pa.field("brand_new", pa.float64())]), + ) + + def test_rejects_type_mismatch(self, temp_dir): + path = Path(temp_dir) / "type_mismatch.lance" + _write_products(path, rows=2) + + with pytest.raises(ValueError, match="Type mismatch for column 'price'"): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=pa.schema([pa.field("price", pa.int64())]), + ) + + def test_rejects_nullable_mismatch(self, temp_dir): + path = Path(temp_dir) / "nullable_mismatch.lance" + _write_products(path, rows=2) + + with pytest.raises(ValueError, match="Nullability mismatch"): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=pa.schema( + [pa.field("price", pa.float64(), nullable=False)] + ), + ) + + def test_rejects_metadata_output_column(self, temp_dir): + path = Path(temp_dir) / "meta_output.lance" + _write_products(path, rows=2) + + with pytest.raises(ValueError, match="Cannot update metadata column"): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=pa.schema([pa.field("_rowaddr", pa.uint64())]), + ) + + def test_rejects_empty_output_schema(self, temp_dir): + path = Path(temp_dir) / "empty_schema.lance" + _write_products(path, rows=2) + + with pytest.raises(ValueError, match="at least one column"): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=pa.schema([]), + ) + + def test_rejects_metadata_in_read_columns(self, temp_dir): + path = Path(temp_dir) / "meta_read.lance" + _write_products(path, rows=2) + + with pytest.raises(ValueError, match="cannot be requested in 'read_columns'"): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price", "_rowaddr"], + ) + + def test_rejects_unknown_read_column(self, temp_dir): + path = Path(temp_dir) / "unknown_read.lance" + _write_products(path, rows=2) + + with pytest.raises(ValueError, match="do not exist in the target dataset"): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["nope"], + ) + + +class TestPhysicalCorrectness: + def test_preserves_row_address_schema_and_field_ids(self, temp_dir): + path = Path(temp_dir) / "identity.lance" + _write_products(path) + + before_ds = lance.dataset(str(path)) + before_meta = ( + lr.read_lance(str(path), with_metadata=True) + .to_pandas() + .sort_values("id") + .reset_index(drop=True) + ) + before_schema = before_ds.schema + before_field_ids = {f.name(): f.id() for f in before_ds.lance_schema.fields()} + + lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + ) + + after_ds = lance.dataset(str(path)) + after_meta = ( + lr.read_lance(str(path), with_metadata=True) + .to_pandas() + .sort_values("id") + .reset_index(drop=True) + ) + + assert after_ds.schema == before_schema + assert {f.name(): f.id() for f in after_ds.lance_schema.fields()} == ( + before_field_ids + ) + assert after_meta["_rowaddr"].tolist() == before_meta["_rowaddr"].tolist() + assert after_meta["_fragid"].tolist() == before_meta["_fragid"].tolist() + + def test_untouched_columns_keep_their_data_files(self, temp_dir): + path = Path(temp_dir) / "files.lance" + _write_products(path, rows=4, max_rows_per_file=4) + + def files_by_field(ds): + fragment = ds.get_fragments()[0] + mapping = {} + for data_file in fragment.data_files(): + for field_id in data_file.fields: + mapping[field_id] = data_file.path + return mapping + + before_ds = lance.dataset(str(path)) + before = files_by_field(before_ds) + field_ids = {f.name(): f.id() for f in before_ds.lance_schema.fields()} + + lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + ) + + after = files_by_field(lance.dataset(str(path))) + # The rewritten column moved to a new file ... + assert after[field_ids["price"]] != before[field_ids["price"]] + # ... while every other column still points at the original file. + for name in ("id", "status"): + assert after[field_ids[name]] == before[field_ids[name]] + + def test_time_travel_reads_pre_update_values(self, temp_dir): + path = Path(temp_dir) / "time_travel.lance" + _write_products(path, rows=4) + + result = lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + ) + + old = lance.dataset(str(path), version=result.read_version).to_table() + new = lance.dataset(str(path), version=result.version).to_table() + assert old["price"].to_pylist() == [10.0, 20.0, 30.0, 40.0] + assert new["price"].to_pylist() == [20.0, 40.0, 60.0, 80.0] + + def test_repeated_updates_stack(self, temp_dir): + path = Path(temp_dir) / "repeated.lance" + _write_products(path, rows=4) + + for _ in range(3): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + ) + + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["price"] == [80.0, 160.0, 240.0, 320.0] + + def test_deleted_rows_do_not_misalign_columns(self, temp_dir): + path = Path(temp_dir) / "deleted.lance" + _write_products(path, rows=6, max_rows_per_file=3) + + ds = lance.dataset(str(path)) + ds.delete("id = 2") + + lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + ) + + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["id"] == [1, 3, 4, 5, 6] + assert got["price"] == [20.0, 60.0, 80.0, 100.0, 120.0] + + +class TestTransactionBehavior: + def test_commits_a_rewrite_columns_update(self, temp_dir): + path = Path(temp_dir) / "txn_shape.lance" + _write_products(path, rows=4) + + before_ds = lance.dataset(str(path)) + price_field_id = before_ds.lance_schema.field("price").id() + + result = lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + ) + + txn = lance.dataset(str(path)).read_transaction(result.version) + assert type(txn.operation).__name__ == "Update" + assert txn.operation.update_mode == "rewrite_columns" + assert list(txn.operation.fields_modified) == [price_field_id] + assert txn.read_version == result.read_version + + def test_stale_snapshot_commit_fails(self, temp_dir): + """A stale Update must never be rebased onto a newer conflicting version.""" + from lance.dataset import Transaction + + path = Path(temp_dir) / "stale.lance" + _write_products(path, rows=4, max_rows_per_file=4) + + ds = lance.dataset(str(path)) + stale_version = ds.version + fragment = ds.get_fragments()[0] + update = pa.table( + { + "_rowaddr": pa.array([0, 1, 2, 3], pa.uint64()), + "price": pa.array([1.0, 1.0, 1.0, 1.0], pa.float64()), + } + ) + meta, fields_modified = fragment.update_columns( + update, left_on="_rowaddr", right_on="_rowaddr" + ) + + # A concurrent writer touches the same fragment first. + lance.dataset(str(path)).delete("id = 1") + + op = lance.LanceOperation.Update( + updated_fragments=[meta], + fields_modified=list(fields_modified), + fields_for_preserving_frag_bitmap=[], + update_mode="rewrite_columns", + ) + # Also pins the premise of ``_is_commit_conflict``: pylance maps this + # onto a builtin exception, so the only signal is the message text. + with pytest.raises(OSError, match="[Cc]ommit conflict") as excinfo: + lance.LanceDataset.commit( + str(path), + Transaction(read_version=stale_version, operation=op), + max_retries=5, + ) + assert _is_commit_conflict(excinfo.value) + + def test_concurrent_append_is_safely_rebased(self, temp_dir): + """An Append landing mid-flight must be rebased over, not rejected. + + The append has to happen *after* the fragment rewrite and *before* the + commit, otherwise the update simply reads the post-append snapshot and + nothing concurrent is exercised. + """ + from lance.dataset import Transaction + + path = Path(temp_dir) / "rebase.lance" + _write_products(path, rows=4, max_rows_per_file=4) + + ds = lance.dataset(str(path)) + read_version = ds.version + fragment = ds.get_fragments()[0] + update = pa.table( + { + "_rowaddr": pa.array([0, 1, 2, 3], pa.uint64()), + "price": pa.array([-1.0] * 4, pa.float64()), + } + ) + meta, fields_modified = fragment.update_columns( + update, left_on="_rowaddr", right_on="_rowaddr" + ) + + # Concurrent append on a brand new fragment; no overlap with ours. + extra = pa.table( + { + "id": pa.array([99], pa.int32()), + "price": pa.array([9.0], pa.float64()), + "status": pa.array(["c"]), + } + ) + lance.write_dataset(extra, str(path), mode="append") + assert lance.dataset(str(path)).version > read_version + + op = lance.LanceOperation.Update( + updated_fragments=[meta], + fields_modified=list(fields_modified), + fields_for_preserving_frag_bitmap=[], + update_mode="rewrite_columns", + ) + txn_uuid = "0198aaaa-bbbb-cccc-dddd-eeeeffff0001" + committed = lance.LanceDataset.commit( + str(path), + Transaction(read_version=read_version, operation=op, uuid=txn_uuid), + max_retries=5, + ) + + got = committed.to_table().to_pydict() + # Our column rewrite applied, and the concurrently appended row came + # through untouched. + assert got["price"] == [-1.0, -1.0, -1.0, -1.0, 9.0] + # The caller-fixed UUID survives Lance's rebase, which is what makes + # the CommitOutcomeUnknown recovery procedure usable. + assert committed.read_transaction(committed.version).uuid == txn_uuid + + +class TestResourceOptions: + """The Ray tuning parameters must actually reach Ray. + + ``map_batches`` / ``map_groups`` both end in ``**ray_remote_args``, so + passing a nested ``ray_remote_args={...}`` is rejected by ray.remote as an + unknown option — and only at execution time, once the whole plan has been + submitted. + """ + + def test_transform_ray_remote_args_are_applied(self, temp_dir): + path = Path(temp_dir) / "transform_args.lance" + _write_products(path, rows=4, max_rows_per_file=2) + + lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + transform_ray_remote_args={"num_cpus": 1}, + ) + + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["price"] == [20.0, 40.0, 60.0, 80.0] + + def test_fragment_ray_remote_args_and_concurrency_are_applied(self, temp_dir): + path = Path(temp_dir) / "fragment_args.lance" + _write_products(path, rows=4, max_rows_per_file=2) + + result = lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + transform_batch_size=2, + fragment_ray_remote_args={"num_cpus": 1}, + fragment_concurrency=1, + ) + + assert result.fragments_rewritten == 2 + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["price"] == [20.0, 40.0, 60.0, 80.0] + + +class TestNestedFieldIds: + """Columns whose Lance field has children report *leaf* field ids. + + ``fields_modified`` comes from the rewritten data file, which declares leaf + ids. Comparing it against the top-level field id makes every ``list``-typed + column fail — and only after the whole distributed pass has completed. + """ + + def test_updates_a_list_column(self, temp_dir): + path = Path(temp_dir) / "list_col.lance" + list_type = pa.list_(pa.int32()) + table = pa.table( + { + "id": pa.array([1, 2, 3, 4], pa.int32()), + "tags": pa.array([[1, 2], [3], [4, 5, 6], []], list_type), + } + ) + lance.write_dataset(table, str(path), max_rows_per_file=2) + + def append_marker(batch: pa.Table) -> pa.Table: + return pa.table( + { + "tags": pa.array( + [v + [99] for v in batch["tags"].to_pylist()], list_type + ) + } + ) + + result = lr.update_columns( + str(path), + transform=append_marker, + output_schema=pa.schema([pa.field("tags", list_type)]), + read_columns=["tags"], + ) + + assert result.rows_updated == 4 + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["tags"] == [[1, 2, 99], [3, 99], [4, 5, 6, 99], [99]] + + def test_fields_modified_uses_leaf_ids(self, temp_dir): + path = Path(temp_dir) / "leaf_ids.lance" + list_type = pa.list_(pa.int32()) + table = pa.table( + { + "id": pa.array([1, 2], pa.int32()), + "tags": pa.array([[1], [2]], list_type), + } + ) + lance.write_dataset(table, str(path), max_rows_per_file=2) + + ds = lance.dataset(str(path)) + tags = ds.lance_schema.field("tags") + leaf_ids = [child.id() for child in tags.children()] + # Precondition for this test to mean anything. + assert leaf_ids and leaf_ids != [tags.id()] + + result = lr.update_columns( + str(path), + transform=lambda b: pa.table( + {"tags": pa.array([[7]] * b.num_rows, list_type)} + ), + output_schema=pa.schema([pa.field("tags", list_type)]), + read_columns=["tags"], + ) + + txn = lance.dataset(str(path)).read_transaction(result.version) + assert list(txn.operation.fields_modified) == leaf_ids + + def test_updates_a_fixed_size_list_column(self, temp_dir): + path = Path(temp_dir) / "vector_col.lance" + vec_type = pa.list_(pa.float32(), 2) + table = pa.table( + { + "id": pa.array([1, 2], pa.int32()), + "vec": pa.array([[1.0, 2.0], [3.0, 4.0]], vec_type), + } + ) + lance.write_dataset(table, str(path), max_rows_per_file=2) + + lr.update_columns( + str(path), + transform=lambda b: pa.table( + { + "vec": pa.FixedSizeListArray.from_arrays( + pa.array([0.5] * (b.num_rows * 2), pa.float32()), 2 + ) + } + ), + output_schema=pa.schema([pa.field("vec", vec_type)]), + read_columns=["vec"], + ) + + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["vec"] == [[0.5, 0.5], [0.5, 0.5]] + + +class TestDriverSideRejection: + """Rejections must land before any Ray task writes a file. + + Asserting only the exception type would still pass if these checks moved + into the fragment worker, since the message propagates out of Ray anyway. + Pinning the dataset fingerprint is what actually holds the line. + """ + + @pytest.mark.parametrize( + "output_schema, read_columns, exc, match", + [ + ( + pa.schema([pa.field("brand_new", pa.float64())]), + None, + ValueError, + "non-existent column", + ), + ( + pa.schema([pa.field("price", pa.int64())]), + None, + ValueError, + "Type mismatch", + ), + ( + pa.schema([pa.field("price", pa.float64(), nullable=False)]), + None, + ValueError, + "Nullability mismatch", + ), + ( + pa.schema([pa.field("_rowaddr", pa.uint64())]), + None, + ValueError, + "metadata column", + ), + ( + pa.schema([pa.field("meta.price", pa.float64())]), + None, + ValueError, + "Nested field path", + ), + (pa.schema([]), None, ValueError, "at least one column"), + (PRICE_SCHEMA, ["price", "_rowaddr"], ValueError, "read_columns"), + (PRICE_SCHEMA, ["nope"], ValueError, "do not exist"), + ], + ) + def test_rejects_without_touching_the_dataset( + self, temp_dir, output_schema, read_columns, exc, match + ): + path = Path(temp_dir) / "untouched.lance" + _write_products(path, rows=4, max_rows_per_file=2) + before = _dataset_fingerprint(path) + + with pytest.raises(exc, match=match): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=output_schema, + read_columns=read_columns, + ) + + assert _dataset_fingerprint(path) == before + + def test_stable_row_ids_rejected_without_touching_the_dataset(self, temp_dir): + path = Path(temp_dir) / "stable_untouched.lance" + table = pa.table( + { + "id": pa.array([1, 2], pa.int32()), + "price": pa.array([1.0, 2.0], pa.float64()), + } + ) + lance.write_dataset(table, str(path), enable_stable_row_ids=True) + before = _dataset_fingerprint(path) + + with pytest.raises(NotImplementedError, match="stable row IDs"): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + ) + + assert _dataset_fingerprint(path) == before + + +class TestRejectedScenarios: + def test_rejects_stable_row_ids(self, temp_dir): + path = Path(temp_dir) / "stable_row_ids.lance" + table = pa.table( + { + "id": pa.array([1, 2], pa.int32()), + "price": pa.array([1.0, 2.0], pa.float64()), + } + ) + lance.write_dataset(table, str(path), enable_stable_row_ids=True) + + with pytest.raises(NotImplementedError, match="stable row IDs"): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + ) + + def test_rejects_nested_field_path(self, temp_dir): + path = Path(temp_dir) / "nested_path.lance" + _write_products(path, rows=2) + + with pytest.raises(ValueError, match="Nested field path"): + lr.update_columns( + str(path), + transform=_double_price, + output_schema=pa.schema([pa.field("meta.price", pa.float64())]), + ) + + def test_rejects_struct_target_column(self, temp_dir): + path = Path(temp_dir) / "struct_col.lance" + struct_type = pa.struct([pa.field("v", pa.int32())]) + table = pa.table( + { + "id": pa.array([1, 2], pa.int32()), + "meta": pa.array([{"v": 1}, {"v": 2}], struct_type), + } + ) + lance.write_dataset(table, str(path)) + + with pytest.raises(ValueError, match="nested \\(struct\\) type"): + lr.update_columns( + str(path), + transform=lambda b: b, + output_schema=pa.schema([pa.field("meta", struct_type)]), + ) + + def test_requires_uri_or_namespace(self, temp_dir): + with pytest.raises(ValueError, match="Must provide either 'uri'"): + lr.update_columns( + transform=_double_price, + output_schema=PRICE_SCHEMA, + ) + + +class TestBlobInput: + @staticmethod + def _write_blob_dataset(path, rows=4): + blob_field = pa.field( + "payload", + pa.large_binary(), + metadata={"lance-encoding:blob": "true"}, + ) + schema = pa.schema( + [ + pa.field("id", pa.int32()), + blob_field, + pa.field("size", pa.int64()), + ] + ) + table = pa.table( + { + "id": pa.array(list(range(rows)), pa.int32()), + "payload": pa.array( + [b"x" * (i + 1) for i in range(rows)], pa.large_binary() + ), + "size": pa.array([0] * rows, pa.int64()), + }, + schema=schema, + ) + lance.write_dataset(table, str(path), max_rows_per_file=2) + return table + + def test_legacy_blob_can_be_read_to_compute_a_plain_column(self, temp_dir): + path = Path(temp_dir) / "blob_input.lance" + self._write_blob_dataset(path) + + def payload_size(batch: pa.Table) -> pa.Table: + sizes = [ + len(v) if v is not None else 0 for v in batch["payload"].to_pylist() + ] + return pa.table({"size": pa.array(sizes, pa.int64())}) + + lr.update_columns( + str(path), + transform=payload_size, + output_schema=pa.schema([pa.field("size", pa.int64())]), + read_columns=["payload"], + ) + + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["size"] == [1, 2, 3, 4] + + def test_default_projection_excludes_blob_columns(self, temp_dir): + path = Path(temp_dir) / "blob_default.lance" + self._write_blob_dataset(path) + + def record_projection(batch: pa.Table) -> pa.Table: + # The transform runs in a Ray worker, so the observation has to + # travel back through the data itself. + assert "payload" not in batch.column_names + width = len(batch.column_names) + return pa.table({"size": pa.array([width] * batch.num_rows, pa.int64())}) + + lr.update_columns( + str(path), + transform=record_projection, + output_schema=pa.schema([pa.field("size", pa.int64())]), + ) + + # read_columns=None expands to the non-blob columns only: id + size. + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["size"] == [2, 2, 2, 2] + + def test_rejects_blob_output(self, temp_dir): + path = Path(temp_dir) / "blob_output.lance" + self._write_blob_dataset(path) + + with pytest.raises(ValueError, match="Cannot write blob column 'payload'"): + lr.update_columns( + str(path), + transform=lambda b: b, + output_schema=pa.schema([pa.field("payload", pa.large_binary())]), + read_columns=["payload"], + ) + + def test_blob_v2_column_is_readable_and_excluded_by_default(self, temp_dir): + blob_field = pytest.importorskip("lance").blob_field + blob_array = pytest.importorskip("lance").blob_array + + path = Path(temp_dir) / "blob_v2.lance" + schema = pa.schema( + [ + pa.field("id", pa.int32()), + blob_field("payload"), + pa.field("size", pa.int64()), + ] + ) + table = pa.table( + { + "id": pa.array([0, 1], pa.int32()), + "payload": blob_array([b"ab", b"cdef"]), + "size": pa.array([0, 0], pa.int64()), + }, + schema=schema, + ) + lance.write_dataset( + table, + str(path), + max_rows_per_file=2, + data_storage_version="2.2", # blob v2 requires file version >= 2.2 + ) + + size_schema = pa.schema([pa.field("size", pa.int64())]) + + # Blob v2 is the case the default projection exists to protect: it must + # not be pulled in unless asked for. + lr.update_columns( + str(path), + transform=lambda b: pa.table( + {"size": pa.array([len(b.column_names)] * b.num_rows, pa.int64())} + ), + output_schema=size_schema, + ) + assert lance.dataset(str(path)).to_table().to_pydict()["size"] == [2, 2] + + # ... but it is readable when explicitly requested. + def payload_size(batch: pa.Table) -> pa.Table: + sizes = [ + len(v) if v is not None else 0 for v in batch["payload"].to_pylist() + ] + return pa.table({"size": pa.array(sizes, pa.int64())}) + + lr.update_columns( + str(path), + transform=payload_size, + output_schema=size_schema, + read_columns=["payload"], + ) + assert lance.dataset(str(path)).to_table().to_pydict()["size"] == [2, 4] + + def test_unrelated_blob_column_does_not_block_updates(self, temp_dir): + path = Path(temp_dir) / "blob_unrelated.lance" + self._write_blob_dataset(path) + + lr.update_columns( + str(path), + transform=lambda b: pa.table( + {"size": pc.cast(pc.multiply(b["id"], 100), pa.int64())} + ), + output_schema=pa.schema([pa.field("size", pa.int64())]), + read_columns=["id"], + ) + + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["size"] == [0, 100, 200, 300] + + # The blob column was not rewritten; read_lance reconstructs the bytes. + payloads = ( + lr.read_lance(str(path), columns=["id", "payload"]) + .to_pandas() + .sort_values("id")["payload"] + .tolist() + ) + assert payloads == [b"x", b"xx", b"xxx", b"xxxx"] From 6d12ebc47c89b8582e78616907dbbcbd060ace46 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Mon, 27 Jul 2026 12:13:45 +0900 Subject: [PATCH 02/11] refactor: use fragment-local update path --- lance_ray/__init__.py | 2 - lance_ray/io.py | 362 +++++++++++++---------------------- tests/test_update_columns.py | 30 +-- 3 files changed, 144 insertions(+), 250 deletions(-) diff --git a/lance_ray/__init__.py b/lance_ray/__init__.py index a4a77721..770379f4 100644 --- a/lance_ray/__init__.py +++ b/lance_ray/__init__.py @@ -19,7 +19,6 @@ from .io import ( CommitOutcomeUnknown, UpdateColumnsResult, - UpdateColumnsTransform, add_columns, add_columns_from, merge_columns_from, @@ -43,7 +42,6 @@ "merge_columns_from", "update_columns", "UpdateColumnsResult", - "UpdateColumnsTransform", "CommitOutcomeUnknown", "create_scalar_index", "create_index", diff --git a/lance_ray/io.py b/lance_ray/io.py index 65818843..9a9996c5 100644 --- a/lance_ray/io.py +++ b/lance_ray/io.py @@ -18,7 +18,7 @@ from ray.util.multiprocessing import Pool from .datasink import LanceDatasink -from .datasource import LanceDatasource, blob_field_kind +from .datasource import LanceDatasource, _read_fragments, blob_field_kind from .fragment import prepare_fragment_write_options from .utils import ( get_namespace_kwargs, @@ -41,13 +41,13 @@ logger = logging.getLogger(__name__) -#: Transform contract for :func:`update_columns`. +#: Internal transform contract for :func:`update_columns`. #: -#: Receives a ``pa.Table`` holding only the requested user columns (metadata -#: columns are hidden) and must return a ``pa.Table`` whose column set is -#: exactly ``output_schema.names``, with the same number of rows **in the same -#: order**. Row reordering cannot be detected and silently corrupts data. -UpdateColumnsTransform = Callable[[pa.Table], pa.Table] +#: This is deliberately not exported as a public type alias. The callable +#: receives a ``pa.Table`` holding only user columns and must return a table +#: whose columns are exactly ``output_schema.names``, with the same number of +#: rows **in the same order**. +_UpdateColumnsTransform = Callable[[pa.Table], pa.Table] _METADATA_COLUMNS = frozenset({"_rowaddr", "_fragid", "_rowid"}) @@ -70,9 +70,7 @@ class UpdateColumnsResult: read_version: int version: int - columns: tuple[str, ...] rows_updated: int - fragments_rewritten: int class CommitOutcomeUnknown(RuntimeError): @@ -91,16 +89,10 @@ def __init__( *, transaction_uuid: str, read_version: int, - columns: tuple[str, ...], - field_ids: tuple[int, ...], - fragment_ids: tuple[int, ...], ): super().__init__(message) self.transaction_uuid = transaction_uuid self.read_version = read_version - self.columns = columns - self.field_ids = field_ids - self.fragment_ids = fragment_ids def read_lance( @@ -210,42 +202,6 @@ def read_lance( ) -def _read_resolved_lance( - uri: str, - *, - table_id: Optional[list[str]], - columns: Optional[list[str]], - filter: Optional[str], - storage_options: Optional[dict[str, Any]], - base_store_params: Optional[dict[str, dict[str, Any]]], - dataset_options: Optional[dict[str, Any]], - namespace_impl: Optional[str], - namespace_properties: Optional[dict[str, str]], - with_metadata: bool, -) -> Dataset: - """Read an already-resolved URI while retaining namespace credentials. - - ``read_lance`` deliberately rejects an explicit URI combined with namespace - arguments. ``update_columns`` is different: the driver resolves a - namespace target once to pin its location and version, but worker dataset - opens still need the namespace client for credential refresh. Keep this - narrow internal path so the public validation contract remains unchanged. - """ - datasource = LanceDatasource( - uri=uri, - table_id=table_id, - columns=columns, - filter=filter, - storage_options=storage_options, - base_store_params=base_store_params, - dataset_options=dataset_options, - namespace_impl=namespace_impl, - namespace_properties=namespace_properties, - with_metadata=with_metadata, - ) - return read_datasource(datasource=datasource, ray_remote_args={}) - - def write_lance( ds: Dataset, uri: Optional[str] = None, @@ -1385,19 +1341,18 @@ def _resolve_read_columns( return resolved -def _make_update_transform( - transform: "UpdateColumnsTransform", +def _apply_update_transform( + transform: "_UpdateColumnsTransform", output_schema: pa.Schema, columns: tuple[str, ...], ) -> Callable[[pa.Table], pa.Table]: - """Wrap the user transform: hide metadata columns, validate, re-attach.""" + """Build the per-batch transform used by a fragment-local worker.""" expected = set(columns) column_list = list(columns) non_nullable = [f.name for f in output_schema if not f.nullable] def _wrapped(batch: pa.Table) -> pa.Table: rowaddr = batch.column("_rowaddr") - fragid = batch.column("_fragid") user_batch = batch.drop_columns( [c for c in _METADATA_COLUMNS if c in batch.column_names] ) @@ -1439,19 +1394,89 @@ def _wrapped(batch: pa.Table) -> pa.Table: f"transform produced nulls for non-nullable column '{name}'." ) - return out.append_column("_rowaddr", rowaddr).append_column("_fragid", fragid) + return out.append_column("_rowaddr", rowaddr) return _wrapped -_UPDATE_RESULT_SCHEMA = pa.schema( - [ - pa.field("frag_id", pa.int64()), - pa.field("fragment_meta", pa.binary()), - pa.field("fields_modified", pa.binary()), - pa.field("rows_updated", pa.int64()), - ] -) +def _handle_update_fragment( + uri: str, + transform: "_UpdateColumnsTransform", + output_schema: pa.Schema, + columns: tuple[str, ...], + projection: list[str], + filter: Optional[str], + batch_size: int, + read_version: int, + storage_options: dict[str, Any], + base_store_params: Optional[dict[str, dict[str, Any]]], + namespace_impl: Optional[str], + namespace_properties: Optional[dict[str, str]], + table_id: Optional[list[str]], +) -> Callable[[int], Optional[tuple[int, bytes, bytes, int]]]: + """Create the fragment-local fast-path worker for ``update_columns``. + + The input always originates from the same pinned Lance fragment. Keeping + scan, transform and rewrite in one Ray task avoids the Ray Data shuffle and + regrouping required by ``merge_columns_from`` for externally-created data. + """ + apply_transform = _apply_update_transform(transform, output_schema, columns) + + def _update_fragment(fragment_id: int) -> Optional[tuple[int, bytes, bytes, int]]: + namespace_kwargs = get_namespace_kwargs( + namespace_impl, namespace_properties, table_id + ) + lance_ds = LanceDataset( + uri=uri, + storage_options=storage_options, + base_store_params=base_store_params, + version=read_version, + **namespace_kwargs, + ) + fragment = lance_ds.get_fragment(fragment_id) + if fragment is None: + raise ValueError( + f"Fragment {fragment_id} not found in Lance dataset at {uri}" + ) + + # _read_fragments is the shared scanner used by read_lance. It adds + # _rowaddr, reconstructs explicitly projected Blob columns with + # take_blobs(), and never materializes unrequested Blob columns. + update_batches: list[pa.Table] = [] + rows_updated = 0 + for batch in _read_fragments( + [fragment_id], + lance_ds, + { + "columns": projection, + "filter": filter, + "batch_size": batch_size, + }, + with_metadata=True, + ): + updated = apply_transform(batch) + update_batches.append(updated) + rows_updated += updated.num_rows + + if not update_batches: + # The exact filter is evaluated here. A non-matching fragment + # produces neither files nor transaction metadata. + return None + + update_table = pa.concat_tables(update_batches).combine_chunks() + fragment_meta, fields_modified = fragment.update_columns( + update_table, + left_on="_rowaddr", + right_on="_rowaddr", + ) + return ( + fragment_id, + pickle.dumps(fragment_meta), + pickle.dumps(list(fields_modified)), + rows_updated, + ) + + return _update_fragment def _is_commit_conflict(exc: BaseException) -> bool: @@ -1470,14 +1495,13 @@ def _is_commit_conflict(exc: BaseException) -> bool: def update_columns( uri: Optional[str] = None, *, - transform: "UpdateColumnsTransform", + transform: "_UpdateColumnsTransform", output_schema: pa.Schema, filter: Optional[str] = None, read_columns: Optional[Sequence[str]] = None, - transform_batch_size: Optional[int] = None, - transform_ray_remote_args: Optional[Mapping[str, Any]] = None, - fragment_ray_remote_args: Optional[Mapping[str, Any]] = None, - fragment_concurrency: Optional[int] = None, + batch_size: int = 1024, + ray_remote_args: Optional[Mapping[str, Any]] = None, + concurrency: Optional[int] = None, storage_options: Optional[Mapping[str, Any]] = None, base_store_params: Optional[Mapping[str, Mapping[str, Any]]] = None, namespace_impl: Optional[str] = None, @@ -1486,14 +1510,12 @@ def update_columns( ) -> UpdateColumnsResult: """Overwrite existing columns of a Lance dataset with Ray. - This is the column-level counterpart of :func:`add_columns_from`. It pins a - snapshot, computes new values for existing columns with a distributed - transform, and rewrites only those columns via Lance's - ``RewriteColumns`` update mode: rows do not move, ``_rowaddr`` is preserved, - and untouched columns keep their original data files. It fits "most rows, - few columns" backfills; it is a poor fit for updating a handful of rows - scattered across many fragments, because a fragment's target column is - rewritten in full even for a single matching row. + A pinned snapshot is processed one fragment per Ray task. Each task scans, + filters, transforms and rewrites its own fragment, avoiding the Ray Data + shuffle required to write an externally-created Dataset back to Lance. + The operation uses Lance's ``RewriteColumns`` update mode: rows do not + move, ``_rowaddr`` is preserved, and untouched columns keep their original + data files. Examples: >>> import lance_ray as lr @@ -1528,16 +1550,12 @@ def update_columns( read_columns: Columns handed to the transform. ``None`` expands to all top-level non-blob columns; blob columns must be requested explicitly because reading them materializes their raw bytes. - transform_batch_size: Arrow batch size for the transform stage, i.e. the - size of a single UDF/model-inference input. ``None`` uses Ray's - default. - transform_ray_remote_args: ``ray.remote`` args for the transform tasks - only, e.g. ``{"num_gpus": 1}``. - fragment_ray_remote_args: ``ray.remote`` args for the per-fragment - rewrite tasks only; these are CPU/memory/IO bound. - fragment_concurrency: Maximum number of fragments rewritten - concurrently. This is the knob that bounds fragment-worker peak - memory; it does not bound data accumulated by the transform stage. + batch_size: Maximum rows in one scanner/transform batch. The final + update table still accumulates all matching rows for a fragment. + ray_remote_args: ``ray.remote`` options for the complete fragment task, + for example ``{"num_gpus": 1}``. + concurrency: Maximum number of fragment tasks running concurrently. + This bounds the number of fragment update tables held in memory. storage_options: Storage options for the dataset. base_store_params: Runtime object-store parameters keyed by base URI. Required to read blob v2 columns backed by an external base. @@ -1596,131 +1614,38 @@ def update_columns( columns, field_ids = _resolve_update_targets(lance_ds, output_schema) projection = _resolve_read_columns(lance_ds, read_columns) - fragments_in_lance = {f.metadata.id for f in lance_ds.get_fragments()} - - ray_ds = _read_resolved_lance( + fragment_ids = [f.metadata.id for f in lance_ds.get_fragments()] + worker = _handle_update_fragment( uri, - columns=projection, - filter=filter, - dataset_options={"version": resolved_read_version}, - storage_options=dict(resolved_storage_options), - base_store_params=base_store_params_dict, - namespace_impl=namespace_impl, - namespace_properties=namespace_properties_dict, - table_id=table_id_list, - with_metadata=True, - ) - - map_batches_kwargs: dict[str, Any] = {} - if transform_batch_size is not None: - map_batches_kwargs["batch_size"] = transform_batch_size - if transform_ray_remote_args: - # map_batches ends in **ray_remote_args, so the options are splatted - # into the call rather than nested under a 'ray_remote_args' key. - map_batches_kwargs.update(transform_ray_remote_args) - - ray_ds = ray_ds.map_batches( - _make_update_transform(transform, output_schema, columns), - batch_format="pyarrow", - **map_batches_kwargs, - ) - - # Capture closure variables for worker tasks. - _uri = uri - _storage_options = dict(resolved_storage_options) - _base_store_params = base_store_params_dict - _namespace_impl = namespace_impl - _namespace_properties = namespace_properties_dict - _table_id = table_id_list - _read_version = resolved_read_version - _columns = list(columns) - - def _update_one_fragment(group: pa.Table) -> pa.Table: - if group.num_rows == 0: - return _UPDATE_RESULT_SCHEMA.empty_table() - - frag_ids = pc.unique(group.column("_fragid")) - if len(frag_ids) != 1: - raise ValueError( - f"map_groups received {len(frag_ids)} fragment ids in one " - f"group: {frag_ids.to_pylist()}. Expected exactly one." - ) - frag_id = int(frag_ids[0].as_py()) - - rowaddr = group.column("_rowaddr") - if pc.count_distinct(rowaddr).as_py() != group.num_rows: - raise ValueError( - f"Duplicate _rowaddr values for fragment {frag_id}. Each row " - "address must appear at most once; update_columns will not " - "silently pick one of several candidate values." - ) - derived = pc.cast(pc.shift_right(rowaddr, 32), pa.uint64()) - if pc.any(pc.not_equal(derived, frag_id)).as_py(): - raise ValueError( - f"Some _rowaddr values in the group for fragment {frag_id} do " - "not belong to that fragment." - ) - - # No sort: the fragment update is a hash join on _rowaddr, so input - # order carries no meaning, and duplicates were already rejected above. - update_table = group.select(["_rowaddr", *_columns]).combine_chunks() - - local_ns_kwargs = get_namespace_kwargs( - _namespace_impl, _namespace_properties, _table_id - ) - local_ds = LanceDataset( - uri=_uri, - storage_options=_storage_options, - base_store_params=_base_store_params, - version=_read_version, - **local_ns_kwargs, - ) - fragment = local_ds.get_fragment(frag_id) - if fragment is None: - raise ValueError(f"Fragment {frag_id} not found in Lance dataset at {_uri}") - - fragment_meta, fields_modified = fragment.update_columns( - update_table, - left_on="_rowaddr", - right_on="_rowaddr", - ) - - return pa.table( - { - "frag_id": pa.array([frag_id], type=pa.int64()), - "fragment_meta": pa.array( - [pickle.dumps(fragment_meta)], type=pa.binary() - ), - "fields_modified": pa.array( - [pickle.dumps(list(fields_modified))], type=pa.binary() - ), - "rows_updated": pa.array([group.num_rows], type=pa.int64()), - }, - schema=_UPDATE_RESULT_SCHEMA, - ) - - map_groups_kwargs: dict[str, Any] = {} - if fragment_ray_remote_args: - # Same as above: map_groups ends in **ray_remote_args. - map_groups_kwargs.update(fragment_ray_remote_args) - if fragment_concurrency is not None: - map_groups_kwargs["concurrency"] = fragment_concurrency - - result_ds = ray_ds.groupby("_fragid").map_groups( - _update_one_fragment, - batch_format="pyarrow", - **map_groups_kwargs, + transform, + output_schema, + columns, + projection, + filter, + batch_size, + resolved_read_version, + dict(resolved_storage_options), + base_store_params_dict, + namespace_impl, + namespace_properties_dict, + table_id_list, ) + pool = Pool(processes=concurrency, ray_remote_args=dict(ray_remote_args or {})) + try: + results = pool.map_async(worker, fragment_ids, chunksize=1).get() + except Exception as exc: + raise RuntimeError(f"Failed to update columns: {exc}") from exc + finally: + pool.close() + pool.join() - rows = result_ds.take_all() + rows = [row for row in results if row is not None] if not rows: # The filter matched nothing: no files were written, no transaction. return UpdateColumnsResult( read_version=resolved_read_version, version=resolved_read_version, - columns=columns, rows_updated=0, - fragments_rewritten=0, ) updated_fragments = [] @@ -1728,27 +1653,21 @@ def _update_one_fragment(group: pa.Table) -> pa.Table: rows_updated = 0 observed_field_ids: Optional[list[int]] = None - for row in rows: - frag_id = int(row["frag_id"]) - if frag_id not in fragments_in_lance: - raise ValueError( - f"Fragment {frag_id} is not part of the pinned snapshot " - f"(version {resolved_read_version}) of {uri}" - ) + for frag_id, fragment_meta_bytes, fields_modified_bytes, updated_rows in rows: if frag_id in seen_frag_ids: - raise ValueError(f"Duplicate fragment {frag_id} in map_groups output") + raise ValueError(f"Duplicate fragment {frag_id} in worker output") seen_frag_ids.add(frag_id) - fragment_meta = pickle.loads(row["fragment_meta"]) + fragment_meta = pickle.loads(fragment_meta_bytes) if fragment_meta.id != frag_id: raise ValueError( f"Fragment rewrite changed the fragment id: expected {frag_id}, " f"got {fragment_meta.id}" - ) + ) updated_fragments.append(fragment_meta) - rows_updated += int(row["rows_updated"]) + rows_updated += updated_rows - worker_field_ids = sorted(pickle.loads(row["fields_modified"])) + worker_field_ids = sorted(pickle.loads(fields_modified_bytes)) if observed_field_ids is None: observed_field_ids = worker_field_ids elif observed_field_ids != worker_field_ids: @@ -1790,19 +1709,14 @@ def _update_one_fragment(group: pa.Table) -> pa.Table: op=op, read_version=resolved_read_version, transaction_uuid=transaction_uuid, - storage_options=_storage_options, + storage_options=dict(resolved_storage_options), namespace_kwargs=namespace_kwargs, - columns=columns, - field_ids=field_ids, - fragment_ids=tuple(sorted(seen_frag_ids)), ) return UpdateColumnsResult( read_version=resolved_read_version, version=committed, - columns=columns, rows_updated=rows_updated, - fragments_rewritten=len(updated_fragments), ) @@ -1817,9 +1731,6 @@ def _commit_update( transaction_uuid: str, storage_options: dict[str, Any], namespace_kwargs: dict[str, Any], - columns: tuple[str, ...], - field_ids: tuple[int, ...], - fragment_ids: tuple[int, ...], ) -> int: """Commit the RewriteColumns transaction exactly once. @@ -1876,9 +1787,6 @@ def _commit_update( f"{transaction_uuid!r} before re-running this backfill.", transaction_uuid=transaction_uuid, read_version=read_version, - columns=columns, - field_ids=field_ids, - fragment_ids=fragment_ids, ) from exc finally: timer.cancel() diff --git a/tests/test_update_columns.py b/tests/test_update_columns.py index 5b54a1de..988490ed 100644 --- a/tests/test_update_columns.py +++ b/tests/test_update_columns.py @@ -106,9 +106,7 @@ def test_updates_all_rows_across_fragments(self, temp_dir): assert result.read_version == 1 assert result.version == 2 - assert result.columns == ("price",) assert result.rows_updated == 6 - assert result.fragments_rewritten == 3 got = lance.dataset(str(path)).to_table().to_pydict() assert got["price"] == [20.0, 40.0, 60.0, 80.0, 100.0, 120.0] @@ -145,7 +143,6 @@ def test_partial_fragment_coverage_is_allowed(self, temp_dir): ) assert result.rows_updated == 2 - assert result.fragments_rewritten == 1 got = lance.dataset(str(path)).to_table().to_pydict() assert got["price"] == [20.0, 40.0, 30.0, 40.0, 50.0, 60.0] @@ -177,7 +174,6 @@ def bump_both(batch: pa.Table) -> pa.Table: read_columns=["price", "label"], ) - assert result.columns == ("price", "label") got = lance.dataset(str(path)).to_table().to_pydict() assert got["price"] == [10.0, 20.0, 30.0, 40.0] assert got["label"] == ["w!", "x!", "y!", "z!"] @@ -213,7 +209,6 @@ def test_no_op_when_filter_matches_nothing(self, temp_dir): ) assert result.rows_updated == 0 - assert result.fragments_rewritten == 0 assert result.version == result.read_version == before assert lance.dataset(str(path)).version == before @@ -612,16 +607,10 @@ def test_concurrent_append_is_safely_rebased(self, temp_dir): class TestResourceOptions: - """The Ray tuning parameters must actually reach Ray. + """The fragment-local Ray task must receive the tuning parameters.""" - ``map_batches`` / ``map_groups`` both end in ``**ray_remote_args``, so - passing a nested ``ray_remote_args={...}`` is rejected by ray.remote as an - unknown option — and only at execution time, once the whole plan has been - submitted. - """ - - def test_transform_ray_remote_args_are_applied(self, temp_dir): - path = Path(temp_dir) / "transform_args.lance" + def test_ray_remote_args_are_applied(self, temp_dir): + path = Path(temp_dir) / "remote_args.lance" _write_products(path, rows=4, max_rows_per_file=2) lr.update_columns( @@ -629,14 +618,14 @@ def test_transform_ray_remote_args_are_applied(self, temp_dir): transform=_double_price, output_schema=PRICE_SCHEMA, read_columns=["price"], - transform_ray_remote_args={"num_cpus": 1}, + ray_remote_args={"num_cpus": 1}, ) got = lance.dataset(str(path)).to_table().to_pydict() assert got["price"] == [20.0, 40.0, 60.0, 80.0] - def test_fragment_ray_remote_args_and_concurrency_are_applied(self, temp_dir): - path = Path(temp_dir) / "fragment_args.lance" + def test_concurrency_and_batch_size_are_applied(self, temp_dir): + path = Path(temp_dir) / "concurrency.lance" _write_products(path, rows=4, max_rows_per_file=2) result = lr.update_columns( @@ -644,12 +633,11 @@ def test_fragment_ray_remote_args_and_concurrency_are_applied(self, temp_dir): transform=_double_price, output_schema=PRICE_SCHEMA, read_columns=["price"], - transform_batch_size=2, - fragment_ray_remote_args={"num_cpus": 1}, - fragment_concurrency=1, + batch_size=2, + ray_remote_args={"num_cpus": 1}, + concurrency=1, ) - assert result.fragments_rewritten == 2 got = lance.dataset(str(path)).to_table().to_pydict() assert got["price"] == [20.0, 40.0, 60.0, 80.0] From ad6fd6151171bf424b9c17a45be3a2dad0792203 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Mon, 27 Jul 2026 12:34:29 +0900 Subject: [PATCH 03/11] refactor: delegate update commits to Lance --- lance_ray/__init__.py | 2 - lance_ray/io.py | 142 +++-------------------------------- tests/test_update_columns.py | 9 +-- 3 files changed, 12 insertions(+), 141 deletions(-) diff --git a/lance_ray/__init__.py b/lance_ray/__init__.py index 770379f4..816ac9ad 100644 --- a/lance_ray/__init__.py +++ b/lance_ray/__init__.py @@ -17,7 +17,6 @@ from .fragment import LanceFragmentWriter from .index import create_index, create_scalar_index, optimize_indices from .io import ( - CommitOutcomeUnknown, UpdateColumnsResult, add_columns, add_columns_from, @@ -42,7 +41,6 @@ "merge_columns_from", "update_columns", "UpdateColumnsResult", - "CommitOutcomeUnknown", "create_scalar_index", "create_index", "optimize_indices", diff --git a/lance_ray/io.py b/lance_ray/io.py index 9a9996c5..6816f7c5 100644 --- a/lance_ray/io.py +++ b/lance_ray/io.py @@ -4,7 +4,6 @@ import logging import pickle -import uuid as uuid_module from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, Optional @@ -52,8 +51,7 @@ _METADATA_COLUMNS = frozenset({"_rowaddr", "_fragid", "_rowid"}) # Lance's built-in commit performs a conflict-checked rebase on every attempt. -# Bounded here only to limit tail latency (0 would still rebase once, it just -# would not retry). +# Bound retries only to limit tail latency. _UPDATE_COMMIT_MAX_RETRIES = 5 @@ -61,11 +59,8 @@ class UpdateColumnsResult: """Outcome of a distributed :func:`update_columns` run. - The transaction UUID is intentionally absent: after a successful commit - ``version`` is the better handle, and ``dataset.read_transaction(version)`` - recovers the full transaction (including its UUID). The UUID only matters - when it is unknown whether a version was produced at all, so it is carried - by :class:`CommitOutcomeUnknown` instead. + ``version`` identifies the committed snapshot. When the filter matches no + rows, no transaction is created and ``version == read_version``. """ read_version: int @@ -73,28 +68,6 @@ class UpdateColumnsResult: rows_updated: int -class CommitOutcomeUnknown(RuntimeError): - """Raised when a commit neither clearly succeeded nor clearly failed. - - The transaction file is written *before* the manifest, so its presence does - not prove the commit landed. To confirm, enumerate versions after - ``read_version`` and match ``dataset.read_transaction(version).uuid`` - against :attr:`transaction_uuid`. Do not re-run the backfill until the - outcome is established. - """ - - def __init__( - self, - message: str, - *, - transaction_uuid: str, - read_version: int, - ): - super().__init__(message) - self.transaction_uuid = transaction_uuid - self.read_version = read_version - - def read_lance( uri: Optional[str] = None, *, @@ -1479,19 +1452,6 @@ def _update_fragment(fragment_id: int) -> Optional[tuple[int, bytes, bytes, int] return _update_fragment -def _is_commit_conflict(exc: BaseException) -> bool: - """Whether a commit failure is a definite conflict (nothing was written). - - pylance maps every Lance error onto a builtin exception type, so the only - signal available is the message text produced by ``lance-core``'s - ``CommitConflict`` / ``RetryableCommitConflict`` / ``IncompatibleTransaction`` - variants. Anything we cannot positively identify is treated as an unknown - outcome, which is the conservative direction. - """ - message = str(exc).lower() - return "commit conflict" in message or "incompatible transaction" in message - - def update_columns( uri: Optional[str] = None, *, @@ -1568,9 +1528,6 @@ def update_columns( is a successful no-op: no transaction is created and ``version`` equals ``read_version``. - Raises: - CommitOutcomeUnknown: The commit neither clearly succeeded nor clearly - failed (timeout, dropped connection). Confirm before re-running. """ validate_uri_or_namespace(uri, namespace_impl, list(table_id) if table_id else None) @@ -1694,106 +1651,27 @@ def update_columns( # Unmodified fragments are intentionally not resubmitted: Lance's Update # transaction merges by fragment id and carries the rest through untouched. - transaction_uuid = str(uuid_module.uuid4()) logger.info( - "Committing update_columns: columns=%s read_version=%s fragments=%s " - "transaction_uuid=%s", + "Committing update_columns: columns=%s read_version=%s fragments=%s", list(columns), resolved_read_version, len(updated_fragments), - transaction_uuid, ) - - committed = _commit_update( - uri=uri, - op=op, + committed = LanceDataset.commit( + uri, + op, read_version=resolved_read_version, - transaction_uuid=transaction_uuid, + max_retries=_UPDATE_COMMIT_MAX_RETRIES, storage_options=dict(resolved_storage_options), - namespace_kwargs=namespace_kwargs, + **namespace_kwargs, ) return UpdateColumnsResult( read_version=resolved_read_version, - version=committed, + version=committed.version, rows_updated=rows_updated, ) - -_COMMIT_SLOW_WARNING_S = 300.0 - - -def _commit_update( - *, - uri: str, - op: "LanceOperation.Update", - read_version: int, - transaction_uuid: str, - storage_options: dict[str, Any], - namespace_kwargs: dict[str, Any], -) -> int: - """Commit the RewriteColumns transaction exactly once. - - Retries are delegated to Lance, whose commit loop re-runs - ``TransactionRebase.check_txn`` against every concurrent transaction and - only rebases when that check passes. lance-ray's own - ``_commit_with_retry`` must not be used here: it compares fragment id sets - and then advances ``read_version`` on its own, which would let an Update - built from stale fragment metadata overwrite a concurrent writer. - """ - import threading - - from lance.dataset import Transaction - - txn = Transaction( - read_version=read_version, - operation=op, - uuid=transaction_uuid, - ) - - timer = threading.Timer( - _COMMIT_SLOW_WARNING_S, - logger.warning, - args=( - "update_columns commit still pending after %.0fs " - "(transaction_uuid=%s, read_version=%s). If it never returns, the " - "outcome is unknown and must be confirmed against version history " - "before re-running.", - _COMMIT_SLOW_WARNING_S, - transaction_uuid, - read_version, - ), - ) - timer.daemon = True - timer.start() - try: - committed_ds = LanceDataset.commit( - uri, - txn, - max_retries=_UPDATE_COMMIT_MAX_RETRIES, - storage_options=storage_options, - **namespace_kwargs, - ) - except Exception as exc: - if _is_commit_conflict(exc): - # Nothing was committed. The caller must recompute from a fresh - # snapshot; retrying with this stale metadata is never safe. - raise - raise CommitOutcomeUnknown( - "update_columns could not determine whether its commit succeeded. " - "The transaction file is written before the manifest, so its " - "presence proves nothing. Enumerate versions after " - f"{read_version} and match read_transaction(version).uuid against " - f"{transaction_uuid!r} before re-running this backfill.", - transaction_uuid=transaction_uuid, - read_version=read_version, - ) from exc - finally: - timer.cancel() - - return committed_ds.version - - def _validate_write_args( uri: Optional[str], namespace_impl: Optional[str], diff --git a/tests/test_update_columns.py b/tests/test_update_columns.py index 988490ed..dccec92b 100644 --- a/tests/test_update_columns.py +++ b/tests/test_update_columns.py @@ -9,7 +9,6 @@ import pyarrow.compute as pc import pytest import ray -from lance_ray.io import _is_commit_conflict @pytest.fixture @@ -538,15 +537,12 @@ def test_stale_snapshot_commit_fails(self, temp_dir): fields_for_preserving_frag_bitmap=[], update_mode="rewrite_columns", ) - # Also pins the premise of ``_is_commit_conflict``: pylance maps this - # onto a builtin exception, so the only signal is the message text. - with pytest.raises(OSError, match="[Cc]ommit conflict") as excinfo: + with pytest.raises(OSError, match="[Cc]ommit conflict"): lance.LanceDataset.commit( str(path), Transaction(read_version=stale_version, operation=op), max_retries=5, ) - assert _is_commit_conflict(excinfo.value) def test_concurrent_append_is_safely_rebased(self, temp_dir): """An Append landing mid-flight must be rebased over, not rejected. @@ -601,8 +597,7 @@ def test_concurrent_append_is_safely_rebased(self, temp_dir): # Our column rewrite applied, and the concurrently appended row came # through untouched. assert got["price"] == [-1.0, -1.0, -1.0, -1.0, 9.0] - # The caller-fixed UUID survives Lance's rebase, which is what makes - # the CommitOutcomeUnknown recovery procedure usable. + # A caller-fixed UUID also survives Lance's conflict-checked rebase. assert committed.read_transaction(committed.version).uuid == txn_uuid From 6e05fc3c61c02b54a1891b46fded17111b897ae5 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Mon, 27 Jul 2026 15:22:11 +0900 Subject: [PATCH 04/11] refactor: simplify update result --- lance_ray/io.py | 5 +---- tests/test_update_columns.py | 10 +++++----- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/lance_ray/io.py b/lance_ray/io.py index 6816f7c5..e1694221 100644 --- a/lance_ray/io.py +++ b/lance_ray/io.py @@ -60,10 +60,9 @@ class UpdateColumnsResult: """Outcome of a distributed :func:`update_columns` run. ``version`` identifies the committed snapshot. When the filter matches no - rows, no transaction is created and ``version == read_version``. + rows, no transaction is created and ``rows_updated`` is zero. """ - read_version: int version: int rows_updated: int @@ -1600,7 +1599,6 @@ def update_columns( if not rows: # The filter matched nothing: no files were written, no transaction. return UpdateColumnsResult( - read_version=resolved_read_version, version=resolved_read_version, rows_updated=0, ) @@ -1667,7 +1665,6 @@ def update_columns( ) return UpdateColumnsResult( - read_version=resolved_read_version, version=committed.version, rows_updated=rows_updated, ) diff --git a/tests/test_update_columns.py b/tests/test_update_columns.py index dccec92b..1686e41e 100644 --- a/tests/test_update_columns.py +++ b/tests/test_update_columns.py @@ -83,7 +83,6 @@ def test_namespace_only_updates_columns(self, temp_dir): table_id=table_id, ) - assert result.version == result.read_version + 1 assert result.rows_updated == 3 got = lr.read_lance( namespace_impl="dir", @@ -103,7 +102,6 @@ def test_updates_all_rows_across_fragments(self, temp_dir): read_columns=["price"], ) - assert result.read_version == 1 assert result.version == 2 assert result.rows_updated == 6 @@ -208,7 +206,7 @@ def test_no_op_when_filter_matches_nothing(self, temp_dir): ) assert result.rows_updated == 0 - assert result.version == result.read_version == before + assert result.version == before assert lance.dataset(str(path)).version == before def test_transaction_recoverable_from_result_version(self, temp_dir): @@ -440,6 +438,7 @@ def files_by_field(ds): def test_time_travel_reads_pre_update_values(self, temp_dir): path = Path(temp_dir) / "time_travel.lance" _write_products(path, rows=4) + read_version = lance.dataset(str(path)).version result = lr.update_columns( str(path), @@ -448,7 +447,7 @@ def test_time_travel_reads_pre_update_values(self, temp_dir): read_columns=["price"], ) - old = lance.dataset(str(path), version=result.read_version).to_table() + old = lance.dataset(str(path), version=read_version).to_table() new = lance.dataset(str(path), version=result.version).to_table() assert old["price"].to_pylist() == [10.0, 20.0, 30.0, 40.0] assert new["price"].to_pylist() == [20.0, 40.0, 60.0, 80.0] @@ -493,6 +492,7 @@ def test_commits_a_rewrite_columns_update(self, temp_dir): _write_products(path, rows=4) before_ds = lance.dataset(str(path)) + read_version = before_ds.version price_field_id = before_ds.lance_schema.field("price").id() result = lr.update_columns( @@ -506,7 +506,7 @@ def test_commits_a_rewrite_columns_update(self, temp_dir): assert type(txn.operation).__name__ == "Update" assert txn.operation.update_mode == "rewrite_columns" assert list(txn.operation.fields_modified) == [price_field_id] - assert txn.read_version == result.read_version + assert txn.read_version == read_version def test_stale_snapshot_commit_fails(self, temp_dir): """A stale Update must never be rebased onto a newer conflicting version.""" From a4216a8a0cf41d8ae89e8c9213af2ba9e2820a53 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Mon, 27 Jul 2026 15:30:01 +0900 Subject: [PATCH 05/11] test: remove redundant update validation cases --- tests/test_update_columns.py | 129 ----------------------------------- 1 file changed, 129 deletions(-) diff --git a/tests/test_update_columns.py b/tests/test_update_columns.py index 1686e41e..6f846787 100644 --- a/tests/test_update_columns.py +++ b/tests/test_update_columns.py @@ -209,25 +209,6 @@ def test_no_op_when_filter_matches_nothing(self, temp_dir): assert result.version == before assert lance.dataset(str(path)).version == before - def test_transaction_recoverable_from_result_version(self, temp_dir): - path = Path(temp_dir) / "txn_lookup.lance" - _write_products(path, rows=2, max_rows_per_file=2) - - result = lr.update_columns( - str(path), - transform=_double_price, - output_schema=PRICE_SCHEMA, - read_columns=["price"], - ) - - # The result deliberately omits the UUID: `version` is enough to - # recover the whole transaction, including its uuid. - assert not hasattr(result, "transaction_uuid") - txn = lance.dataset(str(path)).read_transaction(result.version) - assert txn is not None - assert txn.uuid - - class TestTransformContract: def test_transform_does_not_see_metadata_columns(self, temp_dir): path = Path(temp_dir) / "hidden_meta.lance" @@ -286,88 +267,6 @@ def test_rejects_bad_transform_output(self, temp_dir, bad_transform, match): read_columns=["price"], ) - def test_rejects_non_existent_output_column(self, temp_dir): - path = Path(temp_dir) / "missing_col.lance" - _write_products(path, rows=2) - - with pytest.raises(ValueError, match="non-existent column 'brand_new'"): - lr.update_columns( - str(path), - transform=_double_price, - output_schema=pa.schema([pa.field("brand_new", pa.float64())]), - ) - - def test_rejects_type_mismatch(self, temp_dir): - path = Path(temp_dir) / "type_mismatch.lance" - _write_products(path, rows=2) - - with pytest.raises(ValueError, match="Type mismatch for column 'price'"): - lr.update_columns( - str(path), - transform=_double_price, - output_schema=pa.schema([pa.field("price", pa.int64())]), - ) - - def test_rejects_nullable_mismatch(self, temp_dir): - path = Path(temp_dir) / "nullable_mismatch.lance" - _write_products(path, rows=2) - - with pytest.raises(ValueError, match="Nullability mismatch"): - lr.update_columns( - str(path), - transform=_double_price, - output_schema=pa.schema( - [pa.field("price", pa.float64(), nullable=False)] - ), - ) - - def test_rejects_metadata_output_column(self, temp_dir): - path = Path(temp_dir) / "meta_output.lance" - _write_products(path, rows=2) - - with pytest.raises(ValueError, match="Cannot update metadata column"): - lr.update_columns( - str(path), - transform=_double_price, - output_schema=pa.schema([pa.field("_rowaddr", pa.uint64())]), - ) - - def test_rejects_empty_output_schema(self, temp_dir): - path = Path(temp_dir) / "empty_schema.lance" - _write_products(path, rows=2) - - with pytest.raises(ValueError, match="at least one column"): - lr.update_columns( - str(path), - transform=_double_price, - output_schema=pa.schema([]), - ) - - def test_rejects_metadata_in_read_columns(self, temp_dir): - path = Path(temp_dir) / "meta_read.lance" - _write_products(path, rows=2) - - with pytest.raises(ValueError, match="cannot be requested in 'read_columns'"): - lr.update_columns( - str(path), - transform=_double_price, - output_schema=PRICE_SCHEMA, - read_columns=["price", "_rowaddr"], - ) - - def test_rejects_unknown_read_column(self, temp_dir): - path = Path(temp_dir) / "unknown_read.lance" - _write_products(path, rows=2) - - with pytest.raises(ValueError, match="do not exist in the target dataset"): - lr.update_columns( - str(path), - transform=_double_price, - output_schema=PRICE_SCHEMA, - read_columns=["nope"], - ) - - class TestPhysicalCorrectness: def test_preserves_row_address_schema_and_field_ids(self, temp_dir): path = Path(temp_dir) / "identity.lance" @@ -818,34 +717,6 @@ def test_stable_row_ids_rejected_without_touching_the_dataset(self, temp_dir): class TestRejectedScenarios: - def test_rejects_stable_row_ids(self, temp_dir): - path = Path(temp_dir) / "stable_row_ids.lance" - table = pa.table( - { - "id": pa.array([1, 2], pa.int32()), - "price": pa.array([1.0, 2.0], pa.float64()), - } - ) - lance.write_dataset(table, str(path), enable_stable_row_ids=True) - - with pytest.raises(NotImplementedError, match="stable row IDs"): - lr.update_columns( - str(path), - transform=_double_price, - output_schema=PRICE_SCHEMA, - ) - - def test_rejects_nested_field_path(self, temp_dir): - path = Path(temp_dir) / "nested_path.lance" - _write_products(path, rows=2) - - with pytest.raises(ValueError, match="Nested field path"): - lr.update_columns( - str(path), - transform=_double_price, - output_schema=pa.schema([pa.field("meta.price", pa.float64())]), - ) - def test_rejects_struct_target_column(self, temp_dir): path = Path(temp_dir) / "struct_col.lance" struct_type = pa.struct([pa.field("v", pa.int32())]) From 2650a567f72799c697b18222a4bb43b32ce8fa16 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Mon, 27 Jul 2026 16:17:28 +0900 Subject: [PATCH 06/11] feat: accept record batch update transforms --- lance_ray/io.py | 46 ++++++++++---------- tests/test_update_columns.py | 84 ++++++++++++++++++++++++------------ 2 files changed, 80 insertions(+), 50 deletions(-) diff --git a/lance_ray/io.py b/lance_ray/io.py index e1694221..220b3819 100644 --- a/lance_ray/io.py +++ b/lance_ray/io.py @@ -42,11 +42,11 @@ #: Internal transform contract for :func:`update_columns`. #: -#: This is deliberately not exported as a public type alias. The callable -#: receives a ``pa.Table`` holding only user columns and must return a table -#: whose columns are exactly ``output_schema.names``, with the same number of -#: rows **in the same order**. -_UpdateColumnsTransform = Callable[[pa.Table], pa.Table] +#: This is deliberately not exported as a public type alias. The transform +#: receives a ``pa.RecordBatch`` holding only user columns and must return a +#: record batch whose columns are exactly ``output_schema.names``, with the +#: same number of rows **in the same order**. +_UpdateColumnsTransform = BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch] _METADATA_COLUMNS = frozenset({"_rowaddr", "_fragid", "_rowid"}) @@ -1324,30 +1324,29 @@ def _apply_update_transform( non_nullable = [f.name for f in output_schema if not f.nullable] def _wrapped(batch: pa.Table) -> pa.Table: - rowaddr = batch.column("_rowaddr") + # Scanner output is one RecordBatch. Combining makes that invariant + # explicit after Blob reconstruction has replaced a column. + batch = batch.combine_chunks() + rowaddr = batch.column("_rowaddr").chunk(0) user_batch = batch.drop_columns( [c for c in _METADATA_COLUMNS if c in batch.column_names] - ) + ).to_batches()[0] result = transform(user_batch) - if isinstance(result, pa.RecordBatch): - raise TypeError( - "transform must return pa.Table, got pa.RecordBatch. Use " - "pa.Table.from_batches([rb]) to convert it." - ) - if not isinstance(result, pa.Table): + if not isinstance(result, pa.RecordBatch): raise TypeError( - f"transform must return pa.Table, got {type(result).__name__}." + "transform must return pa.RecordBatch, got " + f"{type(result).__name__}." ) - if result.num_rows != batch.num_rows: + if result.num_rows != user_batch.num_rows: raise ValueError( f"transform changed the row count: got {result.num_rows} rows " - f"for an input batch of {batch.num_rows}. The transform must be " + f"for an input batch of {user_batch.num_rows}. The transform must be " "a row-order-preserving batch mapping; filtering, sorting, " "exploding or aggregating inside the transform is not allowed." ) - actual = set(result.column_names) + actual = set(result.schema.names) if actual != expected: unexpected = sorted(actual - expected) missing = sorted(expected - actual) @@ -1358,7 +1357,7 @@ def _wrapped(batch: pa.Table) -> pa.Table: # Column order and types come from output_schema, not from the user's # return value; cast is safe=True so precision loss raises. - out = result.select(column_list).cast(output_schema) + out = pa.Table.from_batches([result]).select(column_list).cast(output_schema) for name in non_nullable: if out.column(name).null_count: @@ -1480,8 +1479,10 @@ def update_columns( >>> import lance_ray as lr >>> import pyarrow as pa >>> import pyarrow.compute as pc - >>> def bump_price(batch: pa.Table) -> pa.Table: - ... return pa.table({"price": pc.multiply(batch["price"], 1.1)}) + >>> def bump_price(batch: pa.RecordBatch) -> pa.RecordBatch: + ... return pa.RecordBatch.from_pydict( + ... {"price": pc.multiply(batch["price"], 1.1)} + ... ) >>> lr.update_columns( # doctest: +SKIP ... "/tmp/products.lance", ... transform=bump_price, @@ -1493,8 +1494,9 @@ def update_columns( Args: uri: The path to the target Lance dataset. If omitted, provide ``namespace_impl`` and ``table_id`` to resolve it from a namespace. - transform: A callable taking a ``pa.Table`` of the requested columns and - returning a ``pa.Table`` whose columns are exactly + transform: A callable or :class:`lance.udf.BatchUDF` taking a + ``pa.RecordBatch`` of the requested columns and returning a + ``pa.RecordBatch`` whose columns are exactly ``output_schema.names``. It **must preserve row count and row order**; reordering cannot be detected and silently writes values to the wrong rows. Do not filter, sort, join, deduplicate, aggregate or diff --git a/tests/test_update_columns.py b/tests/test_update_columns.py index 6f846787..4f52730d 100644 --- a/tests/test_update_columns.py +++ b/tests/test_update_columns.py @@ -9,6 +9,7 @@ import pyarrow.compute as pc import pytest import ray +from lance.udf import BatchUDF @pytest.fixture @@ -35,8 +36,12 @@ def _write_products(path, rows=6, max_rows_per_file=2): return table -def _double_price(batch: pa.Table) -> pa.Table: - return pa.table({"price": pc.multiply(batch["price"], 2.0)}) +def _record_batch(data, schema=None) -> pa.RecordBatch: + return pa.RecordBatch.from_pydict(data, schema=schema) + + +def _double_price(batch: pa.RecordBatch) -> pa.RecordBatch: + return _record_batch({"price": pc.multiply(batch["price"], 2.0)}) def _dataset_fingerprint(path): @@ -154,8 +159,8 @@ def test_updates_multiple_columns(self, temp_dir): ) lance.write_dataset(table, str(path), max_rows_per_file=2) - def bump_both(batch: pa.Table) -> pa.Table: - return pa.table( + def bump_both(batch: pa.RecordBatch) -> pa.RecordBatch: + return _record_batch( { "price": pc.multiply(batch["price"], 10.0), "label": pc.binary_join_element_wise(batch["label"], "!", ""), @@ -179,8 +184,8 @@ def test_transform_may_read_columns_it_does_not_update(self, temp_dir): path = Path(temp_dir) / "aux_read.lance" _write_products(path, rows=4) - def price_from_id(batch: pa.Table) -> pa.Table: - return pa.table({"price": pc.cast(batch["id"], pa.float64())}) + def price_from_id(batch: pa.RecordBatch) -> pa.RecordBatch: + return _record_batch({"price": pc.cast(batch["id"], pa.float64())}) lr.update_columns( str(path), @@ -214,10 +219,11 @@ def test_transform_does_not_see_metadata_columns(self, temp_dir): path = Path(temp_dir) / "hidden_meta.lance" _write_products(path, rows=4) - def assert_no_metadata(batch: pa.Table) -> pa.Table: + def assert_no_metadata(batch: pa.RecordBatch) -> pa.RecordBatch: + assert isinstance(batch, pa.RecordBatch) for hidden in ("_rowaddr", "_fragid", "_rowid"): assert hidden not in batch.column_names - return pa.table({"price": pc.multiply(batch["price"], 3.0)}) + return _record_batch({"price": pc.multiply(batch["price"], 3.0)}) lr.update_columns( str(path), @@ -229,29 +235,49 @@ def assert_no_metadata(batch: pa.Table) -> pa.Table: got = lance.dataset(str(path)).to_table().to_pydict() assert got["price"] == [30.0, 60.0, 90.0, 120.0] + def test_accepts_batch_udf(self, temp_dir): + path = Path(temp_dir) / "batch_udf.lance" + _write_products(path, rows=2, max_rows_per_file=2) + + udf = BatchUDF( + lambda batch: _record_batch( + {"price": pc.multiply(batch["price"], 4.0)} + ), + output_schema=PRICE_SCHEMA, + ) + lr.update_columns( + str(path), + transform=udf, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + ) + + assert lance.dataset(str(path)).to_table()["price"].to_pylist() == [ + 40.0, + 80.0, + ] + @pytest.mark.parametrize( "bad_transform, match", [ ( - lambda b: pa.table( + lambda b: _record_batch( {"price": pc.multiply(b["price"], 2.0), "extra": b["price"]} ), "Unexpected: \\['extra'\\]", ), - (lambda b: pa.table({"nope": b["price"]}), "missing: \\['price'\\]"), + (lambda b: _record_batch({"nope": b["price"]}), "missing: \\['price'\\]"), ( - # Ray may hand the transform single-row batches, so drop a row - # by doubling instead of slicing: that always changes the count. - lambda b: pa.concat_tables([b.select(["price"])] * 2), + lambda b: _record_batch({"price": [1.0] * (b.num_rows * 2)}), "changed the row count", ), ( - lambda b: pa.RecordBatch.from_pydict({"price": [1.0] * b.num_rows}), - "must return pa.Table, got pa.RecordBatch", + lambda b: pa.table({"price": [1.0] * b.num_rows}), + "must return pa.RecordBatch, got Table", ), ( lambda b: {"price": [1.0] * b.num_rows}, - "must return pa.Table, got dict", + "must return pa.RecordBatch, got dict", ), ], ) @@ -555,8 +581,8 @@ def test_updates_a_list_column(self, temp_dir): ) lance.write_dataset(table, str(path), max_rows_per_file=2) - def append_marker(batch: pa.Table) -> pa.Table: - return pa.table( + def append_marker(batch: pa.RecordBatch) -> pa.RecordBatch: + return _record_batch( { "tags": pa.array( [v + [99] for v in batch["tags"].to_pylist()], list_type @@ -594,7 +620,7 @@ def test_fields_modified_uses_leaf_ids(self, temp_dir): result = lr.update_columns( str(path), - transform=lambda b: pa.table( + transform=lambda b: _record_batch( {"tags": pa.array([[7]] * b.num_rows, list_type)} ), output_schema=pa.schema([pa.field("tags", list_type)]), @@ -617,7 +643,7 @@ def test_updates_a_fixed_size_list_column(self, temp_dir): lr.update_columns( str(path), - transform=lambda b: pa.table( + transform=lambda b: _record_batch( { "vec": pa.FixedSizeListArray.from_arrays( pa.array([0.5] * (b.num_rows * 2), pa.float32()), 2 @@ -775,11 +801,11 @@ def test_legacy_blob_can_be_read_to_compute_a_plain_column(self, temp_dir): path = Path(temp_dir) / "blob_input.lance" self._write_blob_dataset(path) - def payload_size(batch: pa.Table) -> pa.Table: + def payload_size(batch: pa.RecordBatch) -> pa.RecordBatch: sizes = [ len(v) if v is not None else 0 for v in batch["payload"].to_pylist() ] - return pa.table({"size": pa.array(sizes, pa.int64())}) + return _record_batch({"size": pa.array(sizes, pa.int64())}) lr.update_columns( str(path), @@ -795,12 +821,14 @@ def test_default_projection_excludes_blob_columns(self, temp_dir): path = Path(temp_dir) / "blob_default.lance" self._write_blob_dataset(path) - def record_projection(batch: pa.Table) -> pa.Table: + def record_projection(batch: pa.RecordBatch) -> pa.RecordBatch: # The transform runs in a Ray worker, so the observation has to # travel back through the data itself. assert "payload" not in batch.column_names width = len(batch.column_names) - return pa.table({"size": pa.array([width] * batch.num_rows, pa.int64())}) + return _record_batch( + {"size": pa.array([width] * batch.num_rows, pa.int64())} + ) lr.update_columns( str(path), @@ -857,7 +885,7 @@ def test_blob_v2_column_is_readable_and_excluded_by_default(self, temp_dir): # not be pulled in unless asked for. lr.update_columns( str(path), - transform=lambda b: pa.table( + transform=lambda b: _record_batch( {"size": pa.array([len(b.column_names)] * b.num_rows, pa.int64())} ), output_schema=size_schema, @@ -865,11 +893,11 @@ def test_blob_v2_column_is_readable_and_excluded_by_default(self, temp_dir): assert lance.dataset(str(path)).to_table().to_pydict()["size"] == [2, 2] # ... but it is readable when explicitly requested. - def payload_size(batch: pa.Table) -> pa.Table: + def payload_size(batch: pa.RecordBatch) -> pa.RecordBatch: sizes = [ len(v) if v is not None else 0 for v in batch["payload"].to_pylist() ] - return pa.table({"size": pa.array(sizes, pa.int64())}) + return _record_batch({"size": pa.array(sizes, pa.int64())}) lr.update_columns( str(path), @@ -885,7 +913,7 @@ def test_unrelated_blob_column_does_not_block_updates(self, temp_dir): lr.update_columns( str(path), - transform=lambda b: pa.table( + transform=lambda b: _record_batch( {"size": pc.cast(pc.multiply(b["id"], 100), pa.int64())} ), output_schema=pa.schema([pa.field("size", pa.int64())]), From 6e283b95b794b60cfbb34a00cbfc9983b383d6ce Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Mon, 27 Jul 2026 18:36:30 +0900 Subject: [PATCH 07/11] docs: document distributed update columns --- docs/src/data-evolution.md | 70 ++++++++++++++++++++++++++++++++++++++ docs/src/index.md | 4 +-- 2 files changed, 72 insertions(+), 2 deletions(-) diff --git a/docs/src/data-evolution.md b/docs/src/data-evolution.md index 25102c9e..d9b449f4 100644 --- a/docs/src/data-evolution.md +++ b/docs/src/data-evolution.md @@ -30,3 +30,73 @@ Add columns to an existing Lance dataset using Ray's distributed processing. - `concurrency`: Optional number of concurrent processes **Returns:** None + +## `update_columns` + +```python +from lance_ray import update_columns +import pyarrow as pa +import pyarrow.compute as pc + + +def increase_price(batch: pa.RecordBatch) -> pa.RecordBatch: + return pa.RecordBatch.from_pydict( + {"price": pc.multiply(batch["price"], 1.1)} + ) + + +result = update_columns( + "products.lance", + transform=increase_price, + output_schema=pa.schema([pa.field("price", pa.float64())]), + filter="status = 'active'", + read_columns=["price"], +) +print(result.version, result.rows_updated) +``` + +Overwrite existing columns with a distributed, fragment-local Ray transform. +Lance-Ray pins a dataset snapshot and processes each fragment in one Ray task. +The task scans the fragment, applies `filter` and `transform`, then rewrites +only the requested columns. Untouched columns retain their original data files. + +`transform` accepts and returns a `pyarrow.RecordBatch`. A `lance.udf.BatchUDF` +is also accepted. Its result must contain exactly the fields in `output_schema`. +The transform must preserve both row count and row order: do not filter, sort, +join, deduplicate, aggregate, or explode rows inside it. + +**Parameters:** + +- `uri`: Path to the target Lance dataset. Alternatively, resolve the table + with `namespace_impl` and `table_id`. +- `transform`: A `RecordBatch` transform or `BatchUDF` that produces replacement + values. +- `output_schema`: Required schema for the replacement columns. Every field + must already exist in the dataset, and its type and nullability must match. +- `filter`: Optional Lance filter expression. Only matching rows receive new + values; the containing fragment is nevertheless rewritten. +- `read_columns`: Columns supplied to `transform`. When omitted, all top-level + non-Blob columns are read. Request Blob columns explicitly; they are passed + to the transform as raw `LargeBinary` bytes and cannot be written by this API. +- `batch_size`: Maximum rows in each scanner and transform batch. A fragment's + matching update rows are still accumulated before the fragment is written. +- `ray_remote_args`: Ray resource options for each fragment task, such as + `{"num_gpus": 1}`. +- `concurrency`: Maximum number of fragment tasks running at once. Lower it to + bound aggregate fragment-update memory. +- `storage_options`, `base_store_params`, `namespace_impl`, + `namespace_properties`, `table_id`: Dataset storage and namespace options. + +**Returns:** `UpdateColumnsResult(version, rows_updated)`. A filter that +matches no rows succeeds without a transaction; `rows_updated` is `0` and +`version` remains unchanged. + +### Limitations and operational notes + +- Datasets with stable row IDs are rejected. The underlying Python binding does + not expose the updated row offsets required for correct CDF metadata. +- Updating an indexed column is allowed, but current Lance index maintenance + may leave the affected index stale. Rebuild indexes before relying on them + after such an update. +- A sparse filter still causes fragment-wide rewrites. Prefer this API when the + updated rows are reasonably concentrated in their fragments. diff --git a/docs/src/index.md b/docs/src/index.md index 6bb7d816..0ca9b7a4 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -12,7 +12,7 @@ enabling scalable data processing workflows with optimal performance. - **Optimized I/O**: Efficient reading and writing of Lance datasets with Ray integration - **Schema Validation**: Automatic schema compatibility checking between Ray and Lance - **Flexible Filtering**: Support for complex filtering pushdown on distributed Lance data -- **Data Evolution**: Support for data evolution to add new columns and distributedly backfill data using a Ray UDF +- **Data Evolution**: Add or overwrite columns and backfill data using distributed Ray transforms - **Index Maintenance**: Incremental index updates and distributed dataset compaction - **Catalog Integration**: Support for working with Lance datasets stored in various catalog services (e.g. Hive MetaStore, Iceberg REST Catalog, Unity, Gravitino, AWS Glue, etc.) @@ -51,4 +51,4 @@ ray_dataset = read_lance("my_dataset.lance") # Perform distributed operations result = ray_dataset.filter(lambda row: row["value"] > 100).count() print(f"Filtered count: {result}") -``` \ No newline at end of file +``` From 774720cbec45245d48263554bc05b348c458dad4 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Tue, 28 Jul 2026 08:57:12 +0900 Subject: [PATCH 08/11] refactor: stream update column batches --- docs/src/data-evolution.md | 5 ++-- lance_ray/datasource.py | 30 ++++++++++++------- lance_ray/io.py | 57 +++++++++++++++++++++--------------- tests/test_update_columns.py | 23 +++++++++++++++ 4 files changed, 78 insertions(+), 37 deletions(-) diff --git a/docs/src/data-evolution.md b/docs/src/data-evolution.md index d9b449f4..6b180758 100644 --- a/docs/src/data-evolution.md +++ b/docs/src/data-evolution.md @@ -78,8 +78,9 @@ join, deduplicate, aggregate, or explode rows inside it. - `read_columns`: Columns supplied to `transform`. When omitted, all top-level non-Blob columns are read. Request Blob columns explicitly; they are passed to the transform as raw `LargeBinary` bytes and cannot be written by this API. -- `batch_size`: Maximum rows in each scanner and transform batch. A fragment's - matching update rows are still accumulated before the fragment is written. +- `batch_size`: Maximum rows in each scanner and transform batch. Lance receives + the update values as a RecordBatch stream, though its underlying update join + can still materialize a fragment's matching rows. - `ray_remote_args`: Ray resource options for each fragment task, such as `{"num_gpus": 1}`. - `concurrency`: Maximum number of fragment tasks running at once. Lower it to diff --git a/lance_ray/datasource.py b/lance_ray/datasource.py index efba9ce9..a875f7d6 100644 --- a/lance_ray/datasource.py +++ b/lance_ray/datasource.py @@ -330,13 +330,15 @@ def _read_fragments( lance_ds: "lance.LanceDataset", scanner_options: dict[str, Any], with_metadata: bool = False, -) -> Iterator[pa.Table]: + as_record_batches: bool = False, +) -> Iterator[pa.Table | pa.RecordBatch]: """Read Lance fragments in batches. This enhanced reader detects Lance blob-encoded columns and reconstructs raw bytes using the :meth:`LanceDataset.take_blobs` API, returning :class:`pyarrow.LargeBinaryArray` columns instead of the default - struct-based descriptors. + struct-based descriptors. ``as_record_batches`` is for callers, such as + ``update_columns``, that keep the scanner's batch-native execution model. Row ordering is preserved by using per-batch row IDs. @@ -386,19 +388,20 @@ def _read_fragments( for batch in scanner.to_reader(): # Fast path: no blob columns requested in this scan if not blob_columns: - table = pa.Table.from_batches([batch]) - - if with_metadata and "_rowaddr" in table.column_names: - rowaddr_col = table.column("_rowaddr") + if with_metadata and "_rowaddr" in batch.column_names: + rowaddr_col = batch.column("_rowaddr") fragid_values = pc.cast(pc.shift_right(rowaddr_col, 32), pa.uint64()) - table = table.append_column("_fragid", fragid_values) + batch = batch.append_column("_fragid", fragid_values) if not with_metadata: for col in ("_rowaddr", "_fragid"): - if col in table.column_names: - table = table.drop_columns([col]) + if col in batch.column_names: + batch = batch.drop_columns([col]) - yield table + if as_record_batches: + yield batch + else: + yield pa.Table.from_batches([batch]) continue # Build a table so we can manipulate columns easily @@ -547,4 +550,9 @@ def _read_fragments( if col in table.column_names: table = table.drop_columns([col]) - yield table + if as_record_batches: + # Blob replacement currently uses Table column operations. Convert + # back before returning so callers can retain a RecordBatch stream. + yield table.combine_chunks().to_batches()[0] + else: + yield table diff --git a/lance_ray/io.py b/lance_ray/io.py index 220b3819..477d331d 100644 --- a/lance_ray/io.py +++ b/lance_ray/io.py @@ -1317,20 +1317,17 @@ def _apply_update_transform( transform: "_UpdateColumnsTransform", output_schema: pa.Schema, columns: tuple[str, ...], -) -> Callable[[pa.Table], pa.Table]: +) -> Callable[[pa.RecordBatch], pa.RecordBatch]: """Build the per-batch transform used by a fragment-local worker.""" expected = set(columns) column_list = list(columns) non_nullable = [f.name for f in output_schema if not f.nullable] - def _wrapped(batch: pa.Table) -> pa.Table: - # Scanner output is one RecordBatch. Combining makes that invariant - # explicit after Blob reconstruction has replaced a column. - batch = batch.combine_chunks() - rowaddr = batch.column("_rowaddr").chunk(0) + def _wrapped(batch: pa.RecordBatch) -> pa.RecordBatch: + rowaddr = batch.column("_rowaddr") user_batch = batch.drop_columns( [c for c in _METADATA_COLUMNS if c in batch.column_names] - ).to_batches()[0] + ) result = transform(user_batch) @@ -1358,14 +1355,15 @@ def _wrapped(batch: pa.Table) -> pa.Table: # Column order and types come from output_schema, not from the user's # return value; cast is safe=True so precision loss raises. out = pa.Table.from_batches([result]).select(column_list).cast(output_schema) + normalized = out.to_batches()[0] for name in non_nullable: - if out.column(name).null_count: + if normalized.column(name).null_count: raise ValueError( f"transform produced nulls for non-nullable column '{name}'." ) - return out.append_column("_rowaddr", rowaddr) + return normalized.append_column("_rowaddr", rowaddr) return _wrapped @@ -1410,12 +1408,11 @@ def _update_fragment(fragment_id: int) -> Optional[tuple[int, bytes, bytes, int] f"Fragment {fragment_id} not found in Lance dataset at {uri}" ) - # _read_fragments is the shared scanner used by read_lance. It adds + # _read_fragments is the shared scanner used by read_lance. It adds # _rowaddr, reconstructs explicitly projected Blob columns with - # take_blobs(), and never materializes unrequested Blob columns. - update_batches: list[pa.Table] = [] - rows_updated = 0 - for batch in _read_fragments( + # take_blobs(), and never materializes unrequested Blob columns. Keep + # its result as a RecordBatch stream until Lance consumes it. + scanned_batches = _read_fragments( [fragment_id], lance_ds, { @@ -1424,19 +1421,30 @@ def _update_fragment(fragment_id: int) -> Optional[tuple[int, bytes, bytes, int] "batch_size": batch_size, }, with_metadata=True, - ): - updated = apply_transform(batch) - update_batches.append(updated) - rows_updated += updated.num_rows - - if not update_batches: + as_record_batches=True, + ) + try: + first_updated = apply_transform(next(scanned_batches)) + except StopIteration: # The exact filter is evaluated here. A non-matching fragment # produces neither files nor transaction metadata. return None - update_table = pa.concat_tables(update_batches).combine_chunks() + rows_updated = first_updated.num_rows + + def updated_batches() -> Iterator[pa.RecordBatch]: + nonlocal rows_updated + yield first_updated + for scanned_batch in scanned_batches: + updated = apply_transform(scanned_batch) + rows_updated += updated.num_rows + yield updated + + update_reader = pa.RecordBatchReader.from_batches( + first_updated.schema, updated_batches() + ) fragment_meta, fields_modified = fragment.update_columns( - update_table, + update_reader, left_on="_rowaddr", right_on="_rowaddr", ) @@ -1511,8 +1519,9 @@ def update_columns( read_columns: Columns handed to the transform. ``None`` expands to all top-level non-blob columns; blob columns must be requested explicitly because reading them materializes their raw bytes. - batch_size: Maximum rows in one scanner/transform batch. The final - update table still accumulates all matching rows for a fragment. + batch_size: Maximum rows in one scanner/transform batch. Lance consumes + these as a RecordBatch stream, although its underlying update join + can still materialize a fragment's matching rows. ray_remote_args: ``ray.remote`` options for the complete fragment task, for example ``{"num_gpus": 1}``. concurrency: Maximum number of fragment tasks running concurrently. diff --git a/tests/test_update_columns.py b/tests/test_update_columns.py index 4f52730d..81668387 100644 --- a/tests/test_update_columns.py +++ b/tests/test_update_columns.py @@ -114,6 +114,29 @@ def test_updates_all_rows_across_fragments(self, temp_dir): assert got["price"] == [20.0, 40.0, 60.0, 80.0, 100.0, 120.0] assert got["id"] == [1, 2, 3, 4, 5, 6] + def test_updates_one_fragment_over_multiple_record_batches(self, temp_dir): + path = Path(temp_dir) / "multi_batch_fragment.lance" + _write_products(path, rows=6, max_rows_per_file=6) + + result = lr.update_columns( + str(path), + transform=_double_price, + output_schema=PRICE_SCHEMA, + read_columns=["price"], + batch_size=2, + concurrency=1, + ) + + assert result.rows_updated == 6 + assert lance.dataset(str(path)).to_table()["price"].to_pylist() == [ + 20.0, + 40.0, + 60.0, + 80.0, + 100.0, + 120.0, + ] + def test_filter_updates_only_matching_rows(self, temp_dir): path = Path(temp_dir) / "filtered.lance" _write_products(path) From 07c8a8746b7a254cc448af9507c6b9e895c7052d Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Tue, 28 Jul 2026 11:27:39 +0900 Subject: [PATCH 09/11] feat: stream distributed column updates --- docs/src/data-evolution.md | 13 +- lance_ray/io.py | 136 +++++++++--------- tests/test_update_columns.py | 258 ++++++++++++++++++++++++----------- 3 files changed, 264 insertions(+), 143 deletions(-) diff --git a/docs/src/data-evolution.md b/docs/src/data-evolution.md index 6b180758..ad19c807 100644 --- a/docs/src/data-evolution.md +++ b/docs/src/data-evolution.md @@ -48,7 +48,7 @@ def increase_price(batch: pa.RecordBatch) -> pa.RecordBatch: result = update_columns( "products.lance", transform=increase_price, - output_schema=pa.schema([pa.field("price", pa.float64())]), + columns=["price"], filter="status = 'active'", read_columns=["price"], ) @@ -61,7 +61,7 @@ The task scans the fragment, applies `filter` and `transform`, then rewrites only the requested columns. Untouched columns retain their original data files. `transform` accepts and returns a `pyarrow.RecordBatch`. A `lance.udf.BatchUDF` -is also accepted. Its result must contain exactly the fields in `output_schema`. +is also accepted. Its result must contain exactly the fields named in `columns`. The transform must preserve both row count and row order: do not filter, sort, join, deduplicate, aggregate, or explode rows inside it. @@ -71,8 +71,13 @@ join, deduplicate, aggregate, or explode rows inside it. with `namespace_impl` and `table_id`. - `transform`: A `RecordBatch` transform or `BatchUDF` that produces replacement values. -- `output_schema`: Required schema for the replacement columns. Every field - must already exist in the dataset, and its type and nullability must match. +- `columns`: Required names of the columns to overwrite. Each must already exist + in the dataset. Their Arrow type and nullability are taken from the dataset — + `update_columns` cannot change them — and the transform result is cast to them + with `safe=True`. That rejects out-of-range integers, truncating time-unit + conversions, and unparseable strings, but it does **not** catch float + narrowing: returning a `float64` for a `float32` column rounds, and overflows + to `inf`, silently. Produce the column's own type when precision matters. - `filter`: Optional Lance filter expression. Only matching rows receive new values; the containing fragment is nevertheless rewritten. - `read_columns`: Columns supplied to `transform`. When omitted, all top-level diff --git a/lance_ray/io.py b/lance_ray/io.py index 477d331d..7ed1751f 100644 --- a/lance_ray/io.py +++ b/lance_ray/io.py @@ -4,8 +4,9 @@ import logging import pickle -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass +from itertools import chain from typing import TYPE_CHECKING, Any, Literal, Optional import pyarrow as pa @@ -44,7 +45,7 @@ #: #: This is deliberately not exported as a public type alias. The transform #: receives a ``pa.RecordBatch`` holding only user columns and must return a -#: record batch whose columns are exactly ``output_schema.names``, with the +#: record batch whose columns are exactly the requested ``columns``, with the #: same number of rows **in the same order**. _UpdateColumnsTransform = BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch] @@ -1195,20 +1196,30 @@ def _leaf_field_ids(lance_field: Any) -> list[int]: def _resolve_update_targets( lance_ds: LanceDataset, - output_schema: pa.Schema, -) -> tuple[tuple[str, ...], tuple[int, ...]]: - """Validate ``output_schema`` against the target dataset. + columns: Sequence[str], +) -> tuple[tuple[int, ...], pa.Schema]: + """Validate ``columns`` against the target dataset. Everything here runs on the driver, before any Ray task starts, so schema errors never surface after part of the fragments have been rewritten. - Returns the target column names (in ``output_schema`` order) and the *leaf* - Lance field ids they cover, which the driver later cross-checks against the - ``fields_modified`` reported by every worker. + Returns the *leaf* Lance field ids the columns cover -- which the driver + later cross-checks against the ``fields_modified`` reported by every worker + -- and the output schema the transform result is normalized to, holding the + requested columns in the requested order. That schema is taken from the + target dataset rather than the caller: update_columns only overwrites + existing columns, so their Arrow type and nullability are already fixed. """ - if len(output_schema) == 0: + if isinstance(columns, str): + # str satisfies Sequence[str], so this would otherwise iterate + # characters and update whatever single-letter columns happen to exist. + raise TypeError( + f"'columns' must be a sequence of column names, not a bare string. " + f"Did you mean ['{columns}']?" + ) + if len(columns) == 0: raise ValueError( - "'output_schema' must declare at least one column to update. " + "'columns' must name at least one column to update. " "update_columns only overwrites existing columns; use " "add_columns_from() to add new ones." ) @@ -1218,20 +1229,24 @@ def _resolve_update_targets( target_names = set(target_schema.names) blob_columns = _blob_column_names(target_schema) - names: list[str] = [] + fields: list[pa.Field] = [] field_ids: list[int] = [] seen: set[str] = set() - for field in output_schema: - name = field.name + for name in columns: + if not isinstance(name, str): + raise TypeError( + f"'columns' must contain column names as str, got " + f"{type(name).__name__}." + ) if name in seen: - raise ValueError(f"Duplicate column '{name}' in 'output_schema'.") + raise ValueError(f"Duplicate column '{name}' in 'columns'.") seen.add(name) if name in _METADATA_COLUMNS: raise ValueError( f"Cannot update metadata column '{name}'. Metadata columns are " - "managed by lance-ray and must not appear in 'output_schema'." + "managed by lance-ray and must not appear in 'columns'." ) if "." in name: raise ValueError( @@ -1256,28 +1271,16 @@ def _resolve_update_targets( f"Column '{name}' has a nested (struct) type, which is not " "supported by update_columns yet." ) - if target_field.type != field.type: - raise ValueError( - f"Type mismatch for column '{name}': target dataset has " - f"{target_field.type}, 'output_schema' declares {field.type}." - ) - if target_field.nullable != field.nullable: - raise ValueError( - f"Nullability mismatch for column '{name}': target dataset has " - f"nullable={target_field.nullable}, 'output_schema' declares " - f"nullable={field.nullable}. Arrow cast does not check " - "nullability, so this must match exactly." - ) lance_field = lance_schema.field(name) if lance_field is None: raise ValueError( f"Column '{name}' has no Lance field id; cannot update it." ) - names.append(name) + fields.append(target_field) field_ids.extend(_leaf_field_ids(lance_field)) - return tuple(names), tuple(field_ids) + return tuple(field_ids), pa.schema(fields) def _resolve_read_columns( @@ -1316,11 +1319,10 @@ def _resolve_read_columns( def _apply_update_transform( transform: "_UpdateColumnsTransform", output_schema: pa.Schema, - columns: tuple[str, ...], ) -> Callable[[pa.RecordBatch], pa.RecordBatch]: """Build the per-batch transform used by a fragment-local worker.""" - expected = set(columns) - column_list = list(columns) + column_list = list(output_schema.names) + expected = set(column_list) non_nullable = [f.name for f in output_schema if not f.nullable] def _wrapped(batch: pa.RecordBatch) -> pa.RecordBatch: @@ -1333,8 +1335,7 @@ def _wrapped(batch: pa.RecordBatch) -> pa.RecordBatch: if not isinstance(result, pa.RecordBatch): raise TypeError( - "transform must return pa.RecordBatch, got " - f"{type(result).__name__}." + f"transform must return pa.RecordBatch, got {type(result).__name__}." ) if result.num_rows != user_batch.num_rows: raise ValueError( @@ -1348,12 +1349,12 @@ def _wrapped(batch: pa.RecordBatch) -> pa.RecordBatch: unexpected = sorted(actual - expected) missing = sorted(expected - actual) raise ValueError( - "transform output columns must match 'output_schema' exactly. " + "transform output columns must match 'columns' exactly. " f"Unexpected: {unexpected}; missing: {missing}." ) - # Column order and types come from output_schema, not from the user's - # return value; cast is safe=True so precision loss raises. + # Column order and types come from the target dataset, not from the + # user's return value; cast is safe=True so precision loss raises. out = pa.Table.from_batches([result]).select(column_list).cast(output_schema) normalized = out.to_batches()[0] @@ -1372,7 +1373,6 @@ def _handle_update_fragment( uri: str, transform: "_UpdateColumnsTransform", output_schema: pa.Schema, - columns: tuple[str, ...], projection: list[str], filter: Optional[str], batch_size: int, @@ -1389,7 +1389,7 @@ def _handle_update_fragment( scan, transform and rewrite in one Ray task avoids the Ray Data shuffle and regrouping required by ``merge_columns_from`` for externally-created data. """ - apply_transform = _apply_update_transform(transform, output_schema, columns) + apply_transform = _apply_update_transform(transform, output_schema) def _update_fragment(fragment_id: int) -> Optional[tuple[int, bytes, bytes, int]]: namespace_kwargs = get_namespace_kwargs( @@ -1423,25 +1423,34 @@ def _update_fragment(fragment_id: int) -> Optional[tuple[int, bytes, bytes, int] with_metadata=True, as_record_batches=True, ) - try: - first_updated = apply_transform(next(scanned_batches)) - except StopIteration: + first_scanned = next(scanned_batches, None) + if first_scanned is None: # The exact filter is evaluated here. A non-matching fragment - # produces neither files nor transaction metadata. + # produces neither files nor transaction metadata. Lance rejects an + # empty reader outright ("HashJoiner: No data"), so this batch has to + # be pulled before the rewrite starts; it is then fed straight back + # into the stream below, so nothing is scanned twice. return None - rows_updated = first_updated.num_rows + # apply_transform emits output_schema plus the _rowaddr it carries + # through, so the reader schema is known without transforming a batch. + reader_schema = pa.schema( + [ + *output_schema, + pa.field("_rowaddr", first_scanned.schema.field("_rowaddr").type), + ] + ) + rows_updated = 0 def updated_batches() -> Iterator[pa.RecordBatch]: nonlocal rows_updated - yield first_updated - for scanned_batch in scanned_batches: + for scanned_batch in chain([first_scanned], scanned_batches): updated = apply_transform(scanned_batch) rows_updated += updated.num_rows yield updated update_reader = pa.RecordBatchReader.from_batches( - first_updated.schema, updated_batches() + reader_schema, updated_batches() ) fragment_meta, fields_modified = fragment.update_columns( update_reader, @@ -1462,7 +1471,7 @@ def update_columns( uri: Optional[str] = None, *, transform: "_UpdateColumnsTransform", - output_schema: pa.Schema, + columns: Sequence[str], filter: Optional[str] = None, read_columns: Optional[Sequence[str]] = None, batch_size: int = 1024, @@ -1494,7 +1503,7 @@ def update_columns( >>> lr.update_columns( # doctest: +SKIP ... "/tmp/products.lance", ... transform=bump_price, - ... output_schema=pa.schema([pa.field("price", pa.float64())]), + ... columns=["price"], ... filter="status = 'active'", ... read_columns=["price"], ... ) @@ -1504,15 +1513,20 @@ def update_columns( ``namespace_impl`` and ``table_id`` to resolve it from a namespace. transform: A callable or :class:`lance.udf.BatchUDF` taking a ``pa.RecordBatch`` of the requested columns and returning a - ``pa.RecordBatch`` whose columns are exactly - ``output_schema.names``. It **must preserve row count and row - order**; reordering cannot be detected and silently writes values to - the wrong rows. Do not filter, sort, join, deduplicate, aggregate or - explode inside the transform. - output_schema: Schema of the transform output. Required, so that column - existence, Arrow type and nullability are checked before any Ray - task starts. Only names, types and nullability are compared; Lance - field ids and field metadata are resolved from the target dataset. + ``pa.RecordBatch`` whose columns are exactly ``columns``. It **must + preserve row count and row order**; reordering cannot be detected + and silently writes values to the wrong rows. Do not filter, sort, + join, deduplicate, aggregate or explode inside the transform. + columns: Names of the existing columns to overwrite. Required, and + checked against the target dataset before any Ray task starts. The + output Arrow type and nullability of each column are taken from the + dataset -- update_columns cannot change them -- and the transform + result is cast to them with ``safe=True``. That rejects out-of-range + integers, truncating time-unit conversions and unparseable strings, + but it does **not** catch float narrowing: returning a ``float64`` + for a ``float32`` column rounds, and overflows to ``inf``, silently. + Produce the column's own type in the transform when precision + matters. filter: A Lance filter expression. Only matching rows get new values; unmatched rows keep their old values. Note that the physical rewrite is still fragment-wide, so a very sparse filter is a poor fit. @@ -1578,7 +1592,7 @@ def update_columns( "(lance-format/lance#6734)." ) - columns, field_ids = _resolve_update_targets(lance_ds, output_schema) + field_ids, output_schema = _resolve_update_targets(lance_ds, columns) projection = _resolve_read_columns(lance_ds, read_columns) fragment_ids = [f.metadata.id for f in lance_ds.get_fragments()] @@ -1586,7 +1600,6 @@ def update_columns( uri, transform, output_schema, - columns, projection, filter, batch_size, @@ -1629,7 +1642,7 @@ def update_columns( raise ValueError( f"Fragment rewrite changed the fragment id: expected {frag_id}, " f"got {fragment_meta.id}" - ) + ) updated_fragments.append(fragment_meta) rows_updated += updated_rows @@ -1680,6 +1693,7 @@ def update_columns( rows_updated=rows_updated, ) + def _validate_write_args( uri: Optional[str], namespace_impl: Optional[str], diff --git a/tests/test_update_columns.py b/tests/test_update_columns.py index 81668387..367bc043 100644 --- a/tests/test_update_columns.py +++ b/tests/test_update_columns.py @@ -18,9 +18,6 @@ def temp_dir(): yield temp_dir -PRICE_SCHEMA = pa.schema([pa.field("price", pa.float64())]) - - def _write_products(path, rows=6, max_rows_per_file=2): """A small multi-fragment dataset: ids 1..rows, alternating status.""" table = pa.table( @@ -81,7 +78,7 @@ def test_namespace_only_updates_columns(self, temp_dir): result = lr.update_columns( transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], namespace_impl="dir", namespace_properties={"root": temp_dir}, @@ -103,7 +100,7 @@ def test_updates_all_rows_across_fragments(self, temp_dir): result = lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], ) @@ -121,7 +118,7 @@ def test_updates_one_fragment_over_multiple_record_batches(self, temp_dir): result = lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], batch_size=2, concurrency=1, @@ -144,7 +141,7 @@ def test_filter_updates_only_matching_rows(self, temp_dir): result = lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], filter="status = 'a'", read_columns=["price"], ) @@ -162,7 +159,7 @@ def test_partial_fragment_coverage_is_allowed(self, temp_dir): result = lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], filter="id <= 2", read_columns=["price"], ) @@ -193,9 +190,7 @@ def bump_both(batch: pa.RecordBatch) -> pa.RecordBatch: result = lr.update_columns( str(path), transform=bump_both, - output_schema=pa.schema( - [pa.field("price", pa.float64()), pa.field("label", pa.string())] - ), + columns=["price", "label"], read_columns=["price", "label"], ) @@ -213,7 +208,7 @@ def price_from_id(batch: pa.RecordBatch) -> pa.RecordBatch: lr.update_columns( str(path), transform=price_from_id, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["id"], ) @@ -228,7 +223,7 @@ def test_no_op_when_filter_matches_nothing(self, temp_dir): result = lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], filter="id > 1000", read_columns=["price"], ) @@ -237,6 +232,7 @@ def test_no_op_when_filter_matches_nothing(self, temp_dir): assert result.version == before assert lance.dataset(str(path)).version == before + class TestTransformContract: def test_transform_does_not_see_metadata_columns(self, temp_dir): path = Path(temp_dir) / "hidden_meta.lance" @@ -251,7 +247,7 @@ def assert_no_metadata(batch: pa.RecordBatch) -> pa.RecordBatch: lr.update_columns( str(path), transform=assert_no_metadata, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], ) @@ -263,15 +259,13 @@ def test_accepts_batch_udf(self, temp_dir): _write_products(path, rows=2, max_rows_per_file=2) udf = BatchUDF( - lambda batch: _record_batch( - {"price": pc.multiply(batch["price"], 4.0)} - ), - output_schema=PRICE_SCHEMA, + lambda batch: _record_batch({"price": pc.multiply(batch["price"], 4.0)}), + output_schema=pa.schema([pa.field("price", pa.float64())]), ) lr.update_columns( str(path), transform=udf, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], ) @@ -312,10 +306,121 @@ def test_rejects_bad_transform_output(self, temp_dir, bad_transform, match): lr.update_columns( str(path), transform=bad_transform, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], ) + def test_transform_output_is_cast_to_the_dataset_type(self, temp_dir): + """The output type is the dataset's, not whatever the transform built.""" + path = Path(temp_dir) / "cast_output.lance" + _write_products(path, rows=2, max_rows_per_file=2) + + def int_prices(batch: pa.RecordBatch) -> pa.RecordBatch: + return _record_batch({"price": pa.array([1, 2], pa.int32())}) + + lr.update_columns( + str(path), + transform=int_prices, + columns=["price"], + read_columns=["price"], + ) + + table = lance.dataset(str(path)).to_table() + assert table.schema.field("price").type == pa.float64() + assert table.to_pydict()["price"] == [1.0, 2.0] + + def test_rejects_out_of_range_transform_output(self, temp_dir): + """The cast is safe=True, so an out-of-range integer raises. + + Note this does *not* generalize to floats: Arrow's safe cast does not + range-check float narrowing, so a float64 result for a float32 column + rounds (and overflows to inf) silently. See the ``columns`` docstring. + """ + path = Path(temp_dir) / "lossy_output.lance" + table = pa.table( + { + "id": pa.array([1, 2], pa.int32()), + "small": pa.array([1, 2], pa.int8()), + } + ) + lance.write_dataset(table, str(path)) + before = _dataset_fingerprint(path) + + def too_big(batch: pa.RecordBatch) -> pa.RecordBatch: + return _record_batch({"small": pa.array([1000, 2000], pa.int32())}) + + with pytest.raises(RuntimeError, match="Integer value 1000"): + lr.update_columns( + str(path), + transform=too_big, + columns=["small"], + read_columns=["small"], + ) + + assert _dataset_fingerprint(path) == before + + def test_rejects_nulls_for_a_non_nullable_column(self, temp_dir): + """A non-nullable target rejects null output rather than writing it.""" + path = Path(temp_dir) / "non_nullable.lance" + schema = pa.schema( + [ + pa.field("id", pa.int32()), + pa.field("price", pa.float64(), nullable=False), + ] + ) + table = pa.table( + { + "id": pa.array([1, 2], pa.int32()), + "price": pa.array([1.0, 2.0], pa.float64()), + }, + schema=schema, + ) + lance.write_dataset(table, str(path)) + before = _dataset_fingerprint(path) + + def nullify(batch: pa.RecordBatch) -> pa.RecordBatch: + return _record_batch({"price": pa.array([1.0, None], pa.float64())}) + + with pytest.raises(RuntimeError, match="non-nullable"): + lr.update_columns( + str(path), + transform=nullify, + columns=["price"], + read_columns=["price"], + ) + + assert _dataset_fingerprint(path) == before + + def test_updates_a_non_nullable_column(self, temp_dir): + """Non-null output for a non-nullable target round-trips normally.""" + path = Path(temp_dir) / "non_nullable_ok.lance" + schema = pa.schema( + [ + pa.field("id", pa.int32()), + pa.field("price", pa.float64(), nullable=False), + ] + ) + table = pa.table( + { + "id": pa.array([1, 2], pa.int32()), + "price": pa.array([1.0, 2.0], pa.float64()), + }, + schema=schema, + ) + lance.write_dataset(table, str(path)) + + lr.update_columns( + str(path), + transform=_double_price, + columns=["price"], + read_columns=["price"], + ) + + got = lance.dataset(str(path)).to_table() + assert got.schema.field("price").nullable is False + assert got.to_pydict()["price"] == [2.0, 4.0] + + class TestPhysicalCorrectness: def test_preserves_row_address_schema_and_field_ids(self, temp_dir): path = Path(temp_dir) / "identity.lance" @@ -334,7 +439,7 @@ def test_preserves_row_address_schema_and_field_ids(self, temp_dir): lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], ) @@ -372,7 +477,7 @@ def files_by_field(ds): lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], ) @@ -391,7 +496,7 @@ def test_time_travel_reads_pre_update_values(self, temp_dir): result = lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], ) @@ -408,7 +513,7 @@ def test_repeated_updates_stack(self, temp_dir): lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], ) @@ -425,7 +530,7 @@ def test_deleted_rows_do_not_misalign_columns(self, temp_dir): lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], ) @@ -433,6 +538,29 @@ def test_deleted_rows_do_not_misalign_columns(self, temp_dir): assert got["id"] == [1, 3, 4, 5, 6] assert got["price"] == [20.0, 60.0, 80.0, 100.0, 120.0] + def test_filter_with_delete_vector_updates_only_matching_live_rows(self, temp_dir): + """A filtered rewrite must preserve deleted and untouched fragment rows.""" + path = Path(temp_dir) / "deleted_filtered.lance" + _write_products(path, rows=6, max_rows_per_file=3) + + # The first fragment contains ids 1..3. Delete one row from it, then + # update only one of its remaining live rows; the second fragment must + # not be rewritten at all. + lance.dataset(str(path)).delete("id = 2") + + result = lr.update_columns( + str(path), + transform=_double_price, + columns=["price"], + filter="id = 1", + read_columns=["price"], + ) + + assert result.rows_updated == 1 + got = lance.dataset(str(path)).to_table().to_pydict() + assert got["id"] == [1, 3, 4, 5, 6] + assert got["price"] == [20.0, 30.0, 40.0, 50.0, 60.0] + class TestTransactionBehavior: def test_commits_a_rewrite_columns_update(self, temp_dir): @@ -446,7 +574,7 @@ def test_commits_a_rewrite_columns_update(self, temp_dir): result = lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], ) @@ -559,7 +687,7 @@ def test_ray_remote_args_are_applied(self, temp_dir): lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], ray_remote_args={"num_cpus": 1}, ) @@ -574,7 +702,7 @@ def test_concurrency_and_batch_size_are_applied(self, temp_dir): result = lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], read_columns=["price"], batch_size=2, ray_remote_args={"num_cpus": 1}, @@ -616,7 +744,7 @@ def append_marker(batch: pa.RecordBatch) -> pa.RecordBatch: result = lr.update_columns( str(path), transform=append_marker, - output_schema=pa.schema([pa.field("tags", list_type)]), + columns=["tags"], read_columns=["tags"], ) @@ -646,7 +774,7 @@ def test_fields_modified_uses_leaf_ids(self, temp_dir): transform=lambda b: _record_batch( {"tags": pa.array([[7]] * b.num_rows, list_type)} ), - output_schema=pa.schema([pa.field("tags", list_type)]), + columns=["tags"], read_columns=["tags"], ) @@ -673,7 +801,7 @@ def test_updates_a_fixed_size_list_column(self, temp_dir): ) } ), - output_schema=pa.schema([pa.field("vec", vec_type)]), + columns=["vec"], read_columns=["vec"], ) @@ -690,45 +818,21 @@ class TestDriverSideRejection: """ @pytest.mark.parametrize( - "output_schema, read_columns, exc, match", + "columns, read_columns, exc, match", [ - ( - pa.schema([pa.field("brand_new", pa.float64())]), - None, - ValueError, - "non-existent column", - ), - ( - pa.schema([pa.field("price", pa.int64())]), - None, - ValueError, - "Type mismatch", - ), - ( - pa.schema([pa.field("price", pa.float64(), nullable=False)]), - None, - ValueError, - "Nullability mismatch", - ), - ( - pa.schema([pa.field("_rowaddr", pa.uint64())]), - None, - ValueError, - "metadata column", - ), - ( - pa.schema([pa.field("meta.price", pa.float64())]), - None, - ValueError, - "Nested field path", - ), - (pa.schema([]), None, ValueError, "at least one column"), - (PRICE_SCHEMA, ["price", "_rowaddr"], ValueError, "read_columns"), - (PRICE_SCHEMA, ["nope"], ValueError, "do not exist"), + (["brand_new"], None, ValueError, "non-existent column"), + (["_rowaddr"], None, ValueError, "metadata column"), + (["meta.price"], None, ValueError, "Nested field path"), + (["price", "price"], None, ValueError, "Duplicate column"), + ([pa.field("price", pa.float64())], None, TypeError, "as str"), + ("price", None, TypeError, "not a bare string"), + ([], None, ValueError, "at least one column"), + (["price"], ["price", "_rowaddr"], ValueError, "read_columns"), + (["price"], ["nope"], ValueError, "do not exist"), ], ) def test_rejects_without_touching_the_dataset( - self, temp_dir, output_schema, read_columns, exc, match + self, temp_dir, columns, read_columns, exc, match ): path = Path(temp_dir) / "untouched.lance" _write_products(path, rows=4, max_rows_per_file=2) @@ -738,7 +842,7 @@ def test_rejects_without_touching_the_dataset( lr.update_columns( str(path), transform=_double_price, - output_schema=output_schema, + columns=columns, read_columns=read_columns, ) @@ -759,7 +863,7 @@ def test_stable_row_ids_rejected_without_touching_the_dataset(self, temp_dir): lr.update_columns( str(path), transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], ) assert _dataset_fingerprint(path) == before @@ -781,14 +885,14 @@ def test_rejects_struct_target_column(self, temp_dir): lr.update_columns( str(path), transform=lambda b: b, - output_schema=pa.schema([pa.field("meta", struct_type)]), + columns=["meta"], ) def test_requires_uri_or_namespace(self, temp_dir): with pytest.raises(ValueError, match="Must provide either 'uri'"): lr.update_columns( transform=_double_price, - output_schema=PRICE_SCHEMA, + columns=["price"], ) @@ -833,7 +937,7 @@ def payload_size(batch: pa.RecordBatch) -> pa.RecordBatch: lr.update_columns( str(path), transform=payload_size, - output_schema=pa.schema([pa.field("size", pa.int64())]), + columns=["size"], read_columns=["payload"], ) @@ -856,7 +960,7 @@ def record_projection(batch: pa.RecordBatch) -> pa.RecordBatch: lr.update_columns( str(path), transform=record_projection, - output_schema=pa.schema([pa.field("size", pa.int64())]), + columns=["size"], ) # read_columns=None expands to the non-blob columns only: id + size. @@ -871,7 +975,7 @@ def test_rejects_blob_output(self, temp_dir): lr.update_columns( str(path), transform=lambda b: b, - output_schema=pa.schema([pa.field("payload", pa.large_binary())]), + columns=["payload"], read_columns=["payload"], ) @@ -902,8 +1006,6 @@ def test_blob_v2_column_is_readable_and_excluded_by_default(self, temp_dir): data_storage_version="2.2", # blob v2 requires file version >= 2.2 ) - size_schema = pa.schema([pa.field("size", pa.int64())]) - # Blob v2 is the case the default projection exists to protect: it must # not be pulled in unless asked for. lr.update_columns( @@ -911,7 +1013,7 @@ def test_blob_v2_column_is_readable_and_excluded_by_default(self, temp_dir): transform=lambda b: _record_batch( {"size": pa.array([len(b.column_names)] * b.num_rows, pa.int64())} ), - output_schema=size_schema, + columns=["size"], ) assert lance.dataset(str(path)).to_table().to_pydict()["size"] == [2, 2] @@ -925,7 +1027,7 @@ def payload_size(batch: pa.RecordBatch) -> pa.RecordBatch: lr.update_columns( str(path), transform=payload_size, - output_schema=size_schema, + columns=["size"], read_columns=["payload"], ) assert lance.dataset(str(path)).to_table().to_pydict()["size"] == [2, 4] @@ -939,7 +1041,7 @@ def test_unrelated_blob_column_does_not_block_updates(self, temp_dir): transform=lambda b: _record_batch( {"size": pc.cast(pc.multiply(b["id"], 100), pa.int64())} ), - output_schema=pa.schema([pa.field("size", pa.int64())]), + columns=["size"], read_columns=["id"], ) From bab245516c8a4971108849d76028f5aeaba950c9 Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Tue, 28 Jul 2026 11:51:12 +0900 Subject: [PATCH 10/11] chore: satisfy lint on distributed update columns Assert the returned rows_updated in the two tests that discarded it, which also clears ruff F841, and reflow the docs example so `ruff format --check` is clean for the files this branch touches. Claude-Session: https://claude.ai/code/session_016cKDCcCvDW9YVxT7Ea7Qtb --- docs/src/data-evolution.md | 4 +--- tests/test_update_columns.py | 2 ++ 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/src/data-evolution.md b/docs/src/data-evolution.md index ad19c807..7f93a077 100644 --- a/docs/src/data-evolution.md +++ b/docs/src/data-evolution.md @@ -40,9 +40,7 @@ import pyarrow.compute as pc def increase_price(batch: pa.RecordBatch) -> pa.RecordBatch: - return pa.RecordBatch.from_pydict( - {"price": pc.multiply(batch["price"], 1.1)} - ) + return pa.RecordBatch.from_pydict({"price": pc.multiply(batch["price"], 1.1)}) result = update_columns( diff --git a/tests/test_update_columns.py b/tests/test_update_columns.py index 367bc043..d64e1cbf 100644 --- a/tests/test_update_columns.py +++ b/tests/test_update_columns.py @@ -194,6 +194,7 @@ def bump_both(batch: pa.RecordBatch) -> pa.RecordBatch: read_columns=["price", "label"], ) + assert result.rows_updated == 4 got = lance.dataset(str(path)).to_table().to_pydict() assert got["price"] == [10.0, 20.0, 30.0, 40.0] assert got["label"] == ["w!", "x!", "y!", "z!"] @@ -709,6 +710,7 @@ def test_concurrency_and_batch_size_are_applied(self, temp_dir): concurrency=1, ) + assert result.rows_updated == 4 got = lance.dataset(str(path)).to_table().to_pydict() assert got["price"] == [20.0, 40.0, 60.0, 80.0] From 240c8a98c3e0d963d85ffb52604862f964b403ec Mon Sep 17 00:00:00 2001 From: fanng <“fanng@apache.org”> Date: Tue, 28 Jul 2026 12:29:42 +0900 Subject: [PATCH 11/11] fix: make update_columns transforms picklable by Ray workers Ray workers cannot import the test module -- it imports pytest, which is absent from the worker environment -- so every transform that referenced a test-module global failed to deserialize. Move the two helpers that cross the Ray boundary into _ray_test_support, which exists precisely to be importable from workers, and replace the bare asserts inside two transforms with explicit raises: pytest rewrites `assert` into calls on `_pytest`, which the pickled closure then drags along. Claude-Session: https://claude.ai/code/session_016cKDCcCvDW9YVxT7Ea7Qtb --- tests/_ray_test_support.py | 25 +++++++++++++++++++++++++ tests/test_update_columns.py | 32 +++++++++++++++++++------------- 2 files changed, 44 insertions(+), 13 deletions(-) diff --git a/tests/_ray_test_support.py b/tests/_ray_test_support.py index 75a5717b..a2f73798 100644 --- a/tests/_ray_test_support.py +++ b/tests/_ray_test_support.py @@ -114,3 +114,28 @@ def setup_worker() -> None: """ patch_memory_profiler() patch_psutil_for_containers() + + +# --------------------------------------------------------------------------- +# Transform helpers shared by tests that ship a callable to a Ray worker. +# +# These must live here rather than in a ``test_*.py`` module: cloudpickle +# serializes a nested transform by value but its module-level references *by +# reference*, so a transform calling a helper defined next to the tests forces +# the worker to import that test module -- and with it ``pytest``, which is not +# available in the worker environment. +# --------------------------------------------------------------------------- + + +def record_batch(data, schema=None): + """Build a ``pa.RecordBatch`` from a dict, importable from Ray workers.""" + import pyarrow as pa + + return pa.RecordBatch.from_pydict(data, schema=schema) + + +def double_price(batch): + """Double the ``price`` column of a batch.""" + import pyarrow.compute as pc + + return record_batch({"price": pc.multiply(batch["price"], 2.0)}) diff --git a/tests/test_update_columns.py b/tests/test_update_columns.py index d64e1cbf..460b27e9 100644 --- a/tests/test_update_columns.py +++ b/tests/test_update_columns.py @@ -11,6 +11,11 @@ import ray from lance.udf import BatchUDF +# Ray workers cannot import this module (it imports pytest), so any helper a +# transform references must come from a module they *can* import. +from _ray_test_support import double_price as _double_price +from _ray_test_support import record_batch as _record_batch + @pytest.fixture def temp_dir(): @@ -33,14 +38,6 @@ def _write_products(path, rows=6, max_rows_per_file=2): return table -def _record_batch(data, schema=None) -> pa.RecordBatch: - return pa.RecordBatch.from_pydict(data, schema=schema) - - -def _double_price(batch: pa.RecordBatch) -> pa.RecordBatch: - return _record_batch({"price": pc.multiply(batch["price"], 2.0)}) - - def _dataset_fingerprint(path): """Version plus every data file, to prove nothing was written.""" ds = lance.dataset(str(path)) @@ -240,9 +237,16 @@ def test_transform_does_not_see_metadata_columns(self, temp_dir): _write_products(path, rows=4) def assert_no_metadata(batch: pa.RecordBatch) -> pa.RecordBatch: - assert isinstance(batch, pa.RecordBatch) - for hidden in ("_rowaddr", "_fragid", "_rowid"): - assert hidden not in batch.column_names + # Raise rather than assert: pytest rewrites `assert` into calls on + # `_pytest`, which a Ray worker cannot import, so a bare assert here + # fails to deserialize instead of running. + if not isinstance(batch, pa.RecordBatch): + raise TypeError(f"expected RecordBatch, got {type(batch).__name__}") + seen = [ + c for c in ("_rowaddr", "_fragid", "_rowid") if c in batch.column_names + ] + if seen: + raise ValueError(f"transform saw metadata columns {seen}") return _record_batch({"price": pc.multiply(batch["price"], 3.0)}) lr.update_columns( @@ -952,8 +956,10 @@ def test_default_projection_excludes_blob_columns(self, temp_dir): def record_projection(batch: pa.RecordBatch) -> pa.RecordBatch: # The transform runs in a Ray worker, so the observation has to - # travel back through the data itself. - assert "payload" not in batch.column_names + # travel back through the data itself. Raise rather than assert: + # a bare assert would pull `_pytest` into the pickled closure. + if "payload" in batch.column_names: + raise ValueError("blob column was projected by default") width = len(batch.column_names) return _record_batch( {"size": pa.array([width] * batch.num_rows, pa.int64())}