fix: prevent DataFusion RepartitionExec insertion into IcebergCommitExec

Ensures DataFusion coalesces input partitions automatically before commit, fixing partitioning issues in Iceberg write operations (partial writes).

Author: sgrebnov

crates/integrations/datafusion/src/physical_plan/commit.rs

@@ -139,6 +139,14 @@ impl ExecutionPlan for IcebergCommitExec {
         vec![&self.input]
     }
 
+    fn required_input_distribution(&self) -> Vec<datafusion::physical_plan::Distribution> {
+        vec![datafusion::physical_plan::Distribution::SinglePartition; self.children().len()]
+    }
+
+    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+        vec![false]
+    }
+
     fn with_new_children(
         self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
@@ -265,14 +273,16 @@ mod tests {
     use std::fmt;
     use std::sync::Arc;
 
-    use datafusion::arrow::array::{ArrayRef, RecordBatch, StringArray};
+    use datafusion::arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray, UInt64Array};
     use datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
+    use datafusion::datasource::MemTable;
     use datafusion::execution::context::TaskContext;
     use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
     use datafusion::physical_plan::common::collect;
     use datafusion::physical_plan::execution_plan::Boundedness;
     use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
     use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
+    use datafusion::prelude::*;
     use futures::StreamExt;
     use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};
     use iceberg::spec::{
@@ -283,6 +293,7 @@ mod tests {
 
     use super::*;
     use crate::physical_plan::DATA_FILES_COL_NAME;
+    use crate::table::IcebergTableProvider;
 
     // A mock execution plan that returns record batches with serialized data files
     #[derive(Debug)]
@@ -513,4 +524,99 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_datafusion_execution_partitioned_source() -> Result<(), Box<dyn std::error::Error>>
+    {
+        let catalog = Arc::new(
+            MemoryCatalogBuilder::default()
+                .load(
+                    "memory",
+                    HashMap::from([(
+                        MEMORY_CATALOG_WAREHOUSE.to_string(),
+                        "memory://root".to_string(),
+                    )]),
+                )
+                .await?,
+        );
+
+        let namespace = NamespaceIdent::new("test_namespace".to_string());
+        catalog.create_namespace(&namespace, HashMap::new()).await?;
+
+        let schema = Schema::builder()
+            .with_schema_id(1)
+            .with_fields(vec![
+                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
+                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
+            ])
+            .build()?;
+
+        let table_name = "test_table";
+        let table_creation = TableCreation::builder()
+            .name(table_name.to_string())
+            .schema(schema)
+            .location("memory://root/test_table".to_string())
+            .properties(HashMap::new())
+            .build();
+        let _ = catalog.create_table(&namespace, table_creation).await?;
+
+        let arrow_schema = Arc::new(ArrowSchema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("name", DataType::Utf8, false),
+        ]));
+
+        let batches: Vec<RecordBatch> = (1..4)
+            .map(|idx| {
+                RecordBatch::try_new(arrow_schema.clone(), vec![
+                    Arc::new(Int32Array::from(vec![idx])) as ArrayRef,
+                    Arc::new(StringArray::from(vec![format!("Name{idx}")])) as ArrayRef,
+                ])
+            })
+            .collect::<Result<_, _>>()?;
+
+        // Create DataFusion context with specific partition configuration
+        let mut config = SessionConfig::new();
+        config = config.set_usize("datafusion.execution.target_partitions", 8);
+        let ctx = SessionContext::new_with_config(config);
+
+        // Create multiple partitions - each batch becomes a separate partition
+        let partitions: Vec<Vec<RecordBatch>> =
+            batches.into_iter().map(|batch| vec![batch]).collect();
+        let source_table = Arc::new(MemTable::try_new(Arc::clone(&arrow_schema), partitions)?);
+        ctx.register_table("source_table", source_table)?;
+
+        let table_ident = TableIdent::from_strs(vec!["test_namespace", "test_table"])?;
+        let iceberg_table_provider =
+            IcebergTableProvider::try_new(catalog.clone(), table_ident).await?;
+        ctx.register_table("iceberg_table", Arc::new(iceberg_table_provider))?;
+
+        let insert_plan = ctx
+            .sql("INSERT INTO iceberg_table SELECT * FROM source_table")
+            .await?;
+
+        let physical_plan = insert_plan.create_physical_plan().await?;
+
+        let actual_plan = format!(
+            "{}",
+            datafusion::physical_plan::displayable(physical_plan.as_ref()).indent(false)
+        );
+
+        println!("Physical plan:\n{actual_plan}");
+
+        let expected_plan = "\
+IcebergCommitExec: table=test_namespace.test_table
+  CoalescePartitionsExec
+    IcebergWriteExec: table=test_namespace.test_table
+      DataSourceExec: partitions=3, partition_sizes=[1, 1, 1]";
+
+        assert_eq!(
+            actual_plan.trim(),
+            expected_plan.trim(),
+            "Physical plan does not match expected\n\nExpected:\n{}\n\nActual:\n{}",
+            expected_plan.trim(),
+            actual_plan.trim()
+        );
+
+        Ok(())
+    }
 }

crates/integrations/datafusion/src/physical_plan/write.rs

@@ -144,6 +144,18 @@ impl ExecutionPlan for IcebergWriteExec {
         self
     }
 
+    /// Prevents the introduction of additional `RepartitionExec` and processing input in parallel.
+    /// This guarantees that the input is processed as a single stream, preserving the order of the data.
+    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+        vec![false]
+    }
+
+    fn maintains_input_order(&self) -> Vec<bool> {
+        // Maintains ordering in the sense that the written file will reflect
+        // the ordering of the input.
+        vec![true; self.children().len()]
+    }
+
     fn properties(&self) -> &PlanProperties {
         &self.plan_properties
     }
@@ -284,12 +296,23 @@ impl ExecutionPlan for IcebergWriteExec {
                 .await
                 .map_err(to_datafusion_error)?;
             let mut input_stream = data;
+            let mut total_rows = 0;
+            let mut batch_count = 0;
 
             while let Some(batch) = input_stream.next().await {
-                writer.write(batch?).await.map_err(to_datafusion_error)?;
+                batch_count += 1;
+                let batch = batch?;
+                let rows = batch.num_rows();
+                total_rows += rows;
+                println!("Batch {}: {} rows", batch_count, rows);
+                println!("Batch schema: {:?}", batch.schema());
+                writer.write(batch).await.map_err(to_datafusion_error)?;
             }
+            
+            println!("Total: {} batches, {} rows", batch_count, total_rows);
 
             let data_files = writer.close().await.map_err(to_datafusion_error)?;
+            println!("Wrote {} data files", data_files.len());
 
             // Convert builders to data files and then to JSON strings
             let data_files_strs: Vec<String> = data_files
@@ -300,6 +323,8 @@ impl ExecutionPlan for IcebergWriteExec {
                 })
                 .collect::<DFResult<Vec<String>>>()?;
 
+            println!("Serialized data files: {:?}", data_files_strs);
+
             Self::make_result_batch(data_files_strs)
         })
         .boxed();
@@ -617,4 +642,6 @@ mod tests {
 
         Ok(())
     }
-}
+
+    
+}