feat: integrate StatisticsCollector into parquet writes

- Import StatisticsCollector in parquet.rs
- Add stats field to DataFileMetadata with with_stats() method
- Update as_record_batch to use full stats if available
- Update write_parquet_file to collect and attach stats
- Update mod.rs write_parquet to pass stats_columns
- Update write tests to expect full stats output
- Fix write-table example API call

Author: DrakeLin

kernel/examples/write-table/src/main.rs

@@ -94,9 +94,9 @@ async fn try_main() -> DeltaResult<()> {
         .with_data_change(true);
 
     // Write the data using the engine
-    let write_context = Arc::new(txn.get_write_context());
+    let write_context = txn.get_write_context()?;
     let file_metadata = engine
-        .write_parquet(&sample_data, write_context.as_ref(), HashMap::new())
+        .write_parquet(&sample_data, &write_context, HashMap::new())
         .await?;
 
     // Add the file metadata to the transaction

kernel/src/engine/default/mod.rs

@@ -217,7 +217,12 @@ impl<E: TaskExecutor> DefaultEngine<E> {
         )?;
         let physical_data = logical_to_physical_expr.evaluate(data)?;
         self.parquet
-            .write_parquet_file(write_context.target_dir(), physical_data, partition_values)
+            .write_parquet_file(
+                write_context.target_dir(),
+                physical_data,
+                partition_values,
+                write_context.stats_columns(),
+            )
             .await
     }
 }

kernel/src/engine/default/parquet.rs

@@ -7,7 +7,7 @@ use std::sync::Arc;
 use delta_kernel_derive::internal_api;
 
 use crate::arrow::array::builder::{MapBuilder, MapFieldNames, StringBuilder};
