fix: correct GetStructField null handling for null parent structs (value + nullability) (#4523)

Author: schenksj

native/spark-expr/src/struct_funcs/get_struct_field.rs

@@ -15,7 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{Array, StructArray};
+use arrow::array::{make_array, Array, ArrayRef, StructArray};
+use arrow::buffer::NullBuffer;
 use arrow::datatypes::{DataType, Field, Schema};
 use arrow::record_batch::RecordBatch;
 use datafusion::common::{DataFusionError, Result as DataFusionResult, ScalarValue};
@@ -59,6 +60,27 @@ impl GetStructField {
             ))),
         }
     }
+
+    /// Extract field `ordinal` from a struct array, propagating the parent struct's null mask.
+    ///
+    /// Spark semantics: a field of a NULL struct is NULL. Arrow stores a StructArray's child
+    /// arrays with their own validity, INDEPENDENT of the parent struct's null buffer -- so the
+    /// raw child value at a row where the struct itself is null can be non-null (e.g. parquet
+    /// files where a logically-null struct column still has a populated child buffer). Returning
+    /// the child column verbatim then makes `isnotnull(struct.field)` wrongly true for a null
+    /// struct. Union the struct's null mask into the child's (null where the struct is null OR
+    /// the child is null).
+    fn project_field(struct_array: &StructArray, ordinal: usize) -> DataFusionResult<ArrayRef> {
+        let child = struct_array.column(ordinal);
+        match struct_array.nulls() {
+            Some(_) => {
+                let combined = NullBuffer::union(struct_array.nulls(), child.nulls());
+                let data = child.to_data().into_builder().nulls(combined).build()?;
+                Ok(make_array(data))
+            }
+            None => Ok(Arc::clone(child)),
+        }
+    }
 }
 
 impl PhysicalExpr for GetStructField {
@@ -75,7 +97,15 @@ impl PhysicalExpr for GetStructField {
     }
 
     fn nullable(&self, input_schema: &Schema) -> DataFusionResult<bool> {
-        Ok(self.child_field(input_schema)?.is_nullable())
+        // A field extracted from a struct is nullable if EITHER the field itself is declared
+        // nullable OR the parent struct can be null -- a field of a null struct is null (Spark
+        // semantics, enforced by `project_field` unioning the parent null mask). Reporting only
+        // the field's own nullability under-declares: a non-nullable field of a nullable struct
+        // then carries the parent's nulls while claiming non-nullable, which fails Arrow's
+        // RecordBatch validation downstream with "declared as non-nullable but contains null
+        // values" (e.g. once the projected column reaches a shuffle/sort). Mirrors Spark's
+        // `GetStructField.nullable = child.nullable || field.nullable`.
+        Ok(self.child.nullable(input_schema)? || self.child_field(input_schema)?.is_nullable())
     }
 
     fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult<ColumnarValue> {
@@ -88,12 +118,13 @@ impl PhysicalExpr for GetStructField {
                     .downcast_ref::<StructArray>()
                     .expect("A struct is expected");
 
-                Ok(ColumnarValue::Array(Arc::clone(
-                    struct_array.column(self.ordinal),
-                )))
+                Ok(ColumnarValue::Array(Self::project_field(
+                    struct_array,
+                    self.ordinal,
+                )?))
             }
             ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) => Ok(ColumnarValue::Array(
-                Arc::clone(struct_array.column(self.ordinal)),
+                Self::project_field(&struct_array, self.ordinal)?,
             )),
             value => Err(DataFusionError::Execution(format!(
                 "Expected a struct array, got {value:?}"
@@ -125,3 +156,79 @@ impl Display for GetStructField {
         )
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::Int64Array;
+    use arrow::datatypes::Fields;
+    use datafusion::physical_expr::expressions::Column;
+
+    // A field of a NULL struct must be NULL (Spark semantics) even when the child buffer holds a
+    // non-null value at that row -- Arrow stores child validity independently of the parent
+    // struct's null mask, so a logically-null struct column read from parquet can still carry a
+    // populated child buffer. Without propagating the parent null mask, `isnotnull(struct.field)`
+    // wrongly evaluates TRUE for a null struct.
+    #[test]
+    fn field_of_null_struct_is_null() {
+        // Child is non-null at every row; the struct itself is null at rows 1 and 3.
+        let child = Arc::new(Int64Array::from(vec![10_i64, 20, 30, 40])) as ArrayRef;
+        let fields: Fields = Fields::from(vec![Field::new("version", DataType::Int64, true)]);
+        let nulls = NullBuffer::from(vec![true, false, true, false]);
+        let struct_array = StructArray::new(fields.clone(), vec![child], Some(nulls));
+        let schema = Schema::new(vec![Field::new("cm", DataType::Struct(fields), true)]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(struct_array)]).unwrap();
+
+        let expr = GetStructField::new(Arc::new(Column::new("cm", 0)), 0);
+        let out = expr
+            .evaluate(&batch)
+            .unwrap()
+            .into_array(batch.num_rows())
+            .unwrap();
+        let out = out.as_any().downcast_ref::<Int64Array>().unwrap();
+
+        assert!(!out.is_null(0) && out.value(0) == 10);
+        assert!(out.is_null(1), "field of a null struct must be null");
+        assert!(!out.is_null(2) && out.value(2) == 30);
+        assert!(out.is_null(3), "field of a null struct must be null");
+    }
+
+    // A NON-nullable field of a NULLABLE struct must report `nullable() == true`: `project_field`
+    // unions the parent struct's null mask, so the projected column carries nulls wherever the
+    // struct is null. Reporting the field's own (non-nullable) flag would make the output schema
+    // lie, failing Arrow RecordBatch validation downstream with "declared as non-nullable but
+    // contains null values" once the column reaches a shuffle/sort.
+    #[test]
+    fn non_nullable_field_of_nullable_struct_is_nullable() {
+        // `size` is declared non-nullable, but the enclosing struct is nullable.
+        let inner: Fields = Fields::from(vec![Field::new("size", DataType::Int64, false)]);
+        let schema = Schema::new(vec![Field::new(
+            "add",
+            DataType::Struct(inner),
+            /* struct nullable */ true,
+        )]);
+
+        let expr = GetStructField::new(Arc::new(Column::new("add", 0)), 0);
+        assert!(
+            expr.nullable(&schema).unwrap(),
+            "a field of a nullable struct must be nullable even if the field itself is non-nullable"
+        );
+    }
+
+    // A non-nullable field of a NON-nullable struct stays non-nullable (no over-declaring).
+    #[test]
+    fn non_nullable_field_of_non_nullable_struct_stays_non_nullable() {
+        let inner: Fields = Fields::from(vec![Field::new("size", DataType::Int64, false)]);
+        let schema = Schema::new(vec![Field::new(
+            "add",
+            DataType::Struct(inner),
+            /* struct nullable */ false,
+        )]);
+
+        let expr = GetStructField::new(Arc::new(Column::new("add", 0)), 0);
+        assert!(
+            !expr.nullable(&schema).unwrap(),
+            "a non-nullable field of a non-nullable struct must remain non-nullable"
+        );
+    }
+}

spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala

@@ -125,6 +125,37 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("GetStructField: non-nullable field of a nullable struct (Delta action-frame shape)") {
+    // Repro for the under-declared `GetStructField` nullability that crashed Comet's native
+    // execution with "Column '...' is declared as non-nullable but contains null values".
+    //
+    // Models Delta's action frame: each log row is exactly ONE action type, so action columns are
+    // NULLABLE structs, but their inner fields are declared NON-nullable by Delta's typed
+    // SingleAction schema (e.g. `add.size`). We build that exact shape with an explicit in-memory
+    // schema (a Parquet round-trip would mark every field nullable, and a CreateNamedStruct would
+    // be declined). Pushing the struct through a Comet shuffle and projecting the non-nullable
+    // inner field (GetStructField) used to produce a null in a column declared non-nullable, which
+    // Comet's native execution rejected during RecordBatch validation.
+    val schema = StructType(
+      Seq(
+        StructField(
+          "add",
+          StructType(Seq(StructField("size", LongType, nullable = false))),
+          nullable = true),
+        StructField("v", IntegerType, nullable = false)))
+    val rows = (0 until 1000).map(i => Row(if (i % 2 == 0) Row(i.toLong) else null, i))
+    withSQLConf(
+      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_SHUFFLE_WITH_HASH_PARTITIONING_ENABLED.key -> "true") {
+      val df = spark
+        .createDataFrame(spark.sparkContext.parallelize(rows), schema)
+        .repartition(4, col("v")) // materialize the typed struct through a Comet shuffle
+        .select(col("add.size").as("size")) // GetStructField on the non-nullable inner field
+        .repartition(4, col("size")) // re-shuffle: read-back validates the declared schema
+      checkSparkAnswer(df)
+    }
+  }
+
   test("compare true/false to negative zero") {
     Seq(false, true).foreach { dictionary =>
       withSQLConf(