nullcount

DrakeLin · DrakeLin · commit d72d2fc9ff11 · 2026-01-24T04:32:00.000Z
diff --git a/kernel/src/engine/default/stats.rs b/kernel/src/engine/default/stats.rs
@@ -2,32 +2,110 @@
 //!
 //! Provides `collect_stats` to compute statistics for a RecordBatch during file writes.
 
+use std::collections::HashSet;
 use std::sync::Arc;
 
-use crate::arrow::array::{Array, BooleanArray, Int64Array, RecordBatch, StructArray};
-use crate::arrow::datatypes::{DataType, Field};
+use crate::arrow::array::{Array, ArrayRef, BooleanArray, Int64Array, RecordBatch, StructArray};
+use crate::arrow::datatypes::{DataType, Field, Fields};
 use crate::{DeltaResult, Error};
 
+/// Compute null count for a column, handling nested structs.
+fn compute_null_count(column: &ArrayRef) -> DeltaResult<ArrayRef> {
+    match column.data_type() {
+        DataType::Struct(fields) => {
+            let struct_array = column
+                .as_any()
+                .downcast_ref::<StructArray>()
+                .ok_or_else(|| Error::generic("Expected StructArray for struct column"))?;
+            let children: Vec<ArrayRef> = (0..fields.len())
+                .map(|i| compute_null_count(struct_array.column(i)))
+                .collect::<DeltaResult<Vec<_>>>()?;
+            let null_count_fields: Fields = fields
+                .iter()
+                .map(|f| Field::new(f.name(), null_count_data_type(f.data_type()), true))
+                .collect();
+            Ok(Arc::new(
+                StructArray::try_new(null_count_fields, children, None)
+                    .map_err(|e| Error::generic(format!("null count struct: {e}")))?,
+            ))
+        }
+        _ => Ok(Arc::new(Int64Array::from(vec![column.null_count() as i64]))),
+    }
+}
+
+/// Get the data type for null counts of a given data type.
+fn null_count_data_type(data_type: &DataType) -> DataType {
+    match data_type {
+        DataType::Struct(fields) => {
+            let null_count_fields: Vec<Field> = fields
+                .iter()
+                .map(|f| Field::new(f.name(), null_count_data_type(f.data_type()), true))
+                .collect();
+            DataType::Struct(null_count_fields.into())
+        }
+        _ => DataType::Int64,
+    }
+}
+
 /// Collect statistics from a RecordBatch for Delta Lake file statistics.
 ///
 /// Returns a StructArray with the following fields:
 /// - `numRecords`: total row count
+/// - `nullCount`: nested struct with null counts per column (if stats_columns is non-empty)
 /// - `tightBounds`: always true for new file writes
 ///
 /// # Arguments
 /// * `batch` - The RecordBatch to collect statistics from
