fix: propagate update_columns offsets and partial last_updated for RewriteColumns (#6650)

## Summary

* Fixes #6505
* `FileFragment::update_columns` returns `Result<(Fragment, Vec<u32>)>`
(unchanged
public shape). `update_columns_with_offsets` returns
`FragmentUpdateColumnsResult`
(fragment, `fields_modified`, `matched_offsets: RoaringBitmap`) for
callers that need
  physical row indices for stable row-id metadata.
* `HashJoiner::matched_join_rows` — boolean mask for hash hits; used by
  `update_columns_with_offsets` and covered by `test_matched_join_rows`.
* `Operation::Update`: optional `updated_fragment_offsets:
Option<UpdatedFragmentOffsets>`
where `UpdatedFragmentOffsets` wraps `HashMap<u64, RoaringBitmap>`
(newtype with
`Default`, `PartialEq`, manual `DeepSizeOf`). `None` means the caller
did not supply
  offsets.
* Proto (`transaction.proto`): backward-compatible `map<uint64,
UInt32List> updated_fragment_offsets = 9`
  on `Update`; serde round-trip preserves semantics.
* `build_manifest`: when stable row IDs are enabled, `update_mode ==
RewriteColumns`,
and `Some(UpdatedFragmentOffsets(..))` includes a non-empty bitmap for a
fragment,
calls `refresh_row_latest_update_meta_for_partial_frag_rewrite_cols` for
those offsets
  only — unmatched rows and untouched fragments are left unchanged.
* JNI / Java: `FragmentUpdateResult` includes matched row offsets; the
2-arg constructor
`(FragmentMetadata, long[])` delegates to the 3-arg form with an empty
offset array for
  compatibility. JNI uses `update_columns_with_offsets`.
* Python: `update_columns` binding correctly destructures the
`(Fragment, Vec<u32>)` tuple.


## Root cause

For `Operation::Update` with `RewriteColumns`, commits could advance the
dataset version
without advancing `_row_last_updated_at_version` for the rows that were
actually rewritten.
`update_columns` did not report which physical offsets matched, and
`build_manifest` had no
per-fragment offset map to drive the partial refresh. Without that
information the transaction
layer cannot distinguish which rows changed, so the version metadata is
not updated.


## Implementation notes

* `RoaringBitmap` iteration is ascending and duplicate-free; redundant
`sort` / `dedup`
  when building proto lists or offset vectors from bitmaps were removed.
* Call sites that do not populate offsets use `updated_fragment_offsets:
None`.


## Why the protobuf field exists

lance-spark passes `Transaction` through JNI as a protobuf blob: Java
builds a `Transaction`
proto, Rust deserializes it and runs `build_manifest`. Without
`updated_fragment_offsets` on
the wire, the decoded `Operation::Update` would always have
`updated_fragment_offsets: None`
even when matched offsets were computed on the JVM side, and the partial
refresh in
`build_manifest` would silently do nothing.


## Test plan

* `cargo test -p lance test_matched_join_rows` —
`HashJoiner::matched_join_rows`.
* `cargo test -p lance
test_build_manifest_partial_last_updated_rewrite_columns_stable_row_ids`
  — `Dataset::commit` -> `build_manifest`: two fragments, partial
`update_columns_with_offsets`, `Operation::Update` with `RewriteColumns`
and an offset
  map; asserts matched vs unmatched vs untouched row version metadata.
* `cargo test -p lance test_fragment_update` — fragment path with
`Operation::Update` and
  offsets.
* `cargo test -p lance --tests` (or at least `cargo check -p lance
--tests`) and
  `cargo check --manifest-path java/lance-jni/Cargo.toml`.

The `pylance` crate is excluded from the root workspace; validate Python
bindings in the
usual `maturin` / CI flow if you touch `python/`.


## Compatibility

* Rust: `update_columns` signature unchanged;
`update_columns_with_offsets` is additive.
* Java: 2-arg `FragmentUpdateResult` constructor preserved.
* Proto: field 9; older clients ignore unknown fields.

---------

Co-authored-by: Jing chen He <jingh@adobe.com>

Author: jerryjch

java/lance-jni/src/fragment.rs

@@ -48,6 +48,8 @@ pub(crate) struct FragmentMergeResult {
 pub(crate) struct FragmentUpdateResult {
     updated_fragment: Fragment,
     fields_modified: Vec<u32>,
+    /// Physical row offsets that received column updates (from `_rowaddr` low bits).
+    updated_row_offsets: Vec<i64>,
 }
 
 //////////////////
@@ -490,11 +492,13 @@ fn inner_update_column<'local>(
     let reader = unsafe { ArrowArrayStreamReader::from_raw(stream_ptr) }?;
     let left_on_str: String = left_on.extract(env)?;
     let right_on_str: String = right_on.extract(env)?;
-    let (updated_fragment, fields_modified) =
-        RT.block_on(fragment.update_columns(reader, &left_on_str, &right_on_str))?;
+    let r =
+        RT.block_on(fragment.update_columns_with_offsets(reader, &left_on_str, &right_on_str))?;
+    let updated_row_offsets: Vec<i64> = r.matched_offsets.iter().map(|o| o as i64).collect();
     let result = FragmentUpdateResult {
-        updated_fragment,
-        fields_modified,
+        updated_fragment: r.fragment,
+        fields_modified: r.fields_modified,
+        updated_row_offsets,
     };
     result.into_java(env)
 }
@@ -542,7 +546,7 @@ const FRAGMENT_MERGE_RESULT_CLASS: &str = "org/lance/fragment/FragmentMergeResul
 const FRAGMENT_MERGE_RESULT_CONSTRUCTOR_SIG: &str =
     "(Lorg/lance/FragmentMetadata;Lorg/lance/schema/LanceSchema;)V";
 const FRAGMENT_UPDATE_RESULT_CLASS: &str = "org/lance/fragment/FragmentUpdateResult";
-const FRAGMENT_UPDATE_RESULT_CONSTRUCTOR_SIG: &str = "(Lorg/lance/FragmentMetadata;[J)V";
+const FRAGMENT_UPDATE_RESULT_CONSTRUCTOR_SIG: &str = "(Lorg/lance/FragmentMetadata;[J[J)V";
 
 impl IntoJava for &FragmentMergeResult {
     fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result<JObject<'a>> {
@@ -563,12 +567,14 @@ impl IntoJava for &FragmentUpdateResult {
     fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result<JObject<'a>> {
         let java_updated_fragment = self.updated_fragment.into_java(env)?;
         let java_fields_modified = JLance(self.fields_modified.clone()).into_java(env)?;
+        let java_updated_row_offsets = JLance(self.updated_row_offsets.clone()).into_java(env)?;
         Ok(env.new_object(
             FRAGMENT_UPDATE_RESULT_CLASS,
             FRAGMENT_UPDATE_RESULT_CONSTRUCTOR_SIG,
             &[
                 JValueGen::Object(&java_updated_fragment),
                 JValueGen::Object(&java_fields_modified),
+                JValueGen::Object(&java_updated_row_offsets),
             ],
         )?)
     }

java/lance-jni/src/traits.rs

@@ -202,6 +202,14 @@ impl IntoJava for JLance<Vec<u32>> {
     }
 }
 
+impl IntoJava for JLance<Vec<i64>> {
+    fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result<JObject<'a>> {
+        let arr = env.new_long_array(self.0.len() as i32)?;
+        env.set_long_array_region(&arr, 0, &self.0)?;
+        Ok(arr.into())
+    }
+}
+
 impl IntoJava for JLance<usize> {
     fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result<JObject<'a>> {
         Ok(env.new_object("java/lang/Long", "(J)V", &[JValueGen::Long(self.0 as i64)])?)

java/lance-jni/src/transaction.rs

@@ -433,6 +433,7 @@ fn convert_to_java_operation_inner<'local>(
             fields_for_preserving_frag_bitmap,
             update_mode,
             inserted_rows_filter: _,
+            updated_fragment_offsets: _,
         } => {
             let removed_ids: Vec<JLance<i64>> = removed_fragment_ids
                 .iter()
@@ -1222,6 +1223,7 @@ fn convert_to_rust_operation(
                 fields_for_preserving_frag_bitmap,
                 update_mode,
                 inserted_rows_filter: None,
+                updated_fragment_offsets: None,
             }
         }
         "DataReplacement" => {

java/src/main/java/org/lance/fragment/FragmentUpdateResult.java

@@ -26,9 +26,19 @@ public class FragmentUpdateResult {
   private final FragmentMetadata updatedFragment;
   private final long[] fieldsModified;
 
+  /** Local physical row offsets within the fragment that received updates (see RowAddress). */
+  private final long[] updatedRowOffsets;
+
+  /** Two-argument form for callers that do not track per-row offsets; offsets default to empty. */
   public FragmentUpdateResult(FragmentMetadata updatedFragment, long[] updatedFieldIds) {
+    this(updatedFragment, updatedFieldIds, new long[0]);
+  }
+
+  public FragmentUpdateResult(
+      FragmentMetadata updatedFragment, long[] updatedFieldIds, long[] updatedRowOffsets) {
     this.updatedFragment = updatedFragment;
     this.fieldsModified = updatedFieldIds;
+    this.updatedRowOffsets = updatedRowOffsets;
   }
 
   public FragmentMetadata getUpdatedFragment() {
@@ -39,11 +49,17 @@ public long[] getFieldsModified() {
     return fieldsModified;
   }
 
+  /** Physical row offsets (0-based within the fragment) whose columns were rewritten. */
+  public long[] getUpdatedRowOffsets() {
+    return updatedRowOffsets;
+  }
+
   @Override
   public String toString() {
     return MoreObjects.toStringHelper(this)
         .add("fragmentMetadata", updatedFragment)
         .add("updatedFieldIds", fieldsModified)
+        .add("updatedRowOffsets", updatedRowOffsets)
         .toString();
   }
 }

protos/transaction.proto

@@ -227,6 +227,11 @@ message Transaction {
     }
   }
 
+  // Serialized as sorted distinct local physical row offsets within the fragment (0-based).
+  message UInt32List {
+    repeated uint32 values = 1;
+  }
+
   // An operation that updates rows but does not add or remove rows.
   message Update {
     // The fragments that have been removed. These are fragments where all rows
@@ -248,6 +253,8 @@ message Transaction {
     // Filter for checking existence of keys in newly inserted rows, used for conflict detection.
     // Only tracks keys from INSERT operations during merge insert, not updates.
     optional KeyExistenceFilter inserted_rows = 8;
+    // Per-fragment physical row offsets that matched an update_columns hash join (RewriteColumns).
+    map<uint64, UInt32List> updated_fragment_offsets = 9;
   }
 
   // The mode of update operation

python/src/transaction.rs

@@ -299,6 +299,7 @@ impl FromPyObject<'_, '_> for PyLance<Operation> {
                     fields_for_preserving_frag_bitmap,
                     update_mode,
                     inserted_rows_filter: None,
+                    updated_fragment_offsets: None,
                 };
                 Ok(Self(op))
             }

rust/lance/src/dataset/fragment.rs

@@ -13,8 +13,9 @@ use std::sync::Arc;
 
 use arrow::compute::concat_batches;
 use arrow_array::cast::as_primitive_array;
+use arrow_array::types::UInt64Type;
 use arrow_array::{
-    RecordBatch, RecordBatchReader, StructArray, UInt32Array, UInt64Array, new_null_array,
+    Array, RecordBatch, RecordBatchReader, StructArray, UInt32Array, UInt64Array, new_null_array,
 };
 use arrow_schema::Schema as ArrowSchema;
 use datafusion::logical_expr::Expr;
@@ -23,6 +24,7 @@ use futures::future::try_join_all;
 use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt, join, stream};
 use lance_arrow::{RecordBatchExt, SchemaExt};
 use lance_core::datatypes::{OnMissing, OnTypeMismatch, SchemaCompareOptions};
