From 8e695f6f3fdf276af5ade845c344e707678a2e82 Mon Sep 17 00:00:00 2001 From: mck Date: Tue, 25 Aug 2026 10:29:01 +0200 Subject: [PATCH] fix(catalog): apply a MemTable DELETE or UPDATE when the plan runs `EXPLAIN DELETE` and `EXPLAIN UPDATE` changed the rows of a `MemTable`. `handle_explain()` builds the physical plan in order to print it, the physical planner calls the provider hook while it builds the plan, and `MemTable` did the whole row change inside the hook. The returned `DmlResultExec` was a constant node that only reported the count the hook had computed, so the plan text also carried the count. Replace `DmlResultExec` with `MemDmlExec`. The hook now compiles the `WHERE` clause and the assignments, then returns a plan that holds the partitions and the declared sort order of the table. `execute()` applies the operation, clears the sort order, and emits the count. This is the pattern that the provider guide already recommends, and `MemTable` is the reference implementation. Every check of the statement stays in the hook, so an `EXPLAIN` still reports an invalid statement. A plan that runs twice applies the operation twice, as `DataSinkExec` does for an INSERT. The `DmlResultExec: rows_affected=0` lines of `delete.slt` and `update.slt` become `MemDmlExec: op=Delete` and `MemDmlExec: op=Update`. The count is unknown while the plan is built, so it no longer appears in the plan text. Co-Authored-By: Claude Opus 5 --- datafusion/catalog/src/memory/table.rs | 427 +++++++++++------- datafusion/sqllogictest/test_files/delete.slt | 12 +- .../sqllogictest/test_files/dml_delete.slt | 50 ++ .../sqllogictest/test_files/dml_update.slt | 50 ++ datafusion/sqllogictest/test_files/update.slt | 4 +- 5 files changed, 371 insertions(+), 172 deletions(-) diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index d817aa8b7788a..836bf1ad54872 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -33,7 +33,9 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::error::Result; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{Constraints, DFSchema, SchemaExt, not_impl_err, plan_err}; +use datafusion_common::{ + Constraints, DFSchema, SchemaExt, internal_err, not_impl_err, plan_err, +}; use datafusion_datasource::memory::{MemSink, MemorySourceConfig}; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::source::DataSourceExec; @@ -361,71 +363,24 @@ impl MemTable { state: &'a dyn Session, filters: Vec, ) -> BoxFuture<'a, Result>> { - Box::pin(self.delete_from_inner(state, filters)) + Box::pin(ready(self.plan_delete(state, filters))) } - async fn delete_from_inner( + /// Build the plan of a DELETE. The rows change when the plan runs, not here. + fn plan_delete( &self, state: &dyn Session, filters: Vec, ) -> Result> { // Early exit if table has no partitions if self.batches.is_empty() { - return Ok(Arc::new(DmlResultExec::new(0))); + return Ok(self.dml_exec(vec![], vec![], MemDmlOp::Delete)); } - *self.sort_order.lock() = vec![]; - - let mut total_deleted: u64 = 0; let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?; + let filters = compile_filters(filters, &df_schema, state.execution_props())?; - for partition_data in &self.batches { - let mut partition = partition_data.write().await; - let mut new_batches = Vec::with_capacity(partition.len()); - - for batch in partition.iter() { - if batch.num_rows() == 0 { - continue; - } - - // 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()]), - ) - } - }; - - total_deleted += delete_count as u64; - - let filtered_batch = filter_record_batch(batch, &keep_mask)?; - if filtered_batch.num_rows() > 0 { - new_batches.push(filtered_batch); - } - } - - *partition = new_batches; - } - - Ok(Arc::new(DmlResultExec::new(total_deleted))) + Ok(self.dml_exec(self.batches.clone(), filters, MemDmlOp::Delete)) } fn update_boxed<'a>( @@ -434,10 +389,13 @@ impl MemTable { assignments: Vec<(String, Expr)>, filters: Vec, ) -> BoxFuture<'a, Result>> { - Box::pin(self.update_inner(state, assignments, filters)) + Box::pin(ready(self.plan_update(state, assignments, filters))) } - async fn update_inner( + /// Build the plan of an UPDATE. The rows change when the plan runs, not + /// here. Every check of the statement stays here, so that an `EXPLAIN` + /// still reports an invalid statement. + fn plan_update( &self, state: &dyn Session, assignments: Vec<(String, Expr)>, @@ -445,7 +403,7 @@ impl MemTable { ) -> Result> { // Early exit if table has no partitions if self.batches.is_empty() { - return Ok(Arc::new(DmlResultExec::new(0))); + return Ok(self.dml_exec(vec![], vec![], MemDmlOp::Update(HashMap::new()))); } // Validate column names upfront with clear error messages @@ -469,118 +427,208 @@ impl MemTable { // Create physical expressions for assignments upfront (outside batch loop) let physical_assignments: HashMap> = assignments - .iter() + .into_iter() .map(|(name, expr)| { let physical_expr = create_physical_expr( - expr, + &expr, &df_schema, state.execution_props(), &PhysicalPlanningContext::default(), )?; - Ok((name.clone(), physical_expr)) + Ok((name, physical_expr)) }) .collect::>()?; - *self.sort_order.lock() = vec![]; + let filters = compile_filters(filters, &df_schema, state.execution_props())?; - let mut total_updated: u64 = 0; + Ok(self.dml_exec( + self.batches.clone(), + filters, + MemDmlOp::Update(physical_assignments), + )) + } - for partition_data in &self.batches { - let mut partition = partition_data.write().await; - let mut new_batches = Vec::with_capacity(partition.len()); + /// Build the plan node that applies `op` to `partitions`. + fn dml_exec( + &self, + partitions: Vec, + filters: Vec>, + op: MemDmlOp, + ) -> Arc { + Arc::new(MemDmlExec::new(MemDmlState { + partitions, + table_schema: Arc::clone(&self.schema), + sort_order: Arc::clone(&self.sort_order), + filters, + op, + })) + } +} - for batch in partition.iter() { - if batch.num_rows() == 0 { - continue; - } +/// Compile the `WHERE` clause of a DELETE or an UPDATE into physical +/// expressions. An empty result means "match all rows". +fn compile_filters( + filters: Vec, + df_schema: &DFSchema, + execution_props: &datafusion_expr::execution_props::ExecutionProps, +) -> Result>> { + filters + .into_iter() + .map(|filter_expr| { + create_physical_expr( + &filter_expr, + df_schema, + execution_props, + &PhysicalPlanningContext::default(), + ) + }) + .collect() +} - // Evaluate filters - None means "match all rows" - let filter_mask = evaluate_filters_to_mask( - &filters, - batch, - &df_schema, - state.execution_props(), - )?; +/// Delete the rows of `state` that its filters match, and return the number of +/// rows deleted. +async fn apply_delete(state: &MemDmlState) -> Result { + let mut total_deleted: u64 = 0; - 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()]), - ) - } - }; + for partition_data in &state.partitions { + let mut partition = partition_data.write().await; + let mut new_batches = Vec::with_capacity(partition.len()); - total_updated += update_count as u64; + for batch in partition.iter() { + if batch.num_rows() == 0 { + continue; + } - if update_count == 0 { - new_batches.push(batch.clone()); - continue; + // Evaluate filters - None means "match all rows" + let filter_mask = evaluate_filters_to_mask(&state.filters, batch)?; + + 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()]), + ) + } + }; + + total_deleted += delete_count as u64; - let mut new_columns: Vec = - Vec::with_capacity(batch.num_columns()); - - for field in self.schema.fields() { - let column_name = field.name(); - let original_column = - batch.column_by_name(column_name).ok_or_else(|| { - datafusion_common::DataFusionError::Internal(format!( - "Column '{column_name}' not found in batch" - )) - })?; - - let new_column = if let Some(physical_expr) = - physical_assignments.get(column_name.as_str()) - { - // 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 - let new_arr: &dyn Array = new_array.as_ref(); - let orig_arr: &dyn Array = original_column.as_ref(); - zip(&update_mask, &new_arr, &orig_arr)? - } else { - Arc::clone(original_column) - }; - - new_columns.push(new_column); + let filtered_batch = filter_record_batch(batch, &keep_mask)?; + if filtered_batch.num_rows() > 0 { + new_batches.push(filtered_batch); + } + } + + *partition = new_batches; + } + + Ok(total_deleted) +} + +/// Assign a new value to each row of `state` that its filters match, and return +/// the number of rows updated. +async fn apply_update( + state: &MemDmlState, + physical_assignments: &HashMap>, +) -> Result { + let mut total_updated: u64 = 0; + + for partition_data in &state.partitions { + let mut partition = partition_data.write().await; + let mut new_batches = Vec::with_capacity(partition.len()); + + for batch in partition.iter() { + if batch.num_rows() == 0 { + continue; + } + + // Evaluate filters - None means "match all rows" + let filter_mask = evaluate_filters_to_mask(&state.filters, batch)?; + + 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()]), + ) } + }; - let updated_batch = - ArrowRecordBatch::try_new(Arc::clone(&self.schema), new_columns)?; - new_batches.push(updated_batch); + total_updated += update_count as u64; + + if update_count == 0 { + new_batches.push(batch.clone()); + continue; } - *partition = new_batches; + let mut new_columns: Vec = Vec::with_capacity(batch.num_columns()); + + for field in state.table_schema.fields() { + let column_name = field.name(); + let original_column = + batch.column_by_name(column_name).ok_or_else(|| { + datafusion_common::DataFusionError::Internal(format!( + "Column '{column_name}' not found in batch" + )) + })?; + + let new_column = if let Some(physical_expr) = + physical_assignments.get(column_name.as_str()) + { + // 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 + let new_arr: &dyn Array = new_array.as_ref(); + let orig_arr: &dyn Array = original_column.as_ref(); + zip(&update_mask, &new_arr, &orig_arr)? + } else { + Arc::clone(original_column) + }; + + new_columns.push(new_column); + } + + let updated_batch = + ArrowRecordBatch::try_new(Arc::clone(&state.table_schema), new_columns)?; + new_batches.push(updated_batch); } - Ok(Arc::new(DmlResultExec::new(total_updated))) + *partition = new_batches; } + + Ok(total_updated) } /// 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], + filters: &[Arc], batch: &RecordBatch, - df_schema: &DFSchema, - execution_props: &datafusion_expr::execution_props::ExecutionProps, ) -> Result> { if filters.is_empty() { return Ok(None); @@ -588,14 +636,7 @@ fn evaluate_filters_to_mask( let mut combined_mask: Option = None; - for filter_expr in filters { - let physical_expr = create_physical_expr( - filter_expr, - df_schema, - execution_props, - &PhysicalPlanningContext::default(), - )?; - + for physical_expr in filters { let result = physical_expr.evaluate(batch)?; let array = result.into_array(batch.num_rows())?; let bool_array = array @@ -617,16 +658,54 @@ fn evaluate_filters_to_mask( Ok(combined_mask) } -/// Returns a single row with the count of affected rows. +/// The operation that a [`MemDmlExec`] applies to the rows of a [`MemTable`]. #[derive(Debug)] -struct DmlResultExec { - rows_affected: u64, +enum MemDmlOp { + /// Delete each row that the filters match. + Delete, + /// Assign a new value to each row that the filters match. The map holds one + /// expression per assigned column, keyed by column name. + Update(HashMap>), +} + +impl MemDmlOp { + fn as_str(&self) -> &'static str { + match self { + MemDmlOp::Delete => "Delete", + MemDmlOp::Update(_) => "Update", + } + } +} + +/// Everything that a [`MemDmlExec`] needs in order to apply its operation. +/// Each field is a clone of a field of the [`MemTable`], so the plan changes the +/// rows of the table itself. +#[derive(Debug)] +struct MemDmlState { + partitions: Vec, + table_schema: SchemaRef, + sort_order: Arc>>>, + /// The `WHERE` clause of the statement, compiled while the plan was built. + /// An empty list means "match all rows". + filters: Vec>, + op: MemDmlOp, +} + +/// Applies a DELETE or an UPDATE to a [`MemTable`], and returns a single row +/// with the count of affected rows. +/// +/// The rows change in [`ExecutionPlan::execute`], not while the plan is built, +/// so an `EXPLAIN` of the statement leaves the table alone. Each run of the plan +/// applies the operation once more, as [`DataSinkExec`] does for an INSERT. +#[derive(Debug)] +struct MemDmlExec { + state: Arc, schema: SchemaRef, properties: Arc, } -impl DmlResultExec { - fn new(rows_affected: u64) -> Self { +impl MemDmlExec { + fn new(state: MemDmlState) -> Self { let schema = Arc::new(Schema::new(vec![Field::new( "count", DataType::UInt64, @@ -641,14 +720,14 @@ impl DmlResultExec { ); Self { - rows_affected, + state: Arc::new(state), schema, properties: Arc::new(properties), } } } -impl DisplayAs for DmlResultExec { +impl DisplayAs for MemDmlExec { fn fmt_as( &self, t: DisplayFormatType, @@ -658,15 +737,15 @@ impl DisplayAs for DmlResultExec { DisplayFormatType::Default | DisplayFormatType::Verbose | DisplayFormatType::TreeRender => { - write!(f, "DmlResultExec: rows_affected={}", self.rows_affected) + write!(f, "MemDmlExec: op={}", self.state.op.as_str()) } } } } -impl ExecutionPlan for DmlResultExec { +impl ExecutionPlan for MemDmlExec { fn name(&self) -> &str { - "DmlResultExec" + "MemDmlExec" } fn schema(&self) -> SchemaRef { @@ -701,18 +780,38 @@ impl ExecutionPlan for DmlResultExec { fn execute( &self, - _partition: usize, + partition: usize, _context: Arc, ) -> Result { - // Create a single batch with the count - let count_array = UInt64Array::from(vec![self.rows_affected]); - let batch = ArrowRecordBatch::try_new( - Arc::clone(&self.schema), - vec![Arc::new(count_array) as ArrayRef], - )?; + if partition != 0 { + return internal_err!( + "MemDmlExec has one partition, but partition {partition} was requested" + ); + } + + let state = Arc::clone(&self.state); + let schema = Arc::clone(&self.schema); + + // Apply the operation, then emit the count as the single output row. + let stream = futures::stream::once(async move { + // The rows change, so any declared sort order no longer holds. The + // guard drops at the end of this statement, before the first await. + *state.sort_order.lock() = vec![]; + + let rows_affected = match &state.op { + MemDmlOp::Delete => apply_delete(&state).await?, + MemDmlOp::Update(assignments) => { + apply_update(&state, assignments).await? + } + }; + + let count_array = UInt64Array::from(vec![rows_affected]); + Ok(ArrowRecordBatch::try_new( + schema, + vec![Arc::new(count_array) as ArrayRef], + )?) + }); - // Create a stream that yields just this one batch - let stream = futures::stream::iter(vec![Ok(batch)]); Ok(Box::pin(RecordBatchStreamAdapter::new( Arc::clone(&self.schema), stream, diff --git a/datafusion/sqllogictest/test_files/delete.slt b/datafusion/sqllogictest/test_files/delete.slt index 1f33360824393..ca182f07b442f 100644 --- a/datafusion/sqllogictest/test_files/delete.slt +++ b/datafusion/sqllogictest/test_files/delete.slt @@ -36,7 +36,7 @@ logical_plan 02)--TableScan: t1 physical_plan 01)CooperativeExec -02)--DmlResultExec: rows_affected=0 +02)--MemDmlExec: op=Delete # Filtered by existing columns @@ -49,7 +49,7 @@ logical_plan 03)----TableScan: t1 physical_plan 01)CooperativeExec -02)--DmlResultExec: rows_affected=0 +02)--MemDmlExec: op=Delete # Filtered by existing columns, using qualified and unqualified names @@ -62,7 +62,7 @@ logical_plan 03)----TableScan: t1 physical_plan 01)CooperativeExec -02)--DmlResultExec: rows_affected=0 +02)--MemDmlExec: op=Delete # Filtered by a mix of columns and literal predicates @@ -75,7 +75,7 @@ logical_plan 03)----TableScan: t1 physical_plan 01)CooperativeExec -02)--DmlResultExec: rows_affected=0 +02)--MemDmlExec: op=Delete # Deleting by columns that do not exist returns an error @@ -126,7 +126,7 @@ logical_plan 03)----TableScan: t1 physical_plan 01)CooperativeExec -02)--DmlResultExec: rows_affected=0 +02)--MemDmlExec: op=Delete query TT @@ -139,7 +139,7 @@ logical_plan 04)------TableScan: t1 physical_plan 01)CooperativeExec -02)--DmlResultExec: rows_affected=0 +02)--MemDmlExec: op=Delete # Config reset statement ok diff --git a/datafusion/sqllogictest/test_files/dml_delete.slt b/datafusion/sqllogictest/test_files/dml_delete.slt index 3dae431ada377..294f75a19cbb0 100644 --- a/datafusion/sqllogictest/test_files/dml_delete.slt +++ b/datafusion/sqllogictest/test_files/dml_delete.slt @@ -200,3 +200,53 @@ SELECT * FROM test_delete_in; statement ok DROP TABLE test_delete_in; + +# Test that EXPLAIN DELETE does not change the rows +# The rows change when the plan runs, and EXPLAIN prints the plan without +# running it. +statement ok +CREATE TABLE test_explain_delete AS VALUES (1, 'a'), (2, 'b'), (3, 'c'); + +query TT +EXPLAIN DELETE FROM test_explain_delete WHERE column1 > 1; +---- +logical_plan +01)Dml: op=[Delete] table=[test_explain_delete] +02)--Filter: test_explain_delete.column1 > Int64(1) +03)----TableScan: test_explain_delete projection=[column1, column2] +physical_plan +01)CooperativeExec +02)--MemDmlExec: op=Delete + +query TT +EXPLAIN DELETE FROM test_explain_delete; +---- +logical_plan +01)Dml: op=[Delete] table=[test_explain_delete] +02)--TableScan: test_explain_delete projection=[column1, column2] +physical_plan +01)CooperativeExec +02)--MemDmlExec: op=Delete + +query IT rowsort +SELECT * FROM test_explain_delete; +---- +1 a +2 b +3 c + +# EXPLAIN ANALYZE runs the plan, so the rows do go +query TT +EXPLAIN ANALYZE DELETE FROM test_explain_delete WHERE column1 > 1; +---- +Plan with Metrics +01)CooperativeExec, metrics=[] +02)--MemDmlExec: op=Delete, metrics=[] + +query IT rowsort +SELECT * FROM test_explain_delete; +---- +1 a + +statement ok +DROP TABLE test_explain_delete; diff --git a/datafusion/sqllogictest/test_files/dml_update.slt b/datafusion/sqllogictest/test_files/dml_update.slt index 10f74ae3970da..bf18f5e1fc1db 100644 --- a/datafusion/sqllogictest/test_files/dml_update.slt +++ b/datafusion/sqllogictest/test_files/dml_update.slt @@ -284,3 +284,53 @@ SELECT * FROM test_update_div; statement ok DROP TABLE test_update_div; + +# Test that EXPLAIN UPDATE does not change the rows +# The rows change when the plan runs, and EXPLAIN prints the plan without +# running it. +statement ok +CREATE TABLE test_explain_update(id INT, name VARCHAR); + +statement ok +INSERT INTO test_explain_update VALUES (1, 'a'), (2, 'b'), (3, 'c'); + +query TT +EXPLAIN UPDATE test_explain_update SET name = 'z' WHERE id > 1; +---- +logical_plan +01)Dml: op=[Update] table=[test_explain_update] +02)--Projection: test_explain_update.id AS id, Utf8View("z") AS name +03)----Filter: test_explain_update.id > Int32(1) +04)------TableScan: test_explain_update projection=[id] +physical_plan +01)CooperativeExec +02)--MemDmlExec: op=Update + +query IT rowsort +SELECT * FROM test_explain_update; +---- +1 a +2 b +3 c + +# EXPLAIN still reports a statement that cannot be planned +statement error No field named nonexistent +EXPLAIN UPDATE test_explain_update SET nonexistent = 'z'; + +# EXPLAIN ANALYZE runs the plan, so the rows do change +query TT +EXPLAIN ANALYZE UPDATE test_explain_update SET name = 'z' WHERE id > 1; +---- +Plan with Metrics +01)CooperativeExec, metrics=[] +02)--MemDmlExec: op=Update, metrics=[] + +query IT rowsort +SELECT * FROM test_explain_update; +---- +1 a +2 z +3 z + +statement ok +DROP TABLE test_explain_update; diff --git a/datafusion/sqllogictest/test_files/update.slt b/datafusion/sqllogictest/test_files/update.slt index e8fdab6ab18bb..19973bb51f157 100644 --- a/datafusion/sqllogictest/test_files/update.slt +++ b/datafusion/sqllogictest/test_files/update.slt @@ -35,7 +35,7 @@ logical_plan 03)----TableScan: t1 physical_plan 01)CooperativeExec -02)--DmlResultExec: rows_affected=0 +02)--MemDmlExec: op=Update query TT explain update t1 set a=c+1, b=a, c=c+1.0, d=b; @@ -46,7 +46,7 @@ logical_plan 03)----TableScan: t1 physical_plan 01)CooperativeExec -02)--DmlResultExec: rows_affected=0 +02)--MemDmlExec: op=Update statement ok create table t2(a int, b varchar, c double, d int);