fix: Evaluate UPDATE expressions only on matching rows

- Use evaluate_selection() instead of evaluate() for UPDATE assignments
  to avoid evaluating expressions on rows excluded by WHERE clause.
  This prevents errors like divide-by-zero on non-updated rows.

- Extract evaluate_filters_to_mask() helper shared by delete_from/update

- Hoist physical expression creation outside batch loop for efficiency

- Simplify strip_column_qualifiers() signature

- Add regression test for divide-by-zero edge case

Author: ethan-tyler

datafusion/catalog/src/memory/table.rs

@@ -329,50 +329,35 @@ impl TableProvider for MemTable {
                     continue;
                 }
 
-                let filter_mask = if filters.is_empty() {
-                    BooleanArray::from(vec![true; batch.num_rows()])
-                } else {
-                    let mut combined_mask: Option<BooleanArray> = None;
-
-                    for filter_expr in &filters {
-                        let physical_expr = create_physical_expr(
-                            filter_expr,
-                            &df_schema,
-                            state.execution_props(),
-                        )?;
-
-                        let result = physical_expr.evaluate(batch)?;
-                        let array = result.into_array(batch.num_rows())?;
-                        let bool_array = array
-                            .as_any()
-                            .downcast_ref::<BooleanArray>()
-                            .ok_or_else(|| {
-                                datafusion_common::DataFusionError::Internal(
-                                    "Filter did not evaluate to boolean".to_string(),
-                                )
-                            })?
-                            .clone();
-
-                        combined_mask = Some(match combined_mask {
-                            Some(existing) => and(&existing, &bool_array)?,
-                            None => bool_array,
-                        });
+                // Evaluate filters - None means "match all rows"
+                let filter_mask = evaluate_filters_to_mask(
+                    &filters,
+                    batch,
+                    &df_schema,
+                    state.execution_props(),
+                )?;
+
+                let (delete_count, keep_mask) = match filter_mask {
+                    Some(mask) => {
+                        // Count rows where mask is true (will be deleted)
+                        let count = mask.iter().filter(|v| v == &Some(true)).count();
+                        // Keep rows where predicate is false or NULL (SQL three-valued logic)
+                        let keep: BooleanArray =
+                            mask.iter().map(|v| Some(v != Some(true))).collect();
+                        (count, keep)
+                    }
+                    None => {
+                        // No filters = delete all rows
+                        (
+                            batch.num_rows(),
+                            BooleanArray::from(vec![false; batch.num_rows()]),
+                        )
                     }
-
-                    combined_mask.unwrap_or_else(|| {
-                        BooleanArray::from(vec![true; batch.num_rows()])
-                    })
                 };
 
-                let delete_count =
-                    filter_mask.iter().filter(|v| v == &Some(true)).count();
                 total_deleted += delete_count as u64;
 
-                // Keep rows where predicate is false or NULL (SQL three-valued logic)
-                let keep_mask: BooleanArray =
-                    filter_mask.iter().map(|v| Some(v != Some(true))).collect();
                 let filtered_batch = filter_record_batch(batch, &keep_mask)?;
-
                 if filtered_batch.num_rows() > 0 {
                     new_batches.push(filtered_batch);
                 }
@@ -412,15 +397,24 @@ impl TableProvider for MemTable {
             }
         }
 
-        let assignment_map: HashMap<&str, &Expr> = assignments
+        let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?;
+
+        // Create physical expressions for assignments upfront (outside batch loop)
+        let physical_assignments: HashMap<
+            String,
+            Arc<dyn datafusion_physical_plan::PhysicalExpr>,
+        > = assignments
             .iter()
-            .map(|(name, expr)| (name.as_str(), expr))
-            .collect();
+            .map(|(name, expr)| {
+                let physical_expr =
+                    create_physical_expr(expr, &df_schema, state.execution_props())?;
+                Ok((name.clone(), physical_expr))
+            })
+            .collect::<Result<_>>()?;
 
         *self.sort_order.lock() = vec![];
 
         let mut total_updated: u64 = 0;
-        let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?;
 
         for partition_data in &self.batches {
             let mut partition = partition_data.write().await;
@@ -431,54 +425,39 @@ impl TableProvider for MemTable {
                     continue;
                 }
 
-                let filter_mask = if filters.is_empty() {
-                    BooleanArray::from(vec![true; batch.num_rows()])
-                } else {
-                    let mut combined_mask: Option<BooleanArray> = None;
-
-                    for filter_expr in &filters {
-                        let physical_expr = create_physical_expr(
-                            filter_expr,
-                            &df_schema,
-                            state.execution_props(),
-                        )?;
-
-                        let result = physical_expr.evaluate(batch)?;
-                        let array = result.into_array(batch.num_rows())?;
-                        let bool_array = array
-                            .as_any()
-                            .downcast_ref::<BooleanArray>()
-                            .ok_or_else(|| {
-                                datafusion_common::DataFusionError::Internal(
-                                    "Filter did not evaluate to boolean".to_string(),
-                                )
-                            })?
-                            .clone();
-
-                        combined_mask = Some(match combined_mask {
-                            Some(existing) => and(&existing, &bool_array)?,
-                            None => bool_array,
-                        });
+                // Evaluate filters - None means "match all rows"
+                let filter_mask = evaluate_filters_to_mask(
+                    &filters,
+                    batch,
+                    &df_schema,
+                    state.execution_props(),
+                )?;
+
+                let (update_count, update_mask) = match filter_mask {
+                    Some(mask) => {
+                        // Count rows where mask is true (will be updated)
+                        let count = mask.iter().filter(|v| v == &Some(true)).count();
+                        // Normalize mask: only true (not NULL) triggers update
+                        let normalized: BooleanArray =
+                            mask.iter().map(|v| Some(v == Some(true))).collect();
+                        (count, normalized)
+                    }
+                    None => {
+                        // No filters = update all rows
+                        (
+                            batch.num_rows(),
+                            BooleanArray::from(vec![true; batch.num_rows()]),
+                        )
                     }
-
-                    combined_mask.unwrap_or_else(|| {
-                        BooleanArray::from(vec![true; batch.num_rows()])
-                    })
                 };
 
-                let update_count =
-                    filter_mask.iter().filter(|v| v == &Some(true)).count();
                 total_updated += update_count as u64;
 
                 if update_count == 0 {
                     new_batches.push(batch.clone());
                     continue;
                 }
 
-                // Normalize mask: only true (not NULL) triggers update
-                let update_mask: BooleanArray =
-                    filter_mask.iter().map(|v| Some(v == Some(true))).collect();
-
                 let mut new_columns: Vec<ArrayRef> =
                     Vec::with_capacity(batch.num_columns());
 
@@ -491,16 +470,15 @@ impl TableProvider for MemTable {
                             ))
                         })?;
 
-                    let new_column = if let Some(value_expr) =
-                        assignment_map.get(column_name.as_str())
+                    let new_column = if let Some(physical_expr) =
+                        physical_assignments.get(column_name.as_str())
                     {
-                        let physical_expr = create_physical_expr(
-                            value_expr,
-                            &df_schema,
-                            state.execution_props(),
-                        )?;
-
-                        let new_values = physical_expr.evaluate(batch)?;
+                        // Use evaluate_selection to only evaluate on matching rows.
+                        // This avoids errors (e.g., divide-by-zero) on rows that won't
+                        // be updated. The result is scattered back with nulls for
+                        // non-matching rows, which zip() will replace with originals.
+                        let new_values =
+                            physical_expr.evaluate_selection(batch, &update_mask)?;
                         let new_array = new_values.into_array(batch.num_rows())?;
 
                         // Convert to &dyn Array which implements Datum
@@ -526,6 +504,46 @@ impl TableProvider for MemTable {
     }
 }
 
+/// Evaluate filter expressions against a batch and return a combined boolean mask.
+/// Returns None if filters is empty (meaning "match all rows").
+/// The returned mask has true for rows that match the filter predicates.
+fn evaluate_filters_to_mask(
+    filters: &[Expr],
+    batch: &RecordBatch,
+    df_schema: &DFSchema,
+    execution_props: &datafusion_expr::execution_props::ExecutionProps,
+) -> Result<Option<BooleanArray>> {
+    if filters.is_empty() {
+        return Ok(None);
+    }
+
+    let mut combined_mask: Option<BooleanArray> = None;
+
+    for filter_expr in filters {
+        let physical_expr =
+            create_physical_expr(filter_expr, df_schema, execution_props)?;
+
+        let result = physical_expr.evaluate(batch)?;
+        let array = result.into_array(batch.num_rows())?;
+        let bool_array = array
+            .as_any()
+            .downcast_ref::<BooleanArray>()
+            .ok_or_else(|| {
+                datafusion_common::DataFusionError::Internal(
+                    "Filter did not evaluate to boolean".to_string(),
+                )
+            })?
+            .clone();
+
+        combined_mask = Some(match combined_mask {
+            Some(existing) => and(&existing, &bool_array)?,
+            None => bool_array,
+        });
+    }
+
+    Ok(combined_mask)
+}
+
 /// Returns a single row with the count of affected rows.
 #[derive(Debug)]
 struct DmlResultExec {

datafusion/core/src/physical_planner.rs

@@ -1906,6 +1906,7 @@ fn get_physical_expr_pair(
 /// splitting AND conjunctions into individual expressions.
 /// Column qualifiers are stripped so expressions can be evaluated against
 /// the TableProvider's schema.
+///
 fn extract_dml_filters(input: &Arc<LogicalPlan>) -> Result<Vec<Expr>> {
     let mut filters = Vec::new();
 
@@ -1929,7 +1930,7 @@ fn strip_column_qualifiers(expr: Expr) -> Result<Expr> {
         if let Expr::Column(col) = &e
             && col.relation.is_some()
         {
-            // Create unqualified column
+            // Strip the qualifier
             return Ok(Transformed::yes(Expr::Column(Column::new_unqualified(
                 col.name.clone(),
             ))));
@@ -1943,6 +1944,7 @@ fn strip_column_qualifiers(expr: Expr) -> Result<Expr> {
 /// For UPDATE statements, the SQL planner encodes assignments as a projection
 /// over the source table. This function extracts column name and expression pairs
 /// from the projection. Column qualifiers are stripped from the expressions.
+///
 fn extract_update_assignments(input: &Arc<LogicalPlan>) -> Result<Vec<(String, Expr)>> {
     // The UPDATE input plan structure is:
     // Projection(updated columns as expressions with aliases)

datafusion/sqllogictest/test_files/dml_update.slt

@@ -257,3 +257,30 @@ UPDATE test_update_error SET nonexistent = 'value';
 
 statement ok
 DROP TABLE test_update_error;
+
+# Test UPDATE with expression that would error on non-matching rows
+# Regression test: expressions should only be evaluated on rows that match
+# the WHERE clause, not all rows. This prevents divide-by-zero errors
+# on rows that won't be updated.
+statement ok
+CREATE TABLE test_update_div(id INT, divisor INT, result INT);
+
+statement ok
+INSERT INTO test_update_div VALUES (1, 0, 0), (2, 2, 0), (3, 5, 0);
+
+# This should succeed: 1/divisor is only evaluated where divisor != 0
+# Row 1 (divisor=0) is excluded by WHERE clause and expression is NOT evaluated
+query I
+UPDATE test_update_div SET result = 100 / divisor WHERE divisor != 0;
+----
+2
+
+query III rowsort
+SELECT * FROM test_update_div;
+----
+1 0 0
+2 2 50
+3 5 20
+
+statement ok
+DROP TABLE test_update_div;