diff --git a/docs/src/data-evolution.md b/docs/src/data-evolution.md index 25102c9e..7f93a077 100644 --- a/docs/src/data-evolution.md +++ b/docs/src/data-evolution.md @@ -30,3 +30,77 @@ 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, + columns=["price"], + 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 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. + +**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. +- `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 + 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. 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 + 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 +``` diff --git a/lance_ray/__init__.py b/lance_ray/__init__.py index cea756a3..816ac9ad 100644 --- a/lance_ray/__init__.py +++ b/lance_ray/__init__.py @@ -17,10 +17,12 @@ from .fragment import LanceFragmentWriter from .index import create_index, create_scalar_index, optimize_indices from .io import ( + UpdateColumnsResult, 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 +39,8 @@ "add_columns", "add_columns_from", "merge_columns_from", + "update_columns", + "UpdateColumnsResult", "create_scalar_index", "create_index", "optimize_indices", diff --git a/lance_ray/datasource.py b/lance_ray/datasource.py index 436aad89..a875f7d6 100644 --- a/lance_ray/datasource.py +++ b/lance_ray/datasource.py @@ -288,18 +288,57 @@ 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", 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. @@ -319,41 +358,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 @@ -383,19 +388,20 @@ def _is_blob_field(f: pa.Field) -> Optional[str]: 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 @@ -544,4 +550,9 @@ def _is_blob_field(f: pa.Field) -> Optional[str]: 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 df10171a..7ed1751f 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 +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 @@ -15,7 +18,7 @@ from ray.util.multiprocessing import Pool from .datasink import LanceDatasink -from .datasource import LanceDatasource +from .datasource import LanceDatasource, _read_fragments, blob_field_kind from .fragment import prepare_fragment_write_options from .utils import ( get_namespace_kwargs, @@ -36,6 +39,34 @@ | Callable[[pa.RecordBatch], pa.RecordBatch] ) +logger = logging.getLogger(__name__) + +#: Internal transform contract for :func:`update_columns`. +#: +#: 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 the requested ``columns``, with the +#: same number of rows **in the same order**. +_UpdateColumnsTransform = BatchUDF | Callable[[pa.RecordBatch], pa.RecordBatch] + +_METADATA_COLUMNS = frozenset({"_rowaddr", "_fragid", "_rowid"}) + +# Lance's built-in commit performs a conflict-checked rebase on every attempt. +# Bound retries only to limit tail latency. +_UPDATE_COMMIT_MAX_RETRIES = 5 + + +@dataclass(frozen=True) +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 ``rows_updated`` is zero. + """ + + version: int + rows_updated: int + def read_lance( uri: Optional[str] = None, @@ -1141,6 +1172,528 @@ 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, + 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 *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 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( + "'columns' must name 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) + + fields: list[pa.Field] = [] + field_ids: list[int] = [] + seen: set[str] = set() + + 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 '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 'columns'." + ) + 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." + ) + + lance_field = lance_schema.field(name) + if lance_field is None: + raise ValueError( + f"Column '{name}' has no Lance field id; cannot update it." + ) + fields.append(target_field) + field_ids.extend(_leaf_field_ids(lance_field)) + + return tuple(field_ids), pa.schema(fields) + + +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 _apply_update_transform( + transform: "_UpdateColumnsTransform", + output_schema: pa.Schema, +) -> Callable[[pa.RecordBatch], pa.RecordBatch]: + """Build the per-batch transform used by a fragment-local worker.""" + 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: + rowaddr = batch.column("_rowaddr") + user_batch = batch.drop_columns( + [c for c in _METADATA_COLUMNS if c in batch.column_names] + ) + + result = transform(user_batch) + + if not isinstance(result, pa.RecordBatch): + raise TypeError( + f"transform must return pa.RecordBatch, got {type(result).__name__}." + ) + 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 {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.schema.names) + if actual != expected: + unexpected = sorted(actual - expected) + missing = sorted(expected - actual) + raise ValueError( + "transform output columns must match 'columns' exactly. " + f"Unexpected: {unexpected}; missing: {missing}." + ) + + # 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] + + for name in non_nullable: + if normalized.column(name).null_count: + raise ValueError( + f"transform produced nulls for non-nullable column '{name}'." + ) + + return normalized.append_column("_rowaddr", rowaddr) + + return _wrapped + + +def _handle_update_fragment( + uri: str, + transform: "_UpdateColumnsTransform", + output_schema: pa.Schema, + 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) + + 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. Keep + # its result as a RecordBatch stream until Lance consumes it. + scanned_batches = _read_fragments( + [fragment_id], + lance_ds, + { + "columns": projection, + "filter": filter, + "batch_size": batch_size, + }, + with_metadata=True, + as_record_batches=True, + ) + 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. 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 + + # 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 + 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( + reader_schema, updated_batches() + ) + fragment_meta, fields_modified = fragment.update_columns( + update_reader, + left_on="_rowaddr", + right_on="_rowaddr", + ) + return ( + fragment_id, + pickle.dumps(fragment_meta), + pickle.dumps(list(fields_modified)), + rows_updated, + ) + + return _update_fragment + + +def update_columns( + uri: Optional[str] = None, + *, + transform: "_UpdateColumnsTransform", + columns: Sequence[str], + filter: Optional[str] = None, + read_columns: Optional[Sequence[str]] = 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, + namespace_properties: Optional[Mapping[str, str]] = None, + table_id: Optional[Sequence[str]] = None, +) -> UpdateColumnsResult: + """Overwrite existing columns of a Lance dataset with Ray. + + 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 + >>> import pyarrow as pa + >>> import pyarrow.compute as pc + >>> 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, + ... columns=["price"], + ... 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 or :class:`lance.udf.BatchUDF` taking a + ``pa.RecordBatch`` of the requested columns and returning a + ``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. + 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. 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. + 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. + 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``. + + """ + 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)." + ) + + 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()] + worker = _handle_update_fragment( + uri, + transform, + output_schema, + 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 = [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( + version=resolved_read_version, + rows_updated=0, + ) + + updated_fragments = [] + seen_frag_ids: set[int] = set() + rows_updated = 0 + observed_field_ids: Optional[list[int]] = None + + 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 worker output") + seen_frag_ids.add(frag_id) + + 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 += updated_rows + + 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: + 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. + + logger.info( + "Committing update_columns: columns=%s read_version=%s fragments=%s", + list(columns), + resolved_read_version, + len(updated_fragments), + ) + committed = LanceDataset.commit( + uri, + op, + read_version=resolved_read_version, + max_retries=_UPDATE_COMMIT_MAX_RETRIES, + storage_options=dict(resolved_storage_options), + **namespace_kwargs, + ) + + return UpdateColumnsResult( + version=committed.version, + rows_updated=rows_updated, + ) + + def _validate_write_args( uri: Optional[str], namespace_impl: Optional[str], 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 new file mode 100644 index 00000000..460b27e9 --- /dev/null +++ b/tests/test_update_columns.py @@ -0,0 +1,1066 @@ +"""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.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(): + with tempfile.TemporaryDirectory() as temp_dir: + yield temp_dir + + +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 _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, + columns=["price"], + read_columns=["price"], + namespace_impl="dir", + namespace_properties={"root": temp_dir}, + table_id=table_id, + ) + + 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, + columns=["price"], + read_columns=["price"], + ) + + assert result.version == 2 + assert result.rows_updated == 6 + + 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_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, + columns=["price"], + 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) + + result = lr.update_columns( + str(path), + transform=_double_price, + columns=["price"], + 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, + columns=["price"], + filter="id <= 2", + read_columns=["price"], + ) + + assert result.rows_updated == 2 + 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.RecordBatch) -> pa.RecordBatch: + return _record_batch( + { + "price": pc.multiply(batch["price"], 10.0), + "label": pc.binary_join_element_wise(batch["label"], "!", ""), + } + ) + + result = lr.update_columns( + str(path), + transform=bump_both, + columns=["price", "label"], + 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!"] + + 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.RecordBatch) -> pa.RecordBatch: + return _record_batch({"price": pc.cast(batch["id"], pa.float64())}) + + lr.update_columns( + str(path), + transform=price_from_id, + columns=["price"], + 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, + columns=["price"], + filter="id > 1000", + read_columns=["price"], + ) + + assert result.rows_updated == 0 + 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" + _write_products(path, rows=4) + + def assert_no_metadata(batch: pa.RecordBatch) -> pa.RecordBatch: + # 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( + str(path), + transform=assert_no_metadata, + columns=["price"], + read_columns=["price"], + ) + + 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=pa.schema([pa.field("price", pa.float64())]), + ) + lr.update_columns( + str(path), + transform=udf, + columns=["price"], + 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: _record_batch( + {"price": pc.multiply(b["price"], 2.0), "extra": b["price"]} + ), + "Unexpected: \\['extra'\\]", + ), + (lambda b: _record_batch({"nope": b["price"]}), "missing: \\['price'\\]"), + ( + lambda b: _record_batch({"price": [1.0] * (b.num_rows * 2)}), + "changed the row count", + ), + ( + 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.RecordBatch, 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, + 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" + _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, + columns=["price"], + 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, + columns=["price"], + 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) + read_version = lance.dataset(str(path)).version + + result = lr.update_columns( + str(path), + transform=_double_price, + columns=["price"], + read_columns=["price"], + ) + + 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] + + 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, + columns=["price"], + 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, + columns=["price"], + 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] + + 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): + path = Path(temp_dir) / "txn_shape.lance" + _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( + str(path), + transform=_double_price, + columns=["price"], + 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 == 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", + ) + with pytest.raises(OSError, match="[Cc]ommit conflict"): + lance.LanceDataset.commit( + str(path), + Transaction(read_version=stale_version, operation=op), + max_retries=5, + ) + + 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] + # A caller-fixed UUID also survives Lance's conflict-checked rebase. + assert committed.read_transaction(committed.version).uuid == txn_uuid + + +class TestResourceOptions: + """The fragment-local Ray task must receive the tuning parameters.""" + + 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( + str(path), + transform=_double_price, + columns=["price"], + read_columns=["price"], + 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_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( + str(path), + transform=_double_price, + columns=["price"], + read_columns=["price"], + batch_size=2, + ray_remote_args={"num_cpus": 1}, + 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] + + +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.RecordBatch) -> pa.RecordBatch: + return _record_batch( + { + "tags": pa.array( + [v + [99] for v in batch["tags"].to_pylist()], list_type + ) + } + ) + + result = lr.update_columns( + str(path), + transform=append_marker, + columns=["tags"], + 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: _record_batch( + {"tags": pa.array([[7]] * b.num_rows, list_type)} + ), + columns=["tags"], + 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: _record_batch( + { + "vec": pa.FixedSizeListArray.from_arrays( + pa.array([0.5] * (b.num_rows * 2), pa.float32()), 2 + ) + } + ), + columns=["vec"], + 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( + "columns, read_columns, exc, match", + [ + (["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, columns, 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, + columns=columns, + 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, + columns=["price"], + ) + + assert _dataset_fingerprint(path) == before + + +class TestRejectedScenarios: + 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, + 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, + columns=["price"], + ) + + +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.RecordBatch) -> pa.RecordBatch: + sizes = [ + len(v) if v is not None else 0 for v in batch["payload"].to_pylist() + ] + return _record_batch({"size": pa.array(sizes, pa.int64())}) + + lr.update_columns( + str(path), + transform=payload_size, + columns=["size"], + 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.RecordBatch) -> pa.RecordBatch: + # The transform runs in a Ray worker, so the observation has to + # 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())} + ) + + lr.update_columns( + str(path), + transform=record_projection, + columns=["size"], + ) + + # 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, + columns=["payload"], + 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 + ) + + # 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: _record_batch( + {"size": pa.array([len(b.column_names)] * b.num_rows, pa.int64())} + ), + columns=["size"], + ) + assert lance.dataset(str(path)).to_table().to_pydict()["size"] == [2, 2] + + # ... but it is readable when explicitly requested. + 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 _record_batch({"size": pa.array(sizes, pa.int64())}) + + lr.update_columns( + str(path), + transform=payload_size, + columns=["size"], + 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: _record_batch( + {"size": pc.cast(pc.multiply(b["id"], 100), pa.int64())} + ), + columns=["size"], + 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"]