feat: invalidate generated columns on native update

rust/lancedb/src/function/plan_generated_column_invalidation.rs

@@ -5,8 +5,8 @@
 //!
 //! Plans column-wide dependency-epoch advances from a binding snapshot and a
 //! mutation impact. This module does not mutate tables, write metadata, or
-//! execute append/update/delete/merge paths. Native append consumes the plan
-//! through the B4b runtime wiring.
+//! execute append/update/delete/merge paths. Native append and update consume
+//! the plan through the B4b / B4c runtime wiring.
 
 use std::collections::BTreeSet;
 
@@ -20,9 +20,8 @@ pub enum GeneratedColumnMutationImpact {
     RowSetChanged,
     /// Update of the listed stable field IDs (direct and transitive dependents).
     ///
-    /// Reserved for future update/delete/merge invalidation consumers; Native
-    /// append (B4b) only constructs [`Self::RowSetChanged`].
-    #[cfg_attr(not(test), allow(dead_code))]
+    /// Native update (B4c) constructs this impact. Native append (B4b) only
+    /// constructs [`Self::RowSetChanged`].
     UpdatedFields(BTreeSet<i32>),
 }
 

rust/lancedb/src/table.rs

@@ -70,13 +70,13 @@ use self::merge::MergeInsertBuilder;
 
 pub mod add_columns;
 mod add_data;
-mod append_generated_column_invalidation;
 pub mod branch_merge;
 pub mod checkpoint;
 mod create_index;
 pub mod datafusion;
 pub(crate) mod dataset;
 pub mod delete;
+mod generated_column_invalidation;
 pub mod lsm_stats;
 pub mod merge;
 pub mod optimize;
@@ -90,6 +90,8 @@ pub mod write_progress;
 mod append_generated_column_invalidation_contract;
 #[cfg(test)]
 mod schema_metadata_updates_dependency_contract;
+#[cfg(test)]
+mod update_generated_column_invalidation_contract;
 
 use crate::index::waiter::wait_for_index;
 pub use add_columns::AddColumnsBuilder;
