This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24969-40488988ad596c9b093ad60e1453430d803ce33c in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit ba57f8a5db742fa40503ba649796ca870c4d601e Author: Artem Osipov <[email protected]> AuthorDate: Wed Sep 9 10:14:55 2026 +0000 Support INSERT OVERWRITE for MemTable (#24969) ## Which issue does this PR close? - Part of #19617. ## Rationale for this change DataFusion parses and plans `INSERT OVERWRITE`, but `MemTable` rejects every insert operation except append. Users therefore cannot replace the contents of an in-memory table with the standard overwrite operation. ## What changes are included in this PR? - allow `InsertOp::Overwrite` in `MemTable::insert_into` - let `MemSink` replace each target partition after the input stream completes successfully - continue rejecting the distinct `InsertOp::Replace` operation ## What is the testing strategy for this PR? Three unit tests verify that overwrite replaces existing rows, distributes multiple replacement batches across multiple target partitions, and clears the table for empty input. A separate test confirms that `InsertOp::Replace` remains unsupported. Validated with: - `cargo test -p datafusion --lib test_insert_overwrite -- --nocapture` - `cargo clippy -p datafusion -p datafusion-datasource -p datafusion-catalog --all-targets --all-features -- -D warnings` - `cargo fmt --all -- --check` ## Are there any user-facing changes? Yes. `INSERT OVERWRITE` now replaces all existing data in a `MemTable`. There are no breaking public API changes. --- datafusion/catalog/src/memory/table.rs | 5 +- datafusion/core/src/datasource/memory_test.rs | 130 +++++++++++++++++++++++++- datafusion/datasource/src/memory.rs | 24 ++++- 3 files changed, 151 insertions(+), 8 deletions(-) diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index d817aa8b77..1cc7287c32 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -349,10 +349,11 @@ impl MemTable { self.schema() .logically_equivalent_names_and_types(&input.schema())?; - if insert_op != InsertOp::Append { + if insert_op == InsertOp::Replace { return not_impl_err!("{insert_op} not implemented for MemoryTable yet"); } - let sink = MemSink::try_new(self.batches.clone(), Arc::clone(&self.schema))?; + let sink = MemSink::try_new(self.batches.clone(), Arc::clone(&self.schema))? + .with_overwrite(insert_op == InsertOp::Overwrite); Ok(Arc::new(DataSinkExec::new(input, Arc::new(sink), None))) } diff --git a/datafusion/core/src/datasource/memory_test.rs b/datafusion/core/src/datasource/memory_test.rs index cc5ad539da..033a8036f5 100644 --- a/datafusion/core/src/datasource/memory_test.rs +++ b/datafusion/core/src/datasource/memory_test.rs @@ -23,7 +23,7 @@ mod tests { use crate::physical_plan::collect; use crate::prelude::SessionContext; use arrow::array::{AsArray, Int32Array}; - use arrow::datatypes::{DataType, Field, Schema, UInt64Type}; + use arrow::datatypes::{DataType, Field, Int32Type, Schema, UInt64Type}; use arrow::error::ArrowError; use arrow::record_batch::RecordBatch; use arrow_schema::SchemaRef; @@ -318,6 +318,16 @@ mod tests { schema: SchemaRef, initial_data: Vec<Vec<RecordBatch>>, inserted_data: Vec<Vec<RecordBatch>>, + ) -> Result<Vec<Vec<RecordBatch>>> { + experiment_with_insert_op(schema, initial_data, inserted_data, InsertOp::Append) + .await + } + + async fn experiment_with_insert_op( + schema: SchemaRef, + initial_data: Vec<Vec<RecordBatch>>, + inserted_data: Vec<Vec<RecordBatch>>, + insert_op: InsertOp, ) -> Result<Vec<Vec<RecordBatch>>> { let expected_count: u64 = inserted_data .iter() @@ -339,7 +349,7 @@ mod tests { let scan_plan = LogicalPlanBuilder::scan("source", source, None)?.build()?; // Create an insert plan to insert the source data into the initial table let insert_into_table = - LogicalPlanBuilder::insert_into(scan_plan, "t", target, InsertOp::Append)? + LogicalPlanBuilder::insert_into(scan_plan, "t", target, insert_op)? .build()?; // Create a physical plan from the insert plan let plan = session_ctx @@ -479,6 +489,122 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_insert_overwrite_replaces_existing_data() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let initial_batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + let replacement_batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![4, 5]))], + )?; + + let resulting_data = experiment_with_insert_op( + schema, + vec![vec![initial_batch]], + vec![vec![replacement_batch]], + InsertOp::Overwrite, + ) + .await?; + + assert_eq!(resulting_data[0].len(), 1); + assert_eq!( + resulting_data[0][0] + .column(0) + .as_primitive::<Int32Type>() + .values(), + &[4, 5] + ); + Ok(()) + } + + #[tokio::test] + async fn test_insert_overwrite_replaces_multiple_partitions() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = |values| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + ) + }; + + let resulting_data = experiment_with_insert_op( + Arc::clone(&schema), + vec![vec![batch(vec![1])?], vec![batch(vec![2])?]], + vec![vec![ + batch(vec![10])?, + batch(vec![20])?, + batch(vec![30])?, + batch(vec![40])?, + ]], + InsertOp::Overwrite, + ) + .await?; + + assert_eq!(resulting_data.len(), 2); + for (partition, expected) in resulting_data.iter().zip([[10, 30], [20, 40]]) { + let actual = partition + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::<Int32Type>() + .values() + .iter() + .copied() + }) + .collect::<Vec<_>>(); + assert_eq!(actual, expected); + } + Ok(()) + } + + #[tokio::test] + async fn test_insert_overwrite_with_empty_input_clears_table() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let initial_batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + + let resulting_data = experiment_with_insert_op( + schema, + vec![vec![initial_batch]], + vec![vec![]], + InsertOp::Overwrite, + ) + .await?; + + assert!(resulting_data[0].is_empty()); + Ok(()) + } + + #[tokio::test] + async fn test_insert_replace_remains_unsupported() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + + let error = experiment_with_insert_op( + schema, + vec![vec![batch.clone()]], + vec![vec![batch]], + InsertOp::Replace, + ) + .await + .unwrap_err(); + + assert_eq!( + error.strip_backtrace(), + "This feature is not implemented: Replace Into not implemented for MemoryTable yet" + ); + Ok(()) + } + // Test inserting a batch into a MemTable without any partitions #[tokio::test] async fn test_insert_into_zero_partition() -> Result<()> { diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index 4c79cf4a98..0d0d02a37f 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -920,6 +920,7 @@ pub struct MemSink { /// Target locations for writing data batches: Vec<PartitionData>, schema: SchemaRef, + overwrite: bool, } impl Debug for MemSink { @@ -953,7 +954,17 @@ impl MemSink { if batches.is_empty() { return plan_err!("Cannot insert into MemTable with zero partitions"); } - Ok(Self { batches, schema }) + Ok(Self { + batches, + schema, + overwrite: false, + }) + } + + /// Configures whether writes replace the existing data instead of appending to it. + pub fn with_overwrite(mut self, overwrite: bool) -> Self { + self.overwrite = overwrite; + self } } @@ -981,10 +992,15 @@ impl DataSink for MemSink { i = (i + 1) % num_partitions; } - // write the outputs into the batches + // Modify the table only after the input stream has completed successfully. for (target, mut batches) in self.batches.iter().zip(new_batches) { - // Append all the new batches in one go to minimize locking overhead - target.write().await.append(&mut batches); + let mut target = target.write().await; + if self.overwrite { + *target = batches; + } else { + // Append all the new batches in one go to minimize locking overhead + target.append(&mut batches); + } } Ok(row_count as u64) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