-/// * `_stats_columns` - Column names that should have statistics collected (reserved for future use)
+/// * `stats_columns` - Column names that should have statistics collected.
+///   Only these columns will appear in nullCount.
 pub(crate) fn collect_stats(
     batch: &RecordBatch,
-    _stats_columns: &[String],
+    stats_columns: &[String],
 ) -> DeltaResult<StructArray> {
+    let stats_set: HashSet<&str> = stats_columns.iter().map(|s| s.as_str()).collect();
+    let schema = batch.schema();
+
     let mut fields = Vec::new();
     let mut arrays: Vec<Arc<dyn Array>> = Vec::new();
 
     // numRecords
     fields.push(Field::new("numRecords", DataType::Int64, true));
     arrays.push(Arc::new(Int64Array::from(vec![batch.num_rows() as i64])));
 
+    // Collect null counts for each column that's in stats_columns
+    let mut null_count_fields = Vec::new();
+    let mut null_count_arrays = Vec::new();
+
+    for (col_idx, field) in schema.fields().iter().enumerate() {
+        if !stats_set.contains(field.name().as_str()) {
+            continue;
+        }
+
+        let column = batch.column(col_idx);
+
+        // Null count
+        let null_count = compute_null_count(column)?;
+        null_count_fields.push(Field::new(
+            field.name(),
+            null_count.data_type().clone(),
+            true,
+        ));
+        null_count_arrays.push(null_count);
+    }
+
+    // nullCount struct
+    if !null_count_fields.is_empty() {
+        let null_count_struct =
+            StructArray::try_new(null_count_fields.into(), null_count_arrays, None)
+                .map_err(|e| Error::generic(format!("Failed to create nullCount: {e}")))?;
+        fields.push(Field::new(
+            "nullCount",
+            null_count_struct.data_type().clone(),
+            true,
+        ));
+        arrays.push(Arc::new(null_count_struct));
+    }
+
     // tightBounds - always true for new file writes
     fields.push(Field::new("tightBounds", DataType::Boolean, true));
     arrays.push(Arc::new(BooleanArray::from(vec![true])));
@@ -39,7 +117,7 @@ pub(crate) fn collect_stats(
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::arrow::array::Int64Array;
+    use crate::arrow::array::{Int64Array, StringArray};
     use crate::arrow::datatypes::Schema;
 
     #[test]
@@ -89,4 +167,98 @@ mod tests {
             .unwrap();
         assert_eq!(num_records.value(0), 0);
     }
+
+    #[test]
+    fn test_collect_stats_null_counts() {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Int64, false),
+            Field::new("value", DataType::Utf8, true),
+        ]));
+
+        let batch = RecordBatch::try_new(
+            schema,
+            vec![
+                Arc::new(Int64Array::from(vec![1, 2, 3])),
+                Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])),
+            ],
+        )
+        .unwrap();
+
+        let stats = collect_stats(&batch, &["id".to_string(), "value".to_string()]).unwrap();
+
+        // Check nullCount struct
+        let null_count = stats
+            .column_by_name("nullCount")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<StructArray>()
+            .unwrap();
+
+        // id has 0 nulls
+        let id_null_count = null_count
+            .column_by_name("id")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<Int64Array>()
+            .unwrap();
+        assert_eq!(id_null_count.value(0), 0);
+
+        // value has 1 null
+        let value_null_count = null_count
+            .column_by_name("value")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<Int64Array>()
+            .unwrap();
+        assert_eq!(value_null_count.value(0), 1);
+    }
+
+    #[test]
+    fn test_collect_stats_respects_stats_columns() {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Int64, false),
+            Field::new("value", DataType::Utf8, true),
+        ]));
+
+        let batch = RecordBatch::try_new(
+            schema,
+            vec![
+                Arc::new(Int64Array::from(vec![1, 2, 3])),
+                Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])),
+            ],
+        )
+        .unwrap();
+
+        // Only collect stats for "id", not "value"
+        let stats = collect_stats(&batch, &["id".to_string()]).unwrap();
+
+        let null_count = stats
+            .column_by_name("nullCount")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<StructArray>()
+            .unwrap();
+
+        // Only id should be present
+        assert!(null_count.column_by_name("id").is_some());
+        assert!(null_count.column_by_name("value").is_none());
+    }
+
+    #[test]
+    fn test_collect_stats_empty_stats_columns() {
+        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
+
+        let batch =
+            RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1, 2, 3]))]).unwrap();
+
+        // No stats columns requested
+        let stats = collect_stats(&batch, &[]).unwrap();
+
+        // Should still have numRecords and tightBounds
+        assert!(stats.column_by_name("numRecords").is_some());
+        assert!(stats.column_by_name("tightBounds").is_some());
+
+        // Should not have nullCount
+        assert!(stats.column_by_name("nullCount").is_none());
+    }
 }
diff --git a/kernel/tests/write.rs b/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,\"tightBounds\":true}"
+                    "stats": "{\"numRecords\":3,\"nullCount\":{\"number\":0},\"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,\"tightBounds\":true}"
+                    "stats": "{\"numRecords\":3,\"nullCount\":{\"number\":0},\"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,\"tightBounds\":true}"
+                    "stats": "{\"numRecords\":3,\"nullCount\":{\"number\":0},\"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,\"tightBounds\":true}"
+                    "stats": "{\"numRecords\":3,\"nullCount\":{\"number\":0},\"tightBounds\":true}"
                 }
             }),
         ];
