feat: add stats_columns to ParquetHandler + Transaction stats API

- Add stats_columns parameter to write_parquet_file trait
- Add stats_schema(), stats_columns(), get_clustering_columns() to Transaction
- Add stats_columns to WriteContext
- Update get_write_context() to take engine parameter
- Add clustering column support to expected_stats_schema()

Author: DrakeLin

ffi/src/transaction/mod.rs

@@ -288,7 +288,9 @@ mod tests {
                 ))
             };
 
-            let write_context = unsafe { get_write_context(txn_with_engine_info.shallow_copy()) };
+            let write_context = ok_or_panic(unsafe {
+                get_write_context(txn_with_engine_info.shallow_copy(), engine.shallow_copy())
+            });
 
             // Ensure we get the correct schema
             let write_schema = unsafe { get_write_schema(write_context.shallow_copy()) };

ffi/src/transaction/write_context.rs

@@ -1,5 +1,9 @@
+use crate::error::{ExternResult, IntoExternResult};
 use crate::handle::Handle;
-use crate::{kernel_string_slice, AllocateStringFn, NullableCvoid, SharedSchema};
+use crate::{
+    kernel_string_slice, AllocateStringFn, DeltaResult, ExternEngine, NullableCvoid,
+    SharedExternEngine, SharedSchema,
+};
 use delta_kernel::transaction::WriteContext;
 use delta_kernel_ffi_macros::handle_descriptor;
 
@@ -19,13 +23,24 @@ pub struct SharedWriteContext;
 ///
 /// # Safety
 ///
-/// Caller is responsible for passing a [valid][Handle#Validity] transaction handle.
+/// Caller is responsible for passing a [valid][Handle#Validity] transaction handle and engine handle.
 #[no_mangle]
 pub unsafe extern "C" fn get_write_context(
     txn: Handle<ExclusiveTransaction>,
-) -> Handle<SharedWriteContext> {
+    engine: Handle<SharedExternEngine>,
+) -> ExternResult<Handle<SharedWriteContext>> {
     let txn = unsafe { txn.as_ref() };
-    Arc::new(txn.get_write_context()).into()
+    let extern_engine = unsafe { engine.as_ref() };
+    get_write_context_impl(txn, extern_engine).into_extern_result(&extern_engine)
+}
+
+fn get_write_context_impl(
+    txn: &delta_kernel::transaction::Transaction,
+    extern_engine: &dyn ExternEngine,
+) -> DeltaResult<Handle<SharedWriteContext>> {
+    let engine = extern_engine.engine();
+    let write_context = txn.get_write_context(engine.as_ref())?;
+    Ok(Arc::new(write_context).into())
 }
 
 #[no_mangle]

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

@@ -94,7 +94,7 @@ 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 = Arc::new(txn.get_write_context(&engine)?);
     let file_metadata = engine
         .write_parquet(&sample_data, write_context.as_ref(), HashMap::new())
         .await?;

kernel/src/scan/data_skipping/stats_schema.rs

