propagate

Author: DrakeLin

kernel/src/engine/arrow_expression/apply_schema.rs

@@ -5,8 +5,9 @@ use std::sync::Arc;
 use itertools::Itertools;
 
 use crate::arrow::array::{
-    Array, ArrayRef, AsArray, ListArray, MapArray, RecordBatch, StructArray,
+    make_array, Array, ArrayRef, AsArray, ListArray, MapArray, RecordBatch, StructArray,
 };
+use crate::arrow::buffer::NullBuffer;
 use crate::arrow::datatypes::Schema as ArrowSchema;
 use crate::arrow::datatypes::{DataType as ArrowDataType, Field as ArrowField};
 
@@ -18,6 +19,10 @@ use crate::schema::{ArrayType, DataType, MapType, Schema, StructField};
 // Apply a schema to an array. The array _must_ be a `StructArray`. Returns a `RecordBatch where the
 // names of fields, nullable, and metadata in the struct have been transformed to match those in
 // schema specified by `schema`
+//
+// If the struct array has top-level nulls (rows where the entire struct is null), those nulls are
+// propagated to all child columns. This allows expressions like `add.stats_parsed` to return null
+// for rows where `add` is null, rather than erroring.
 pub(crate) fn apply_schema(array: &dyn Array, schema: &DataType) -> DeltaResult<RecordBatch> {
     let DataType::Struct(struct_schema) = schema else {
         return Err(Error::generic(
@@ -26,19 +31,40 @@ pub(crate) fn apply_schema(array: &dyn Array, schema: &DataType) -> DeltaResult<
     };
     let applied = apply_schema_to_struct(array, struct_schema)?;
     let (fields, columns, nulls) = applied.into_parts();
-    if let Some(nulls) = nulls {
-        if nulls.null_count() != 0 {
-            return Err(Error::invalid_struct_data(
-                "Top-level nulls in struct are not supported",
-            ));
+
+    // If there are top-level nulls, propagate them to each child column.
+    // This handles cases where we're extracting a struct from a nullable parent
+    // (e.g., `add.stats_parsed` where `add` can be null for non-add rows).
+    let columns = if let Some(ref struct_nulls) = nulls {
+        if struct_nulls.null_count() > 0 {
+            columns
+                .into_iter()
+                .map(|column| propagate_nulls_to_column(&column, struct_nulls))
+                .collect()
+        } else {
+            columns
         }
-    }
+    } else {
+        columns
+    };
+
     Ok(RecordBatch::try_new(
         Arc::new(ArrowSchema::new(fields)),
         columns,
     )?)
 }
 
+/// Propagate a parent null buffer to a column, combining it with the column's existing nulls.
+fn propagate_nulls_to_column(column: &ArrayRef, parent_nulls: &NullBuffer) -> ArrayRef {
+    let data = column.to_data();
+    let combined_nulls = NullBuffer::union(Some(parent_nulls), data.nulls());
+    let builder = data.into_builder().nulls(combined_nulls);
+    // SAFETY: We're only adding more nulls to an existing valid array.
+    // The union can only grow the set of NULL rows, preserving data validity.
+    let data = unsafe { builder.build_unchecked() };
+    make_array(data)
+}
+
 // helper to transform an arrow field+col into the specified target type. If `rename` is specified
 // the field will be renamed to the contained `str`.
 fn new_field_with_metadata(
@@ -191,6 +217,7 @@ mod apply_schema_validation_tests {
     use std::sync::Arc;
 
     use crate::arrow::array::{Int32Array, StructArray};
+    use crate::arrow::buffer::{BooleanBuffer, NullBuffer};
     use crate::arrow::datatypes::{
         DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema,
     };
@@ -238,4 +265,75 @@ mod apply_schema_validation_tests {
             StructField::new("b", DataType::INTEGER, false),
         ])
     }
+
+    /// Test that top-level struct nulls are propagated to child columns.
+    ///
+    /// This simulates a Delta log scenario where each row is one action type (add, remove, etc.).
+    /// When extracting `add.stats_parsed`, rows where `add` is null (e.g., remove actions) should
+    /// return null for all child columns rather than erroring.
+    #[test]
+    fn test_apply_schema_propagates_top_level_nulls() {
+        // Create a struct array with 4 rows where rows 1 and 3 have top-level nulls.
+        // This simulates: [add_action, remove_action, add_action, remove_action]
+        // where remove_action rows have null for the entire struct.
+        let field_a = ArrowField::new("a", ArrowDataType::Int32, true);
+        let field_b = ArrowField::new("b", ArrowDataType::Int32, true);
+        let schema = ArrowSchema::new(vec![field_a, field_b]);
+
+        // Child column data - note these have values even for rows that will be struct-null
+        // because Arrow stores child data independently of the struct's null buffer.
+        let a_data = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4)]);
+        let b_data = Int32Array::from(vec![Some(10), Some(20), Some(30), Some(40)]);
+
+        // Top-level nulls: rows 0 and 2 are valid, rows 1 and 3 are null
+        let null_buffer = NullBuffer::new(BooleanBuffer::from(vec![true, false, true, false]));
+
+        let struct_array = StructArray::try_new(
+            schema.fields.clone(),
+            vec![Arc::new(a_data), Arc::new(b_data)],
+            Some(null_buffer),
+        )
+        .unwrap();
+
+        // Target schema with nullable fields
+        let target_schema = DataType::Struct(Box::new(StructType::new_unchecked([
+            StructField::new("a", DataType::INTEGER, true),
+            StructField::new("b", DataType::INTEGER, true),
+        ])));
+
+        // Apply schema - this should propagate top-level nulls to child columns
+        let result = apply_schema(&struct_array, &target_schema).unwrap();
+
+        assert_eq!(result.num_rows(), 4);
+        assert_eq!(result.num_columns(), 2);
+
+        // Verify column "a" has nulls propagated from the struct's null buffer
+        let col_a = result.column(0);
+        assert!(col_a.is_valid(0), "Row 0 should be valid");
+        assert!(col_a.is_null(1), "Row 1 should be null (struct was null)");
+        assert!(col_a.is_valid(2), "Row 2 should be valid");
+        assert!(col_a.is_null(3), "Row 3 should be null (struct was null)");
+
+        // Verify column "b" has nulls propagated from the struct's null buffer
+        let col_b = result.column(1);
+        assert!(col_b.is_valid(0), "Row 0 should be valid");
+        assert!(col_b.is_null(1), "Row 1 should be null (struct was null)");
+        assert!(col_b.is_valid(2), "Row 2 should be valid");
+        assert!(col_b.is_null(3), "Row 3 should be null (struct was null)");
+
+        // Verify the actual values for valid rows
+        let col_a = col_a
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .expect("column a should be Int32Array");
+        let col_b = col_b
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .expect("column b should be Int32Array");
+
+        assert_eq!(col_a.value(0), 1);
+        assert_eq!(col_a.value(2), 3);
+        assert_eq!(col_b.value(0), 10);
+        assert_eq!(col_b.value(2), 30);
+    }
 }

