feat: distributed update_columns for overwriting existing columns

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.

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}")
-```
+```

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",

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