diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index d817aa8b7788a..1cc7287c32cf2 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 cc5ad539dae71..36956aa19d57b 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>, inserted_data: Vec>, + ) -> Result>> { + experiment_with_insert_op(schema, initial_data, inserted_data, InsertOp::Append) + .await + } + + async fn experiment_with_insert_op( + schema: SchemaRef, + initial_data: Vec>, + inserted_data: Vec>, + insert_op: InsertOp, ) -> Result>> { 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,81 @@ 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::() + .values(), + &[4, 5] + ); + 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 7c10dba981c82..5caee6656661b 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -879,6 +879,7 @@ pub struct MemSink { /// Target locations for writing data batches: Vec, schema: SchemaRef, + overwrite: bool, } impl Debug for MemSink { @@ -912,7 +913,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 } } @@ -940,10 +951,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)