diff --git a/datafusion_iceberg/src/table/mod.rs b/datafusion_iceberg/src/table/mod.rs index d7e9c050..605676ad 100644 --- a/datafusion_iceberg/src/table/mod.rs +++ b/datafusion_iceberg/src/table/mod.rs @@ -385,12 +385,16 @@ impl TableProvider for DataFusionTable { if !self.schema().equivalent_names_and_types(&input.schema()) { return plan_err!("Inserting query must have the same schema with the table."); } - let InsertOp::Append = insert_op else { - return not_impl_err!("Overwrite not implemented for MemoryTable yet"); + let write_operation = match insert_op { + InsertOp::Append => IcebergWriteOperation::Append, + InsertOp::Overwrite => IcebergWriteOperation::Overwrite, + InsertOp::Replace => { + return not_impl_err!("REPLACE INTO is not implemented for Iceberg tables"); + } }; Ok(Arc::new(DataSinkExec::new( input, - Arc::new(self.clone().into_data_sink()), + Arc::new(self.clone().into_data_sink(write_operation)), None, ))) } @@ -1291,18 +1295,30 @@ impl DisplayAs for DataFusionTable { } } +#[derive(Debug, Clone, Copy)] +enum IcebergWriteOperation { + Append, + Overwrite, +} + #[derive(Debug)] -pub(crate) struct IcebergDataSink(DataFusionTable); +pub(crate) struct IcebergDataSink { + table: DataFusionTable, + operation: IcebergWriteOperation, +} impl DataFusionTable { - pub(crate) fn into_data_sink(self) -> IcebergDataSink { - IcebergDataSink(self) + fn into_data_sink(self, operation: IcebergWriteOperation) -> IcebergDataSink { + IcebergDataSink { + table: self, + operation, + } } } impl DisplayAs for IcebergDataSink { fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { - self.0.fmt_as(t, f) + self.table.fmt_as(t, f) } } @@ -1315,7 +1331,7 @@ impl DataSink for IcebergDataSink { ) -> Result { // Clone the table from the read lock let mut table = { - let lock = self.0.tabular.read().unwrap(); + let lock = self.table.tabular.read().unwrap(); if let Tabular::Table(table) = lock.deref() { Ok(table.clone()) } else { @@ -1325,7 +1341,7 @@ impl DataSink for IcebergDataSink { }; let metadata_files = - write_parquet_data_files(&table, data, context, self.0.branch.as_deref()).await?; + write_parquet_data_files(&table, data, context, self.table.branch.as_deref()).await?; let written_rows = metadata_files.iter().try_fold(0_u64, |total, file| { let file_rows = u64::try_from(*file.record_count()).map_err(|_| { DataFusionError::Execution(format!( @@ -1340,15 +1356,18 @@ impl DataSink for IcebergDataSink { }) })?; - table - .new_transaction(self.0.branch.as_deref()) - .append_data(metadata_files) + let transaction = table.new_transaction(self.table.branch.as_deref()); + let transaction = match self.operation { + IcebergWriteOperation::Append => transaction.append_data(metadata_files), + IcebergWriteOperation::Overwrite => transaction.replace(metadata_files), + }; + transaction .commit() .await .map_err(DataFusionIcebergError::from)?; // Acquire write lock and overwrite the old table with the new one - let mut lock = self.0.tabular.write().unwrap(); + let mut lock = self.table.tabular.write().unwrap(); *lock = Tabular::Table(table); Ok(written_rows) @@ -1357,7 +1376,7 @@ impl DataSink for IcebergDataSink { None } fn schema(&self) -> &SchemaRef { - &self.0.schema + &self.table.schema } } diff --git a/datafusion_iceberg/tests/integration_df_dml.rs b/datafusion_iceberg/tests/integration_df_dml.rs index 2876707b..b6a978dd 100644 --- a/datafusion_iceberg/tests/integration_df_dml.rs +++ b/datafusion_iceberg/tests/integration_df_dml.rs @@ -4,9 +4,8 @@ //! //! Notes on divergences: //! - DataFusion + iceberg-rust supports `INSERT INTO ... VALUES`, -//! `INSERT INTO ... SELECT`, but does not currently support -//! `INSERT OVERWRITE` or partition-targeted overwrites. Those upstream -//! scenarios are documented in `integration_df_dml_unsupported.rs`. +//! `INSERT INTO ... SELECT`, and full-table `INSERT OVERWRITE`. +//! Partition-targeted overwrites are not currently supported. //! - DELETE / UPDATE / MERGE are not supported by iceberg-rust yet — they //! live in `integration_df_unsupported_row_dml.rs`. @@ -169,3 +168,77 @@ async fn integration_df_insert_then_select_specific_column() { let sum = execute_scalar_i64(&ctx, "SELECT SUM(id) FROM warehouse.dml_select_col.t").await; assert_eq!(sum, 60); } + +#[tokio::test] +async fn integration_df_insert_overwrite_replaces_all_rows() { + let ctx = boot_df_stack().await; + setup_target_table(&ctx, "dml_overwrite", "target").await; + setup_target_table(&ctx, "dml_overwrite", "source").await; + execute_sql( + &ctx, + "INSERT INTO warehouse.dml_overwrite.target VALUES \ + (1, 'old-a'), (2, 'old-b'), (3, 'old-c')", + ) + .await; + execute_sql( + &ctx, + "INSERT INTO warehouse.dml_overwrite.source VALUES \ + (20, 'new-a'), (30, 'new-b')", + ) + .await; + + execute_sql( + &ctx, + "INSERT OVERWRITE INTO warehouse.dml_overwrite.target \ + SELECT id, label FROM warehouse.dml_overwrite.source", + ) + .await; + + let count = + execute_scalar_i64(&ctx, "SELECT COUNT(*) FROM warehouse.dml_overwrite.target").await; + let sum = execute_scalar_i64(&ctx, "SELECT SUM(id) FROM warehouse.dml_overwrite.target").await; + let exact_rows = execute_scalar_i64( + &ctx, + "SELECT COUNT(*) FROM warehouse.dml_overwrite.target \ + WHERE (id = 20 AND label = 'new-a') OR (id = 30 AND label = 'new-b')", + ) + .await; + let old_rows = execute_scalar_i64( + &ctx, + "SELECT COUNT(*) FROM warehouse.dml_overwrite.target WHERE id IN (1, 2, 3)", + ) + .await; + assert_eq!((count, sum, exact_rows, old_rows), (2, 50, 2, 0)); +} + +#[tokio::test] +async fn integration_df_insert_overwrite_can_empty_table() { + let ctx = boot_df_stack().await; + setup_target_table(&ctx, "dml_overwrite_empty", "target").await; + setup_target_table(&ctx, "dml_overwrite_empty", "source").await; + execute_sql( + &ctx, + "INSERT INTO warehouse.dml_overwrite_empty.target VALUES \ + (1, 'old-a'), (2, 'old-b')", + ) + .await; + execute_sql( + &ctx, + "INSERT INTO warehouse.dml_overwrite_empty.source VALUES (10, 'source')", + ) + .await; + + execute_sql( + &ctx, + "INSERT OVERWRITE INTO warehouse.dml_overwrite_empty.target \ + SELECT id, label FROM warehouse.dml_overwrite_empty.source WHERE id < 0", + ) + .await; + + let count = execute_scalar_i64( + &ctx, + "SELECT COUNT(*) FROM warehouse.dml_overwrite_empty.target", + ) + .await; + assert_eq!(count, 0); +}