i luv cleaning up

Author: CTTY

crates/integrations/datafusion/Cargo.toml

@@ -37,6 +37,7 @@ iceberg = { workspace = true }
 parquet = { workspace = true }
 tokio = { workspace = true }
 serde_json = { workspace = true }
+uuid = { workspace = true }
 
 [dev-dependencies]
 expect-test = { workspace = true }

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

@@ -29,8 +29,11 @@ use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
-use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, execute_input_stream,
+};
 use futures::StreamExt;
+use iceberg::arrow::schema_to_arrow_schema;
 use iceberg::spec::{DataFile, DataFileFormat, DataFileSerde, FormatVersion};
 use iceberg::table::Table;
 use iceberg::writer::CurrentFileStatus;
@@ -39,6 +42,7 @@ use iceberg::writer::file_writer::location_generator::{
 };
 use iceberg::writer::file_writer::{FileWriter, FileWriterBuilder, ParquetWriterBuilder};
 use parquet::file::properties::WriterProperties;
+use uuid::Uuid;
 
 use crate::to_datafusion_error;
 
@@ -159,14 +163,24 @@ impl ExecutionPlan for IcebergWriteExec {
             self.table.file_io().clone(),
             DefaultLocationGenerator::new(self.table.metadata().clone())
                 .map_err(to_datafusion_error)?,
-            // todo actual filename
-            DefaultFileNameGenerator::new("what".to_string(), None, DataFileFormat::Parquet),
+            // todo filename prefix/suffix should be configurable
+            DefaultFileNameGenerator::new(
+                "datafusion".to_string(),
+                Some(Uuid::now_v7().to_string()),
+                DataFileFormat::Parquet,
+            ),
         )
         .build();
 
-        // todo repartition
-        let data = self.input.execute(partition, context)?;
-        let result_schema = Arc::clone(&self.result_schema);
+        let data = execute_input_stream(
+            Arc::clone(&self.input),
+            Arc::new(
+                schema_to_arrow_schema(self.table.metadata().current_schema())
+                    .map_err(to_datafusion_error)?,
+            ),
+            partition,
+            Arc::clone(&context),
+        )?;
 
         // todo non-default partition spec?
         let spec_id = self.table.metadata().default_partition_spec_id();
@@ -175,7 +189,6 @@ impl ExecutionPlan for IcebergWriteExec {
 
         let stream = futures::stream::once(async move {
             let mut writer = parquet_writer_fut.await.map_err(to_datafusion_error)?;
-
             let mut input_stream = data;
 
             while let Some(batch_res) = input_stream.next().await {
@@ -187,27 +200,45 @@ impl ExecutionPlan for IcebergWriteExec {
             let data_file_builders = writer.close().await.map_err(to_datafusion_error)?;
 
             // Convert builders to data files
-            let data_files = data_file_builders
+            let data_files: DFResult<Vec<DataFile>> = data_file_builders
                 .into_iter()
-                .map(|mut builder| builder.partition_spec_id(spec_id).build().unwrap())
-                .collect::<Vec<DataFile>>();
+                .map(|mut builder| {
+                    builder.partition_spec_id(spec_id).build().map_err(|e| {
+                        DataFusionError::Execution(format!("Failed to build data file: {}", e))
+                    })
+                })
+                .collect();
+            let data_files = data_files?;
 
-            let data_files = data_files
+            let data_files: DFResult<Vec<String>> = data_files
                 .into_iter()
                 .map(|f| {
-                    let serde = DataFileSerde::try_from(f, &partition_type, is_version_1).unwrap();
-                    let json = serde_json::to_string(&serde).unwrap();
+                    // Convert to DataFileSerde
+                    let serde =
+                        DataFileSerde::try_from(f, &partition_type, is_version_1).map_err(|e| {
+                            DataFusionError::Execution(format!(
+                                "Failed to convert to DataFileSerde: {}",
+                                e
+                            ))
+                        })?;
+
+                    // Serialize to JSON
+                    let json = serde_json::to_string(&serde).map_err(|e| {
+                        DataFusionError::Execution(format!("Failed to serialize to JSON: {}", e))
+                    })?;
+
                     println!("Serialized data file: {}", json); // todo remove log
-                    json
+                    Ok(json)
                 })
-                .collect::<Vec<String>>();
+                .collect();
+            let data_files = data_files?;
 
             Ok(Self::make_result_batch(count, data_files)?)
         })
         .boxed();
 
         Ok(Box::pin(RecordBatchStreamAdapter::new(
-            result_schema,
+            Arc::clone(&self.result_schema),
             stream,
         )))
     }