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,
  )))
  }