You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/16 02:57:37 UTC
(datafusion) branch main updated: Support INSERT OVERWRITE for MemTable (#24969)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 62dc7ac37b Support INSERT OVERWRITE for MemTable (#24969)
62dc7ac37b is described below
commit 62dc7ac37b31c6c9f11e91880f93eeae80277dfa
Author: Artem Osipov <59...@users.noreply.github.com>
AuthorDate: Wed Sep 16 02:48:11 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: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org