implement children and with_new_children for write node, fix fmt

Author: CTTY

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

@@ -27,7 +27,7 @@ use datafusion::arrow::datatypes::{
 use datafusion::common::Result as DFResult;
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
-use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
@@ -61,7 +61,7 @@ pub(crate) struct IcebergWriteExec {
 
 impl IcebergWriteExec {
     pub fn new(table: Table, input: Arc<dyn ExecutionPlan>, schema: ArrowSchemaRef) -> Self {
-        let plan_properties = Self::compute_properties(&input, schema.clone());
+        let plan_properties = Self::compute_properties(&input, schema);
 
         Self {
             table,
@@ -77,7 +77,7 @@ impl IcebergWriteExec {
     ) -> PlanProperties {
         PlanProperties::new(
             EquivalenceProperties::new(schema),
-            input.output_partitioning().clone(),
+            Partitioning::UnknownPartitioning(input.output_partitioning().partition_count()),
             input.pipeline_behavior(),
             input.boundedness(),
         )
@@ -147,14 +147,24 @@ impl ExecutionPlan for IcebergWriteExec {
     }
 
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
-        vec![]
+        vec![&self.input]
     }
 
     fn with_new_children(
         self: Arc<Self>,
-        _children: Vec<Arc<dyn ExecutionPlan>>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> DFResult<Arc<dyn ExecutionPlan>> {
-        Ok(self)
+        if children.len() != 1 {
+            return Err(DataFusionError::Internal(
+                "IcebergWriteExec expects exactly one child".to_string(),
+            ));
+        }
+
+        Ok(Arc::new(Self::new(
+            self.table.clone(),
+            Arc::clone(&children[0]),
+            self.schema(),
+        )))
     }
 
     fn execute(

crates/integrations/datafusion/src/table/mod.rs

@@ -162,8 +162,7 @@ impl TableProvider for IcebergTableProvider {
     fn supports_filters_pushdown(
         &self,
         filters: &[&Expr],
-    ) -> std::result::Result<Vec<TableProviderFilterPushDown>, DataFusionError>
-    {
+    ) -> std::result::Result<Vec<TableProviderFilterPushDown>, DataFusionError> {
         // Push down all filters, as a single source of truth, the scanner will drop the filters which couldn't be push down
         Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
     }
@@ -174,10 +173,16 @@ impl TableProvider for IcebergTableProvider {
         input: Arc<dyn ExecutionPlan>,
         _insert_op: InsertOp,
     ) -> DFResult<Arc<dyn ExecutionPlan>> {
-        if !self.table.metadata().default_partition_spec().is_unpartitioned() {
+        if !self
+            .table
+            .metadata()
+            .default_partition_spec()
+            .is_unpartitioned()
+        {
             // TODO add insert into support for partitioned tables
             return Err(DataFusionError::NotImplemented(
-                "IcebergTableProvider::insert_into does not support partitioned tables yet".to_string()
+                "IcebergTableProvider::insert_into does not support partitioned tables yet"
+                    .to_string(),
             ));
         }