-use crate::arrow::array::{Int64Array, RecordBatch, StringArray, StructArray};
+use crate::arrow::array::{Array, Int64Array, RecordBatch, StringArray, StructArray};
 use crate::arrow::datatypes::{DataType, Field};
 use crate::parquet::arrow::arrow_reader::{
     ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder,
@@ -23,9 +23,10 @@ use object_store::{DynObjectStore, ObjectStore};
 use uuid::Uuid;
 
 use super::file_stream::{FileOpenFuture, FileOpener, FileStream};
+use super::stats::StatisticsCollector;
 use super::UrlExt;
 use crate::engine::arrow_conversion::{TryFromArrow as _, TryIntoArrow as _};
-use crate::engine::arrow_data::ArrowEngineData;
+use crate::engine::arrow_data::{extract_record_batch, ArrowEngineData};
 use crate::engine::arrow_utils::{
     fixup_parquet_read, generate_mask, get_requested_indices, ordering_needs_row_indexes,
     RowIndexBuilder,
@@ -54,16 +55,25 @@ pub struct DataFileMetadata {
     file_meta: FileMeta,
     // NB: We use usize instead of u64 since arrow uses usize for record batch sizes
     num_records: usize,
+    /// Collected statistics for this file (optional).
+    stats: Option<StructArray>,
 }
 
 impl DataFileMetadata {
     pub fn new(file_meta: FileMeta, num_records: usize) -> Self {
         Self {
             file_meta,
             num_records,
+            stats: None,
         }
     }
 
+    /// Set the collected statistics for this file.
+    pub fn with_stats(mut self, stats: StructArray) -> Self {
+        self.stats = Some(stats);
+        self
+    }
+
     /// Convert DataFileMetadata into a record batch which matches the schema returned by
     /// [`add_files_schema`].
     ///
@@ -81,6 +91,7 @@ impl DataFileMetadata {
                     size,
                 },
             num_records,
+            stats,
         } = self;
         // create the record batch of the write metadata
         let path = Arc::new(StringArray::from(vec![location.to_string()]));
@@ -104,20 +115,53 @@ impl DataFileMetadata {
             .map_err(|_| Error::generic("Failed to convert parquet metadata 'size' to i64"))?;
         let size = Arc::new(Int64Array::from(vec![size]));
         let modification_time = Arc::new(Int64Array::from(vec![*last_modified]));
-        let stats = Arc::new(StructArray::try_new_with_length(
-            vec![Field::new("numRecords", DataType::Int64, true)].into(),
-            vec![Arc::new(Int64Array::from(vec![*num_records as i64]))],
-            None,
-            1,
-        )?);
 
-        Ok(Box::new(ArrowEngineData::new(RecordBatch::try_new(
-            Arc::new(
-                crate::transaction::BASE_ADD_FILES_SCHEMA
-                    .as_ref()
-                    .try_into_arrow()?,
+        // Use full stats if available, otherwise just numRecords
+        let stats_array: Arc<StructArray> = if let Some(full_stats) = stats {
+            Arc::new(full_stats.clone())
+        } else {
+            Arc::new(StructArray::try_new_with_length(
+                vec![Field::new("numRecords", DataType::Int64, true)].into(),
+                vec![Arc::new(Int64Array::from(vec![*num_records as i64]))],
+                None,
+                1,
+            )?)
+        };
+
+        // Build schema dynamically based on stats
+        let stats_field = Field::new("stats", stats_array.data_type().clone(), true);
+        let schema = crate::arrow::datatypes::Schema::new(vec![
+            Field::new("path", crate::arrow::datatypes::DataType::Utf8, false),
+            Field::new(
+                "partitionValues",
+                crate::arrow::datatypes::DataType::Map(
+                    Arc::new(Field::new(
+                        "key_value",
+                        crate::arrow::datatypes::DataType::Struct(
+                            vec![
+                                Field::new("key", crate::arrow::datatypes::DataType::Utf8, false),
+                                Field::new("value", crate::arrow::datatypes::DataType::Utf8, true),
+                            ]
+                            .into(),
+                        ),
+                        false,
+                    )),
+                    false,
+                ),
+                false,
+            ),
+            Field::new("size", crate::arrow::datatypes::DataType::Int64, false),
+            Field::new(
+                "modificationTime",
+                crate::arrow::datatypes::DataType::Int64,
+                false,
             ),
-            vec![path, partitions, size, modification_time, stats],
+            stats_field,
+        ]);
+
+        Ok(Box::new(ArrowEngineData::new(RecordBatch::try_new(
+            Arc::new(schema),
+            vec![path, partitions, size, modification_time, stats_array],
         )?)))
     }
 }
@@ -201,8 +245,22 @@ impl<E: TaskExecutor> DefaultParquetHandler<E> {
         path: &url::Url,
         data: Box<dyn EngineData>,
         partition_values: HashMap<String, String>,
+        stats_columns: &[String],
     ) -> DeltaResult<Box<dyn EngineData>> {
-        let parquet_metadata = self.write_parquet(path, data).await?;
+        // Collect statistics from the data during write
+        let record_batch = extract_record_batch(data.as_ref())?;
+
+        // Initialize stats collector and update with this batch
+        let mut stats_collector = StatisticsCollector::new(record_batch.schema(), stats_columns);
+        stats_collector.update(record_batch, None)?; // No mask for new file writes
+        let stats = stats_collector.finalize()?;
+
+        // Write the parquet file
+        let mut parquet_metadata = self.write_parquet(path, data).await?;
+
+        // Attach the collected statistics
+        parquet_metadata = parquet_metadata.with_stats(stats);
+
         parquet_metadata.as_record_batch(&partition_values)
     }
 }
@@ -684,6 +742,7 @@ mod tests {
                     size,
                 },
             num_records,
+            ..
         } = write_metadata;
         let expected_location = Url::parse("memory:///data/").unwrap();
 

kernel/tests/write.rs

@@ -393,7 +393,7 @@ async fn test_append() -> Result<(), Box<dyn std::error::Error>> {
                     "size": size,
                     "modificationTime": 0,
                     "dataChange": true,
-                    "stats": "{\"numRecords\":3}"
+                    "stats": "{\"numRecords\":3,\"nullCount\":{\"number\":0},\"minValues\":{\"number\":1},\"maxValues\":{\"number\":3},\"tightBounds\":true}"
                 }
             }),
             json!({
@@ -403,7 +403,7 @@ async fn test_append() -> Result<(), Box<dyn std::error::Error>> {
                     "size": size,
                     "modificationTime": 0,
                     "dataChange": true,
-                    "stats": "{\"numRecords\":3}"
+                    "stats": "{\"numRecords\":3,\"nullCount\":{\"number\":0},\"minValues\":{\"number\":4},\"maxValues\":{\"number\":6},\"tightBounds\":true}"
                 }
             }),
         ];
@@ -601,7 +601,7 @@ async fn test_append_partitioned() -> Result<(), Box<dyn std::error::Error>> {
                     "size": size,
                     "modificationTime": 0,
                     "dataChange": false,
-                    "stats": "{\"numRecords\":3}"
+                    "stats": "{\"numRecords\":3,\"nullCount\":{\"number\":0},\"minValues\":{\"number\":1},\"maxValues\":{\"number\":3},\"tightBounds\":true}"
                 }
             }),
             json!({
@@ -613,7 +613,7 @@ async fn test_append_partitioned() -> Result<(), Box<dyn std::error::Error>> {
                     "size": size,
                     "modificationTime": 0,
                     "dataChange": false,
-                    "stats": "{\"numRecords\":3}"
+                    "stats": "{\"numRecords\":3,\"nullCount\":{\"number\":0},\"minValues\":{\"number\":4},\"maxValues\":{\"number\":6},\"tightBounds\":true}"
                 }
             }),
         ];
@@ -1078,6 +1078,7 @@ async fn test_append_variant() -> Result<(), Box<dyn std::error::Error>> {
             write_context.target_dir(),
             Box::new(ArrowEngineData::new(data.clone())),
             HashMap::new(),
+            write_context.stats_columns(),
         )
         .await?;
 
@@ -1251,6 +1252,7 @@ async fn test_shredded_variant_read_rejection() -> Result<(), Box<dyn std::error
             write_context.target_dir(),
             Box::new(ArrowEngineData::new(data.clone())),
             HashMap::new(),
+            write_context.stats_columns(),
         )
         .await?;