+use lance_core::utils::address::RowAddress;
 use lance_core::utils::deletion::DeletionVector;
 use lance_core::utils::tokio::get_num_compute_intensive_cpus;
 use lance_core::{Error, Result, cache::CacheKey, datatypes::Schema};
@@ -48,6 +50,7 @@ use lance_table::utils::stream::{
     ReadBatchFutStream, ReadBatchTask, ReadBatchTaskStream, RowIdAndDeletesConfig,
     wrap_with_row_id_and_delete,
 };
+use roaring::RoaringBitmap;
 
 use self::write::FragmentCreateBuilder;
 
@@ -61,6 +64,16 @@ use crate::dataset::Dataset;
 use crate::dataset::fragment::session::FragmentSession;
 use crate::io::deletion::read_dataset_deletion_file;
 
+/// Result of [`FileFragment::update_columns_with_offsets`]: updated fragment metadata, modified field ids,
+/// and physical row offsets that matched the join (for stable row-id version metadata).
+#[derive(Debug, Clone)]
+pub struct FragmentUpdateColumnsResult {
+    pub fragment: Fragment,
+    pub fields_modified: Vec<u32>,
+    /// Physical row offsets (0-based within this fragment) whose columns were rewritten from the right-hand stream.
+    pub matched_offsets: RoaringBitmap,
+}
+
 /// A Fragment of a Lance [`Dataset`].
 ///
 /// The interface is modeled after `pyarrow.dataset.Fragment`.