@@ -93,18 +93,29 @@ use crate::{
 ///   tightBounds: boolean,
 /// }
 /// ```
+/// Generates the expected schema for file statistics.
+///
+/// # Arguments
+///
+/// * `physical_file_schema` - The physical schema of the data files (no partition columns)
+/// * `table_properties` - Table properties containing stats configuration
+/// * `clustering_columns` - List of clustering columns that must always be included
+///   in statistics regardless of `dataSkippingNumIndexedCols` limits. Pass an empty slice
+///   if there are no clustering columns.
 #[allow(unused)]
 pub(crate) fn expected_stats_schema(
     physical_file_schema: &Schema,
     table_properties: &TableProperties,
+    clustering_columns: &[ColumnName],
 ) -> DeltaResult<Schema> {
     let mut fields = Vec::with_capacity(5);
     fields.push(StructField::nullable("numRecords", DataType::LONG));
 
     // generate the base stats schema:
     // - make all fields nullable
     // - include fields according to table properties (num_indexed_cols, stats_columns, ...)
-    let mut base_transform = BaseStatsTransform::new(table_properties);
+    // - always include clustering columns (they don't count against the column limit)
+    let mut base_transform = BaseStatsTransform::new(table_properties, clustering_columns);
     if let Some(base_schema) = base_transform.transform_struct(physical_file_schema) {
         let base_schema = base_schema.into_owned();
 
@@ -184,26 +195,30 @@ impl<'a> SchemaTransform<'a> for NullCountStatsTransform {
 ///   to be used for data skipping statistics. (takes precedence)
 /// * `dataSkippingNumIndexedCols` - used to specify the number of columns
 ///   to be used for data skipping statistics. Defaults to 32.
+/// * Clustering columns are always included regardless of the above limits.
 ///
 /// All fields are nullable.
 #[allow(unused)]
-struct BaseStatsTransform {
+struct BaseStatsTransform<'a> {
     n_columns: Option<DataSkippingNumIndexedCols>,
     added_columns: u64,
     column_names: Option<Vec<ColumnName>>,
+    /// Clustering columns that must always be included (don't count against limit)
+    clustering_columns: &'a [ColumnName],
     path: Vec<String>,
 }
 
-impl BaseStatsTransform {
+impl<'a> BaseStatsTransform<'a> {
     #[allow(unused)]
-    fn new(props: &TableProperties) -> Self {
+    fn new(props: &TableProperties, clustering_columns: &'a [ColumnName]) -> Self {
         // If data_skipping_stats_columns is specified, it takes precedence
         // over data_skipping_num_indexed_cols, even if that is also specified.
         if let Some(column_names) = &props.data_skipping_stats_columns {
             Self {
                 n_columns: None,
                 added_columns: 0,
                 column_names: Some(column_names.clone()),
+                clustering_columns,
                 path: Vec::new(),
             }
         } else {
@@ -214,44 +229,63 @@ impl BaseStatsTransform {
                 n_columns: Some(n_cols),
                 added_columns: 0,
                 column_names: None,
+                clustering_columns,
                 path: Vec::new(),
             }
         }
     }
+
+    /// Checks if the current path matches a clustering column.
+    /// Clustering columns are always included regardless of column limits.
+    fn is_clustering_column(&self) -> bool {
+        let current = ColumnName::new(&self.path);
+        self.clustering_columns.iter().any(|c| c == &current)
+    }
 }
 
-impl<'a> SchemaTransform<'a> for BaseStatsTransform {
+impl<'a, 'b> SchemaTransform<'a> for BaseStatsTransform<'b> {
     fn transform_struct_field(&mut self, field: &'a StructField) -> Option<Cow<'a, StructField>> {
         use Cow::*;
 
+        self.path.push(field.name.clone());
+        let data_type = field.data_type();
+
+        // Clustering columns are always included, regardless of limits
+        let is_clustering = self.is_clustering_column();
+
         // Check if the number of columns is set and if the added columns exceed the limit
-        // In the constructor we assert this will always be None if column_names are specified
-        if let Some(DataSkippingNumIndexedCols::NumColumns(n_cols)) = self.n_columns {
-            if self.added_columns >= n_cols {
-                return None;
+        // Clustering columns bypass this limit
+        if !is_clustering {
+            if let Some(DataSkippingNumIndexedCols::NumColumns(n_cols)) = self.n_columns {
+                if self.added_columns >= n_cols {
+                    self.path.pop();
+                    return None;
+                }
             }
         }
 
-        self.path.push(field.name.clone());
-        let data_type = field.data_type();
-
         // We always traverse struct fields (they don't count against the column limit),
         // but we only include leaf fields if they qualify based on column_names config.
         // When column_names is None, all leaf fields are included (up to n_columns limit).
+        // Clustering columns are always included regardless of column_names config.
         if !matches!(data_type, DataType::Struct(_)) {
-            let should_include = self
-                .column_names
-                .as_ref()
-                .map(|ns| should_include_column(&ColumnName::new(&self.path), ns))
-                .unwrap_or(true);
+            let should_include = is_clustering
+                || self
+                    .column_names
+                    .as_ref()
+                    .map(|ns| should_include_column(&ColumnName::new(&self.path), ns))
+                    .unwrap_or(true);
 
             if !should_include {
                 self.path.pop();
                 return None;
             }
 
-            // Increment count only for leaf columns
-            self.added_columns += 1;
+            // Increment count only for non-clustering leaf columns
+            // Clustering columns don't count against the limit
+            if !is_clustering {
+                self.added_columns += 1;
+            }
         }
 
         let field = match self.transform(&field.data_type)? {
@@ -363,7 +397,7 @@ mod tests {
         let properties: TableProperties = [("key", "value")].into();
         let file_schema = StructType::new_unchecked([StructField::nullable("id", DataType::LONG)]);
 
-        let stats_schema = expected_stats_schema(&file_schema, &properties).unwrap();
+        let stats_schema = expected_stats_schema(&file_schema, &properties, &[]).unwrap();
         let expected = StructType::new_unchecked([
             StructField::nullable("numRecords", DataType::LONG),
             StructField::nullable("nullCount", file_schema.clone()),
@@ -387,7 +421,7 @@ mod tests {
             StructField::not_null("id", DataType::LONG),
             StructField::not_null("user", DataType::Struct(Box::new(user_struct.clone()))),
         ]);
-        let stats_schema = expected_stats_schema(&file_schema, &properties).unwrap();
+        let stats_schema = expected_stats_schema(&file_schema, &properties, &[]).unwrap();
 
         // Expected result: The stats schema should maintain the nested structure
         // but make all fields nullable
@@ -437,7 +471,7 @@ mod tests {
             ),
         ]);
 
-        let stats_schema = expected_stats_schema(&file_schema, &properties).unwrap();
+        let stats_schema = expected_stats_schema(&file_schema, &properties, &[]).unwrap();
 
         let expected_null_nested = StructType::new_unchecked([
             StructField::nullable("name", DataType::LONG),
@@ -486,7 +520,7 @@ mod tests {
             StructField::nullable("user.info", DataType::Struct(Box::new(user_struct.clone()))),
         ]);
 
-        let stats_schema = expected_stats_schema(&file_schema, &properties).unwrap();
+        let stats_schema = expected_stats_schema(&file_schema, &properties, &[]).unwrap();
 
         let expected_nested =
             StructType::new_unchecked([StructField::nullable("name", DataType::STRING)]);
@@ -523,7 +557,7 @@ mod tests {
             StructField::nullable("age", DataType::INTEGER),
         ]);
 
-        let stats_schema = expected_stats_schema(&logical_schema, &properties).unwrap();
+        let stats_schema = expected_stats_schema(&logical_schema, &properties, &[]).unwrap();
 
         let expected_fields =
             StructType::new_unchecked([StructField::nullable("name", DataType::STRING)]);
@@ -557,7 +591,7 @@ mod tests {
             StructField::nullable("metadata", DataType::BINARY),
         ]);
 
-        let stats_schema = expected_stats_schema(&file_schema, &properties).unwrap();
+        let stats_schema = expected_stats_schema(&file_schema, &properties, &[]).unwrap();
 
         // Expected nullCount schema: all fields converted to LONG
         let expected_null_count = StructType::new_unchecked([
@@ -599,7 +633,7 @@ mod tests {
             StructField::nullable("is_deleted", DataType::BOOLEAN), // NOT eligible for min/max
         ]);
 
-        let stats_schema = expected_stats_schema(&file_schema, &properties).unwrap();
+        let stats_schema = expected_stats_schema(&file_schema, &properties, &[]).unwrap();
 
         // Expected nullCount schema: all fields converted to LONG, maintaining structure
         let expected_null_user = StructType::new_unchecked([
@@ -649,7 +683,7 @@ mod tests {
             ),
         ]);
 
-        let stats_schema = expected_stats_schema(&file_schema, &properties).unwrap();
+        let stats_schema = expected_stats_schema(&file_schema, &properties, &[]).unwrap();
 
         // Expected nullCount schema: all fields converted to LONG
         let expected_null_count = StructType::new_unchecked([
@@ -669,4 +703,103 @@ mod tests {
 
         assert_eq!(&expected, &stats_schema);
     }
+
+    #[test]
+    fn test_stats_schema_clustering_columns_bypass_limit() {
+        // Set limit to 1 column, but clustering columns should still be included
+        let properties: TableProperties = [(
+            "delta.dataSkippingNumIndexedCols".to_string(),
+            "1".to_string(),
+        )]
+        .into();
+
+        let file_schema = StructType::new_unchecked([
+            StructField::nullable("id", DataType::LONG),
+            StructField::nullable("name", DataType::STRING),
+            StructField::nullable("cluster_col", DataType::INTEGER),
+        ]);
+
+        // Without clustering columns: only first column (id) should be included
+        let stats_schema_no_cluster =
+            expected_stats_schema(&file_schema, &properties, &[]).unwrap();
+
+        let expected_no_cluster =
+            StructType::new_unchecked([StructField::nullable("id", DataType::LONG)]);
+        let null_count_no_cluster = NullCountStatsTransform
+            .transform_struct(&expected_no_cluster)
+            .unwrap()
+            .into_owned();
+
+        let expected_no_cluster = StructType::new_unchecked([
+            StructField::nullable("numRecords", DataType::LONG),
+            StructField::nullable("nullCount", null_count_no_cluster),
+            StructField::nullable("minValues", expected_no_cluster.clone()),
+            StructField::nullable("maxValues", expected_no_cluster),
+            StructField::nullable("tightBounds", DataType::BOOLEAN),
+        ]);
+        assert_eq!(&expected_no_cluster, &stats_schema_no_cluster);
+
+        // With clustering columns: id (from limit) + cluster_col (clustering) should be included
+        let clustering_cols = vec![ColumnName::new(["cluster_col"])];
+        let stats_schema_with_cluster =
+            expected_stats_schema(&file_schema, &properties, &clustering_cols).unwrap();
+
+        let expected_with_cluster = StructType::new_unchecked([
+            StructField::nullable("id", DataType::LONG),
+            StructField::nullable("cluster_col", DataType::INTEGER),
+        ]);
+        let null_count_with_cluster = NullCountStatsTransform
+            .transform_struct(&expected_with_cluster)
+            .unwrap()
+            .into_owned();
+
+        let expected_with_cluster = StructType::new_unchecked([
+            StructField::nullable("numRecords", DataType::LONG),
+            StructField::nullable("nullCount", null_count_with_cluster),
+            StructField::nullable("minValues", expected_with_cluster.clone()),
+            StructField::nullable("maxValues", expected_with_cluster),
+            StructField::nullable("tightBounds", DataType::BOOLEAN),
+        ]);
+        assert_eq!(&expected_with_cluster, &stats_schema_with_cluster);
+    }
+
+    #[test]
+    fn test_stats_schema_clustering_columns_with_stats_columns() {
+        // When dataSkippingStatsColumns is set, clustering columns should still be included
+        let properties: TableProperties = [(
+            "delta.dataSkippingStatsColumns".to_string(),
+            "id".to_string(),
+        )]
+        .into();
+
+        let file_schema = StructType::new_unchecked([
+            StructField::nullable("id", DataType::LONG),
+            StructField::nullable("name", DataType::STRING),
+            StructField::nullable("cluster_col", DataType::INTEGER),
+        ]);
+
+        // With dataSkippingStatsColumns=id and clustering_columns=[cluster_col],
+        // both should be included
+        let clustering_cols = vec![ColumnName::new(["cluster_col"])];
+        let stats_schema =
+            expected_stats_schema(&file_schema, &properties, &clustering_cols).unwrap();
+
+        let expected_fields = StructType::new_unchecked([
+            StructField::nullable("id", DataType::LONG),
+            StructField::nullable("cluster_col", DataType::INTEGER),
+        ]);
+        let null_count = NullCountStatsTransform
+            .transform_struct(&expected_fields)
+            .unwrap()
+            .into_owned();
+
+        let expected = StructType::new_unchecked([
+            StructField::nullable("numRecords", DataType::LONG),
+            StructField::nullable("nullCount", null_count),
+            StructField::nullable("minValues", expected_fields.clone()),
+            StructField::nullable("maxValues", expected_fields),
+            StructField::nullable("tightBounds", DataType::BOOLEAN),
+        ]);
+        assert_eq!(&expected, &stats_schema);
+    }
 }