Support INSERT OVERWRITE for MemTable

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)))
     }
 

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,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::<Int32Type>()
+                .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<()> {

datafusion/datasource/src/memory.rs

@@ -879,6 +879,7 @@ pub struct MemSink {
     /// Target locations for writing data
     batches: Vec<PartitionData>,
     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)