get row counts from data files directly

Author: CTTY

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

@@ -187,7 +187,7 @@ impl ExecutionPlan for IcebergCommitExec {
         // Process the input streams from all partitions and commit the data files
         let stream = futures::stream::once(async move {
             let mut data_files: Vec<DataFile> = Vec::new();
-            let mut total_count: u64 = 0;
+            let mut total_record_count: u64 = 0;
 
             // Execute and collect results from all partitions of the input plan
             let batches = execute_stream_partitioned(input_plan, context)?;
@@ -197,21 +197,6 @@ impl ExecutionPlan for IcebergCommitExec {
                 while let Some(batch_result) = batch_stream.as_mut().next().await {
                     let batch = batch_result?;
 
-                    let count_array = batch
-                        .column_by_name("count")
-                        .ok_or_else(|| {
-                            DataFusionError::Internal(
-                                "Expected 'count' column in input batch".to_string(),
-                            )
-                        })?
-                        .as_any()
-                        .downcast_ref::<UInt64Array>()
-                        .ok_or_else(|| {
-                            DataFusionError::Internal(
-                                "Expected 'count' column to be UInt64Array".to_string(),
-                            )
-                        })?;
-
                     let files_array = batch
                         .column_by_name("data_files")
                         .ok_or_else(|| {
@@ -230,9 +215,6 @@ impl ExecutionPlan for IcebergCommitExec {
                     // todo remove log
                     println!("files_array to deserialize: {:?}", files_array);
 
-                    // Sum all values in the count_array
-                    total_count += count_array.iter().flatten().sum::<u64>();
-
                     // Deserialize all data files from the StringArray
                     let batch_files: Vec<DataFile> = files_array
                         .into_iter()
@@ -249,6 +231,9 @@ impl ExecutionPlan for IcebergCommitExec {
                         })
                         .collect::<datafusion::common::Result<_>>()?;
 
+                    // add record_counts from the current batch to total record count
+                    total_record_count += batch_files.iter().map(|f| f.record_count()).sum::<u64>();
+
                     // Add all deserialized files to our collection
                     data_files.extend(batch_files);
                 }
@@ -272,7 +257,7 @@ impl ExecutionPlan for IcebergCommitExec {
             //     .await
             //     .map_err(to_datafusion_error)?;
 
-            Self::make_count_batch(total_count)
+            Self::make_count_batch(total_record_count)
         })
         .boxed();
 

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

@@ -20,7 +20,7 @@ use std::fmt::{Debug, Formatter};
 use std::str::FromStr;
 use std::sync::Arc;
 
-use datafusion::arrow::array::{ArrayRef, RecordBatch, StringArray, UInt64Array};
+use datafusion::arrow::array::{ArrayRef, RecordBatch, StringArray};
 use datafusion::arrow::datatypes::{
     DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef,
 };
@@ -45,7 +45,7 @@ use iceberg::writer::file_writer::ParquetWriterBuilder;
 use iceberg::writer::file_writer::location_generator::{
     DefaultFileNameGenerator, DefaultLocationGenerator,
 };
-use iceberg::writer::{CurrentFileStatus, IcebergWriter, IcebergWriterBuilder};
+use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
 use iceberg::{Error, ErrorKind};
 use parquet::file::properties::WriterProperties;
 use uuid::Uuid;
@@ -83,26 +83,22 @@ impl IcebergWriteExec {
         )
     }
 
-    // Create a record batch with count and serialized data files
-    fn make_result_batch(count: u64, data_files: Vec<String>) -> DFResult<RecordBatch> {
-        let count_array = Arc::new(UInt64Array::from(vec![count])) as ArrayRef;
+    // Create a record batch with serialized data files
+    fn make_result_batch(data_files: Vec<String>) -> DFResult<RecordBatch> {
         let files_array = Arc::new(StringArray::from(data_files)) as ArrayRef;
 
-        RecordBatch::try_from_iter_with_nullable(vec![
-            ("count", count_array, false),
-            ("data_files", files_array, false),
-        ])
-        .map_err(|e| {
-            DataFusionError::ArrowError(e, Some("Failed to make result batch".to_string()))
-        })
+        RecordBatch::try_from_iter_with_nullable(vec![("data_files", files_array, false)]).map_err(
+            |e| DataFusionError::ArrowError(e, Some("Failed to make result batch".to_string())),
+        )
     }
 
     fn make_result_schema() -> ArrowSchemaRef {
         // Define a schema.
-        Arc::new(ArrowSchema::new(vec![
-            Field::new("count", DataType::UInt64, false),
-            Field::new("data_files", DataType::Utf8, false),
-        ]))
+        Arc::new(ArrowSchema::new(vec![Field::new(
+            "data_files",
+            DataType::Utf8,
+            false,
+        )]))
     }
 }
 
@@ -238,7 +234,6 @@ impl ExecutionPlan for IcebergWriteExec {
                 writer.write(batch?).await.map_err(to_datafusion_error)?;
             }
 
-            let count = writer.current_row_num() as u64;
             let data_file_builders = writer.close().await.map_err(to_datafusion_error)?;
 
             // Convert builders to data files and then to JSON strings
@@ -255,7 +250,7 @@ impl ExecutionPlan for IcebergWriteExec {
                 })
                 .collect::<DFResult<Vec<String>>>()?;
 
-            Self::make_result_batch(count, data_files)
+            Self::make_result_batch(data_files)
         })
         .boxed();