kernel/src/engine/arrow_expression/tests.rs

@@ -827,6 +827,8 @@ fn test_create_one_mismatching_scalar_types() {
 
 #[test]
 fn test_create_one_not_null_struct() {
+    // Creating a NOT NULL struct field with null values should error.
+    // The error comes from Arrow's RecordBatch validation (non-nullable column has nulls).
     let values: &[Scalar] = &[
         Scalar::Null(KernelDataType::INTEGER),
         Scalar::Null(KernelDataType::INTEGER),
@@ -841,23 +843,25 @@ fn test_create_one_not_null_struct() {
     let handler = ArrowEvaluationHandler;
     assert_result_error_with_message(
         handler.create_one(schema, values),
-        "Invalid struct data: Top-level nulls in struct are not supported",
+        "Column 'a' is declared as non-nullable but contains null values",
     );
 }
 
 #[test]
 fn test_create_one_top_level_null() {
+    // Creating a NOT NULL field with null value should error.
+    // The error comes from Arrow's RecordBatch validation.
     let values = &[Scalar::Null(KernelDataType::INTEGER)];
     let handler = ArrowEvaluationHandler;
 
     let schema = Arc::new(StructType::new_unchecked([StructField::not_null(
         "col_1",
         KernelDataType::INTEGER,
     )]));
-    assert!(matches!(
+    assert_result_error_with_message(
         handler.create_one(schema, values),
-        Err(Error::InvalidStructData(_))
-    ));
+        "Column 'col_1' is declared as non-nullable but contains null values",
+    );
 }
 
 #[test]