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/09 17:43:01 UTC
(datafusion) branch gh-readonly-queue/main/pr-24969-40488988ad596c9b093ad60e1453430d803ce33c created (now ba57f8a5db)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-24969-40488988ad596c9b093ad60e1453430d803ce33c
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at ba57f8a5db Support INSERT OVERWRITE for MemTable (#24969)
This branch includes the following new commits:
new ba57f8a5db Support INSERT OVERWRITE for MemTable (#24969)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: Support INSERT OVERWRITE for MemTable (#24969)
Posted by gi...@apache.org.
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 <59...@users.noreply.github.com>
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: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org