feat: write value stats for append data files (#477)

Author: QuakeWang

crates/integrations/datafusion/src/merge_into.rs

@@ -26,7 +26,9 @@ use std::collections::{HashMap, HashSet};
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
-use datafusion::arrow::array::{Array, Int32Array, RecordBatch, UInt32Array, UInt64Array};
+use datafusion::arrow::array::{
+    Array, ArrayRef, Int32Array, RecordBatch, UInt32Array, UInt64Array,
+};
 use datafusion::arrow::compute;
 use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema};
 use datafusion::error::{DataFusionError, Result as DFResult};
@@ -37,7 +39,7 @@ use datafusion::sql::sqlparser::ast::{
 };
 use futures::TryStreamExt;
 
-use paimon::spec::{datums_to_binary_row, extract_datum_from_arrow, CoreOptions};
+use paimon::spec::{datums_to_binary_row, extract_datum_from_arrow, CoreOptions, DataField};
 use paimon::table::{CopyOnWriteMergeWriter, DataSplitBuilder, Table, WriteBuilder};
 
 use crate::error::to_datafusion_error;
@@ -569,13 +571,6 @@ async fn execute_cow_merge_inner(
 
     // Handle NOT MATCHED → INSERT
     if !clauses.inserts.is_empty() {
-        let table_fields: Vec<String> = table
-            .schema()
-            .fields()
-            .iter()
-            .map(|f| f.name().to_string())
-            .collect();
-
         let insert_sql = if has_target_data {
             format!(
                 "SELECT {s_alias}.* FROM {source_ref} AS {s_alias} \
@@ -595,7 +590,7 @@ async fn execute_cow_merge_inner(
                 &clauses.inserts,
                 s_alias,
                 &[],
-                &table_fields,
+                table.schema().fields(),
                 temp_tracker,
             )
             .await?;
@@ -734,21 +729,14 @@ async fn execute_merge_into_once(
                 injected_columns.push(format!("__upd_{col}"));
             }
         }
-        // Table schema field names for reordering INSERT columns
-        let table_fields: Vec<String> = table
-            .schema()
-            .fields()
-            .iter()
-            .map(|f| f.name().to_string())
-            .collect();
         let mut temp_tracker = TempTableTracker::new(ctx);
         let insert_batches = build_insert_batches(
             ctx,
             &not_matched_batches,
             &parsed.inserts,
             s_alias,
             &injected_columns,
-            &table_fields,
+            table.schema().fields(),
             &mut temp_tracker,
         )
         .await?;
@@ -854,7 +842,7 @@ async fn build_insert_batches(
     inserts: &[MergeInsertClause],
     s_alias: &str,
     injected_columns: &[String],
-    table_fields: &[String],
+    table_fields: &[DataField],
     temp_tracker: &mut TempTableTracker<'_>,
 ) -> DFResult<Vec<RecordBatch>> {
     if not_matched_batches.is_empty() || not_matched_batches.iter().all(|b| b.num_rows() == 0) {
@@ -883,7 +871,7 @@ async fn build_insert_batches_inner(
     inserts: &[MergeInsertClause],
     s_alias: &str,
     tmp_name: &str,
-    table_fields: &[String],
+    table_fields: &[DataField],
 ) -> DFResult<Vec<RecordBatch>> {
     let mut all_batches = Vec::new();
     let mut consumed_predicates: Vec<String> = Vec::new();
@@ -910,12 +898,61 @@ async fn build_insert_batches_inner(
         let sql = format!("SELECT {select_clause} FROM {tmp_name} AS {s_alias}{where_clause}");
 
         let batches = ctx.ctx().sql(&sql).await?.collect().await?;
-        all_batches.extend(batches);
+        for batch in batches {
+            all_batches.push(normalize_insert_batch_to_table_schema(
+                &batch,
+                table_fields,
+            )?);
+        }
     }
 
     Ok(all_batches)
 }
 
+fn normalize_insert_batch_to_table_schema(
+    batch: &RecordBatch,
+    table_fields: &[DataField],
+) -> DFResult<RecordBatch> {
+    if batch.num_columns() != table_fields.len() {
+        return Err(DataFusionError::Plan(format!(
+            "MERGE INSERT output has {} columns but target table has {}",
+            batch.num_columns(),
+            table_fields.len()
+        )));
+    }
+
+    let target_schema =
+        paimon::arrow::build_target_arrow_schema(table_fields).map_err(to_datafusion_error)?;
+    let mut columns = Vec::with_capacity(table_fields.len());
+
+    for (target_idx, field) in table_fields.iter().enumerate() {
+        let column = batch.column(target_idx).clone();
+        let target_type = target_schema.field(target_idx).data_type();
+        let column = cast_insert_column(field.name(), column, target_type)?;
+        columns.push(column);
+    }
+
+    RecordBatch::try_new(target_schema, columns).map_err(DataFusionError::from)
+}
+
+fn cast_insert_column(
+    name: &str,
+    column: ArrayRef,
+    target_type: &ArrowDataType,
+) -> DFResult<ArrayRef> {
+    if column.data_type() == target_type {
+        return Ok(column);
+    }
+
+    compute::cast(column.as_ref(), target_type).map_err(|e| {
+        DataFusionError::Plan(format!(
+            "Cannot cast MERGE INSERT column '{name}' from {:?} to {:?}: {e}",
+            column.data_type(),
+            target_type
+        ))
+    })
+}
+
 /// Remove injected columns from batches, keeping only source columns.
 fn strip_non_source_columns(
     batches: &[RecordBatch],
@@ -946,7 +983,7 @@ fn strip_non_source_columns(
 /// When the INSERT specifies explicit columns (`INSERT (col2, col1) VALUES (expr2, expr1)`),
 /// the output must be reordered to match the table schema so that `write_arrow_batch`
 /// (which reads columns by positional index) maps them correctly.
-fn insert_select_clause(ins: &MergeInsertClause, table_fields: &[String]) -> String {
+fn insert_select_clause(ins: &MergeInsertClause, table_fields: &[DataField]) -> String {
     if ins.columns.is_empty() && ins.value_exprs.is_empty() {
         "*".to_string()
     } else {
@@ -962,11 +999,11 @@ fn insert_select_clause(ins: &MergeInsertClause, table_fields: &[String]) -> Str
         table_fields
             .iter()
             .map(|field| {
-                let key = field.to_lowercase();
+                let key = field.name().to_lowercase();
                 match col_expr_map.get(&key) {
-                    Some(expr) => format!("{expr} AS {}", quote_identifier(field)),
+                    Some(expr) => format!("{expr} AS {}", quote_identifier(field.name())),
                     // Column not in INSERT list — fill with NULL
-                    None => format!("NULL AS {}", quote_identifier(field)),
+                    None => format!("NULL AS {}", quote_identifier(field.name())),
                 }
             })
             .collect::<Vec<_>>()
@@ -1546,7 +1583,7 @@ mod tests {
     use datafusion::sql::sqlparser::parser::Parser;
     use paimon::catalog::{Catalog, Identifier};
     use paimon::io::FileIOBuilder;
-    use paimon::spec::{DataType, IntType, Schema as PaimonSchema, TableSchema};
+    use paimon::spec::{DataField, DataType, IntType, Schema as PaimonSchema, TableSchema};
     use paimon::{CatalogOptions, FileSystemCatalog, Options};
     use tempfile::TempDir;
 
@@ -1613,6 +1650,42 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_normalize_merge_insert_batch_uses_position() {
+        let table_fields = vec![
+            DataField::new(0, "a".to_string(), DataType::Int(IntType::new())),
+            DataField::new(1, "b".to_string(), DataType::Int(IntType::new())),
+        ];
+        let batch = RecordBatch::try_new(
+            Arc::new(Schema::new(vec![
+                Field::new("b", ArrowDataType::Int32, false),
+                Field::new("x", ArrowDataType::Int32, false),
+            ])),
+            vec![
+                Arc::new(Int32Array::from(vec![100])),
+                Arc::new(Int32Array::from(vec![7])),
+            ],
+        )
+        .unwrap();
+
+        let normalized = normalize_insert_batch_to_table_schema(&batch, &table_fields).unwrap();
+        let first = normalized
+            .column(0)
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap();
+        let second = normalized
+            .column(1)
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap();
+
+        assert_eq!(normalized.schema().field(0).name(), "a");
+        assert_eq!(normalized.schema().field(1).name(), "b");
+        assert_eq!(first.value(0), 100);
+        assert_eq!(second.value(0), 7);
+    }
+
     #[test]
     fn test_source_partition_pruning_requires_partition_equality() {
         let merge = parse_merge(

crates/paimon/src/arrow/format/blob.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use super::{FilePredicates, FormatFileReader, FormatFileWriter};
+use super::{FilePredicates, FormatFileReader, FormatFileWriter, FormatWriteResult};
 use crate::arrow::build_target_arrow_schema;
 use crate::io::{FileRead, FileWrite};
 use crate::spec::{BlobDescriptor, DataField, DataType};
@@ -725,7 +725,7 @@ impl FormatFileWriter for BlobFormatWriter {
         Ok(())
     }
 
-    async fn close(mut self: Box<Self>) -> crate::Result<u64> {
+    async fn close(mut self: Box<Self>) -> crate::Result<FormatWriteResult> {
         let index_bytes = encode_delta_varints_write(&self.lengths);
         let index_length = index_bytes.len() as i32;
 
@@ -739,7 +739,7 @@ impl FormatFileWriter for BlobFormatWriter {
 
         let total = self.bytes_written + index_length as u64 + BLOB_FOOTER_SIZE;
         self.writer.close().await?;
-        Ok(total)
+        Ok(FormatWriteResult::new(total))
     }
 }
 

crates/paimon/src/arrow/format/mod.rs

@@ -19,7 +19,7 @@ mod avro;
 pub(crate) mod blob;
 mod mosaic;
 mod orc;
-mod parquet;
+pub(crate) mod parquet;
 mod row;
 mod shredding;
 #[cfg(feature = "vortex")]
@@ -29,6 +29,7 @@ mod vortex;
 pub(crate) use parquet::ParquetFormatWriter;
 
 use crate::io::{FileRead, OutputFile};
+use crate::spec::stats::BinaryTableStats;
 use crate::spec::{DataField, Predicate};
 use crate::table::{ArrowRecordBatchStream, RowRange};
 use crate::Error;
@@ -102,8 +103,40 @@ pub(crate) trait FormatFileWriter: Send {
     async fn flush(&mut self) -> crate::Result<()>;
 
     /// Flush and close the writer, finalizing the file on storage.
-    /// Returns the total number of bytes written.
-    async fn close(self: Box<Self>) -> crate::Result<u64>;
+    async fn close(self: Box<Self>) -> crate::Result<FormatWriteResult>;
+}
+
+pub(crate) struct FormatWriteResult {
+    pub(crate) file_size: u64,
+    pub(crate) value_stats: Option<FormatValueStats>,
+}
+
+pub(crate) struct FormatValueStats {
+    pub(crate) stats: BinaryTableStats,
+    pub(crate) columns: Option<Vec<String>>,
+}
+
+impl FormatWriteResult {
+    pub(crate) fn new(file_size: u64) -> Self {
+        Self {
+            file_size,
+            value_stats: None,
+        }
+    }
+
+    pub(crate) fn with_value_stats(
+        file_size: u64,
+        value_stats: BinaryTableStats,
+        columns: Option<Vec<String>>,
+    ) -> Self {
+        Self {
+            file_size,
+            value_stats: Some(FormatValueStats {
+                stats: value_stats,
+                columns,
+            }),
+        }
+    }
 }
 
 /// Create a format reader based on the file extension.
@@ -180,6 +213,7 @@ pub(crate) async fn create_format_writer(
             output,
             compression,
             zstd_level,
+            format_options.cloned().unwrap_or_default(),
         ));
         shredding::ShreddingFormatWriter::create(
             writer_factory,