@@ -1602,12 +1615,27 @@ impl FileFragment {
         Ok(self)
     }
 
+    /// Same as [`Self::update_columns_with_offsets`] but discards the matched row offsets.
+    /// Use [`Self::update_columns_with_offsets`] if you need per-row version metadata for stable row IDs.
     pub async fn update_columns(
         &mut self,
         right_stream: impl RecordBatchReader + Send + 'static,
         left_on: &str,
         right_on: &str,
     ) -> Result<(Fragment, Vec<u32>)> {
+        let r = self
+            .update_columns_with_offsets(right_stream, left_on, right_on)
+            .await?;
+        Ok((r.fragment, r.fields_modified))
+    }
+
+    /// Same operation as [`Self::update_columns`], and also returns matched physical row offsets for stable row IDs.
+    pub async fn update_columns_with_offsets(
+        &mut self,
+        right_stream: impl RecordBatchReader + Send + 'static,
+        left_on: &str,
+        right_on: &str,
+    ) -> Result<FragmentUpdateColumnsResult> {
         if self.schema().field(left_on).is_none() && left_on != ROW_ID && left_on != ROW_ADDR {
             return Err(Error::invalid_input(format!(
                 "Column {} does not exist in the left side fragment",
@@ -1647,18 +1675,44 @@ impl FileFragment {
         let mut read_columns: Vec<String> =
             write_schema.fields.iter().map(|f| f.name.clone()).collect();
         read_columns.push(left_on.to_string());
+        // Physical positions for matched rows are taken from `_rowaddr` (fragment id + row offset).
+        // The updater scans live rows in physical order; `_rowaddr` encodes the slot index used by row-level version metadata.
+        if !read_columns.iter().any(|n| n.as_str() == ROW_ADDR) {
+            read_columns.push(ROW_ADDR.to_string());
+        }
         let mut updater = self
             .updater(
                 Some(&read_columns),
                 Some((write_schema.clone(), self.schema().clone())),
                 None,
             )
             .await?;
-        // Hash join
+        // Hash join: rows matched on the right-hand stream rewrite columns; track physical offsets via `_rowaddr`.
         let joiner = Arc::new(HashJoiner::try_new(right_stream, right_on).await?);
+        let mut matched_offsets = RoaringBitmap::new();
+        let frag_id_u32 = u32::try_from(self.metadata.id).map_err(|_| {
+            Error::invalid_input(format!(
+                "Fragment id {} does not fit RowAddress fragment id",
+                self.metadata.id
+            ))
+        })?;
         while let Some(batch) = updater.next().await? {
+            let index_column = batch[left_on].clone();
+            let matched = joiner.matched_join_rows(index_column.clone())?;
+            if let Some(addr_col) = batch.column_by_name(ROW_ADDR) {
+                let addrs = as_primitive_array::<UInt64Type>(addr_col.as_ref());
+                for (row_idx, &is_matched) in matched.iter().enumerate().take(batch.num_rows()) {
+                    if !is_matched || addrs.is_null(row_idx) {
+                        continue;
+                    }
+                    let addr = RowAddress::from(addrs.value(row_idx));
+                    if addr.fragment_id() == frag_id_u32 {
+                        matched_offsets.insert(addr.row_offset());
+                    }
+                }
+            }
             let updated_batch = joiner
-                .collect_with_fallback(batch, batch[left_on].clone(), self.dataset())
+                .collect_with_fallback(batch, index_column, self.dataset())
                 .await?;
             updater.update(updated_batch).await?;
         }
@@ -1689,8 +1743,11 @@ impl FileFragment {
             .iter()
             .filter_map(|&i| u32::try_from(i).ok())
             .collect();
-        // Note: updated field should be returned when committing, waiting to be done
-        Ok((updated_fragment, updated_fields))
+        Ok(FragmentUpdateColumnsResult {
+            fragment: updated_fragment,
+            fields_modified: updated_fields,
+            matched_offsets,
+        })
     }
 
     /// Append new columns to the fragment
@@ -2647,12 +2704,13 @@ mod tests {
     use lance_io::{assert_io_eq, assert_io_lt, object_store::ObjectStore};
     use pretty_assertions::assert_eq;
     use rstest::rstest;
+    use std::collections::HashMap;
 
     use super::*;
     use crate::{
         dataset::{
             InsertBuilder,
-            transaction::{Operation, UpdateMode},
+            transaction::{Operation, UpdateMode, UpdatedFragmentOffsets},
         },
         session::Session,
         utils::test::TestDatasetGenerator,
@@ -2851,19 +2909,29 @@ mod tests {
             vec![Ok(update_batch1)].into_iter(),
             schema1,
         ));
-        let (updated_fragment1, fields_modified1) = fragment1
-            .update_columns(right_stream1, ROW_ID, ROW_ID)
+        let u1 = fragment1
+            .update_columns_with_offsets(right_stream1, ROW_ID, ROW_ID)
             .await
             .unwrap();
+        assert_eq!(u1.matched_offsets.iter().count(), 38);
+        assert!(!u1.matched_offsets.contains(0));
+        assert!(!u1.matched_offsets.contains(3));
+        assert!(u1.matched_offsets.contains(1));
+        assert!(u1.matched_offsets.contains(39));
+        let frag_id_1 = u1.fragment.id;
+        let matched_1 = u1.matched_offsets;
         let op1 = Operation::Update {
             removed_fragment_ids: vec![],
-            updated_fragments: vec![updated_fragment1],
+            updated_fragments: vec![u1.fragment],
             new_fragments: vec![],
-            fields_modified: fields_modified1,
+            fields_modified: u1.fields_modified,
             merged_generations: Vec::new(),
             fields_for_preserving_frag_bitmap: vec![],
             update_mode: Some(UpdateMode::RewriteColumns),
             inserted_rows_filter: None,
+            updated_fragment_offsets: Some(UpdatedFragmentOffsets(HashMap::from([(
+                frag_id_1, matched_1,
+            )]))),
         };
         let mut dataset1 = Dataset::commit(
             test_uri,
@@ -2924,19 +2992,27 @@ mod tests {
             vec![Ok(update_batch2)].into_iter(),
             schema2,
         ));
-        let (updated_fragment2, fields_modified2) = fragment2
-            .update_columns(right_stream2, "i", "i1")
+        let u2 = fragment2
+            .update_columns_with_offsets(right_stream2, "i", "i1")
             .await
             .unwrap();
+        assert_eq!(u2.matched_offsets.iter().count(), 38);
+        assert!(!u2.matched_offsets.contains(0));
+        assert!(!u2.matched_offsets.contains(3));
+        let frag_id_2 = u2.fragment.id;
+        let matched_2 = u2.matched_offsets;
         let op = Operation::Update {
             removed_fragment_ids: vec![],
-            updated_fragments: vec![updated_fragment2],
+            updated_fragments: vec![u2.fragment],
             new_fragments: vec![],
-            fields_modified: fields_modified2,
+            fields_modified: u2.fields_modified,
             merged_generations: Vec::new(),
             fields_for_preserving_frag_bitmap: vec![],
             update_mode: Some(UpdateMode::RewriteColumns),
             inserted_rows_filter: None,
+            updated_fragment_offsets: Some(UpdatedFragmentOffsets(HashMap::from([(
+                frag_id_2, matched_2,
+            )]))),
         };
         let dataset2 = Dataset::commit(
             test_uri,