@@ -3270,7 +3272,7 @@ impl BaseTable for NativeTable {
         // Plan after in-memory preprocessing, before any InsertExec file write.
         // Canonical overwrite is PreprocessingOutput.overwrite, not final lance_params.mode.
         let schema_metadata_updates =
-            append_generated_column_invalidation::plan_native_append_generated_column_invalidation(
+            generated_column_invalidation::plan_native_append_generated_column_invalidation(
                 ds.as_ref(),
                 output.overwrite,
             )?;
@@ -3791,7 +3793,7 @@ impl BaseTable for NativeTable {
         let is_overwrite = matches!(write_params.mode, WriteMode::Overwrite);
         // Reject generated-table overwrite before returning an execution plan.
         let schema_metadata_updates =
-            append_generated_column_invalidation::plan_native_append_generated_column_invalidation(
+            generated_column_invalidation::plan_native_append_generated_column_invalidation(
                 dataset.as_ref(),
                 is_overwrite,
             )?;

rust/lancedb/src/table/generated_column_invalidation.rs

@@ -1,13 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
-//! Crate-private Native append wiring for generated-column invalidation (B4b).
+//! Crate-private Native wiring for generated-column invalidation (B4b / B4c).
 //!
 //! Converts the B4a pure planner into one Lance field-metadata patch for Native
-//! append commits. Planning is strict-decode/validate; overwrite of a table with
-//! any generated-column definition fails closed as [`Error::NotSupported`].
+//! append and update commits. Planning is strict-decode/validate; overwrite of a
+//! table with any generated-column definition, and direct writes of generated
+//! outputs via Update, fail closed as [`Error::NotSupported`].
 
-use std::collections::HashMap;
+use std::collections::{BTreeSet, HashMap};
 
 use lance::Dataset;
 use lance::dataset::transaction::{SchemaMetadataUpdates, UpdateMap, UpdateMapEntry};
@@ -49,6 +50,51 @@ pub(super) fn plan_native_append_generated_column_invalidation(
     Ok(Some(planned_invalidation_to_schema_metadata_updates(plan)))
 }
 
+/// Plan Native update invalidation against one exact dataset snapshot.
+///
+/// Strict-decodes and validates every present generated-column definition before
+/// impact calculation, even when `updated_field_ids` does not affect any
+/// generated output. After the global planner succeeds, a target whose snapshot
+/// entry contains generated metadata is rejected as a direct generated-output
+/// write ([`Error::NotSupported`]) before any Update file write. Returns
+/// `Some(patch)` when the impact closure is non-empty, otherwise `None`.
+pub(super) fn plan_native_update_generated_column_invalidation(
+    dataset: &Dataset,
+    updated_field_ids: BTreeSet<i32>,
+) -> Result<Option<SchemaMetadataUpdates>> {
+    let snapshot = generated_column_binding_snapshot_from_dataset(dataset)?;
+    let plan = plan_generated_column_invalidation(
+        &snapshot,
+        &GeneratedColumnMutationImpact::UpdatedFields(updated_field_ids.clone()),
+    )?;
+
+    for field_id in &updated_field_ids {
+        let Some(entry) = snapshot
+            .entries()
+            .iter()
+            .find(|entry| entry.field_id() == *field_id)
+        else {
+            return Err(Error::InvalidInput {
+                message: format!("updated field id {field_id} was not found in the table schema"),
+            });
+        };
+        if entry
+            .field()
+            .metadata()
+            .contains_key(GENERATED_COLUMN_METADATA_KEY)
+        {
+            return Err(Error::NotSupported {
+                message: "Updating generated columns is not supported".to_string(),
+            });
+        }
+    }
+
+    if plan.is_empty() {
+        return Ok(None);
+    }
+    Ok(Some(planned_invalidation_to_schema_metadata_updates(plan)))
+}
+
 /// Convert planner replacements into one non-empty Lance field-metadata patch.
 ///
 /// Each entry is keyed by stable output field ID, uses `replace: false`, and

rust/lancedb/src/table/update.rs

@@ -1,8 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
+use std::collections::BTreeSet;
 use std::sync::Arc;
 
+use lance::Dataset;
 use lance::dataset::UpdateBuilder as LanceUpdateBuilder;
 use serde::{Deserialize, Serialize};
 
@@ -80,27 +82,37 @@ pub(crate) async fn execute_update(
 ) -> Result<UpdateResult> {
     table.dataset.ensure_mutable()?;
 
-    // 1. Snapshot the current dataset
+    // One exact dataset supplies SET/filter planning, stable target field IDs,
+    // generated-column invalidation planning, the Lance UpdateBuilder, and its
+    // transaction basis. Do not call table schema()/version() or another get().
     let dataset = table.dataset.get().await?;
 
-    // 2. Initialize the Lance Core builder
-    let mut builder = LanceUpdateBuilder::new(dataset);
+    let mut builder = LanceUpdateBuilder::new(dataset.clone());
 
-    // 3. Apply the filter (WHERE clause)
     if let Some(predicate) = update.filter {
         builder = builder.update_where(&predicate)?;
     }
 
-    // 4. Apply the columns (SET clause)
-    for (column, value) in update.columns {
-        builder = builder.set(column, &value)?;
+    let columns = update.columns;
+    for (column, value) in &columns {
+        builder = builder.set(column, value)?;
+    }
+
+    // After Lance SET validation, resolve stable field IDs from the same
+    // snapshot and plan invalidation before UpdateJob writes files.
+    let updated_field_ids = updated_stable_field_ids(dataset.as_ref(), &columns)?;
+    if let Some(schema_metadata_updates) =
+        super::generated_column_invalidation::plan_native_update_generated_column_invalidation(
+            dataset.as_ref(),
+            updated_field_ids,
+        )?
+    {
+        builder = builder.with_schema_metadata_updates(schema_metadata_updates)?;
     }
 
-    // 5. Execute the operation (Write new files)
     let operation = builder.build()?;
     let res = operation.execute().await?;
 
-    // 6. Update the table's view of the latest version
     table.dataset.update(res.new_dataset.as_ref().clone());
 
     Ok(UpdateResult {
@@ -109,6 +121,24 @@ pub(crate) async fn execute_update(
     })
 }
 
+/// Resolve exact top-level SET targets to a deterministic set of stable field IDs.
+fn updated_stable_field_ids(
+    dataset: &Dataset,
+    columns: &[(String, String)],
+) -> Result<BTreeSet<i32>> {
+    let mut ids = BTreeSet::new();
+    for (column, _) in columns {
+        let field = dataset
+            .schema()
+            .field(column)
+            .ok_or_else(|| Error::InvalidInput {
+                message: format!("Column '{column}' does not exist in dataset schema"),
+            })?;
+        ids.insert(field.id);
+    }
+    Ok(ids)
+}
+
 #[cfg(test)]
 mod tests {
     use crate::connect;