feat: support NullType output types in codegen dispatch

.github/workflows/pr_build_linux.yml

@@ -477,6 +477,7 @@ jobs:
               org.apache.comet.CometUuidExpressionSuite
               org.apache.comet.serde.CometScalarFunctionSuite
               org.apache.comet.CometFallbackInvarianceSuite
+              org.apache.comet.CometNullTypeCompositionSuite
       fail-fast: false
     name: ${{ matrix.profile.name }} [${{ matrix.suite.name }}]
     runs-on: ubuntu-24.04

.github/workflows/pr_build_macos.yml

@@ -250,6 +250,7 @@ jobs:
               org.apache.comet.CometUuidExpressionSuite
               org.apache.comet.serde.CometScalarFunctionSuite
               org.apache.comet.CometFallbackInvarianceSuite
+              org.apache.comet.CometNullTypeCompositionSuite
 
       fail-fast: false
     name: ${{ matrix.os }}/${{ matrix.profile.name }} [${{ matrix.suite.name }}]

dev/scalastyle-config.xml

@@ -193,6 +193,19 @@ This file is divided into 3 sections:
     ]]></customMessage>
   </check>
 
+  <check customId="arrowstreamwriter" level="error" class="org.scalastyle.file.RegexChecker" enabled="true">
+    <parameters><parameter name="regex">new ArrowStreamWriter</parameter></parameters>
+    <customMessage><![CDATA[
+      Are you sure that you want to use new ArrowStreamWriter? In most cases, you should use
+      Utils.newArrowStreamWriter instead, which repairs the NullType map keys Arrow reports as
+      nullable (see its scaladoc).
+      If you must use new ArrowStreamWriter, wrap the code block with
+      // scalastyle:off arrowstreamwriter
+      new ArrowStreamWriter(...)
+      // scalastyle:on arrowstreamwriter
+    ]]></customMessage>
+  </check>
+
   <check customId="caselocale" level="error" class="org.scalastyle.file.RegexChecker" enabled="true">
     <parameters><parameter name="regex">(\.toUpperCase|\.toLowerCase)(?!(\(|\(Locale.ROOT\)))</parameter></parameters>
     <customMessage><![CDATA[

docs/source/user-guide/latest/scala_java_udfs.md

@@ -45,7 +45,7 @@ This feature is enabled by default. Set `spark.comet.exec.scalaUDF.codegen.enabl
 - Table UDFs and generators.
 - Python `@udf` and Pandas `@pandas_udf`.
 - Hive `GenericUDF` and `SimpleUDF`.
-- `CalendarIntervalType`, `NullType`, and `UserDefinedType` arguments and return types. UDT-typed columns fall back to Spark; to keep execution in the Comet pipeline, store and read the underlying representation directly (e.g. write MLlib `Vector` outputs as `Struct<type: Byte, size: Int, indices: Array<Int>, values: Array<Double>>` rather than `VectorUDT`).
+- `UserDefinedType` arguments and return types, and `NullType` arguments. UDT-typed columns fall back to Spark; to keep execution in the Comet pipeline, store and read the underlying representation directly (e.g. write MLlib `Vector` outputs as `Struct<type: Byte, size: Int, indices: Array<Int>, values: Array<Double>>` rather than `VectorUDT`). A `NullType` _return_ type is supported: Comet writes an all-null Arrow vector for it.
 - Trees whose total nested-field count (output plus all input columns the UDF tree references) exceeds `spark.sql.codegen.maxFields` (default 100). Comet refuses these at plan time and the operator falls back to Spark.
 
 When a UDF is rejected, the reason surfaces through Comet's standard fallback diagnostics; the query still runs on Spark.

native/shuffle/src/remote_schema_tests.rs

@@ -18,9 +18,9 @@
 use crate::{decode_remote_shuffle_batch, CompressionCodec, ShuffleBlockWriter};
 use arrow::array::{
     Array, ArrayRef, BinaryArray, BinaryDictionaryBuilder, DictionaryArray, FixedSizeListArray,
-    Int16Array, Int32Array, LargeListArray, ListArray, MapArray, PrimitiveDictionaryBuilder,
-    RecordBatch, RecordBatchOptions, StringArray, StringDictionaryBuilder, StructArray,
-    UInt16Array,
+    Int16Array, Int32Array, Int64Array, LargeListArray, ListArray, MapArray, NullArray,
+    PrimitiveDictionaryBuilder, RecordBatch, RecordBatchOptions, StringArray,
+    StringDictionaryBuilder, StructArray, UInt16Array,
 };
 use arrow::buffer::{NullBuffer, OffsetBuffer};
 use arrow::datatypes::{
@@ -570,3 +570,52 @@ fn remote_shuffle_preserves_row_count_without_columns() {
     assert_eq!(decoded.num_columns(), 0);
     assert_eq!(decoded.num_rows(), 3);
 }
+
+// A `NullType` column, and one nested under a list, a map value and a struct field, decode
+// unchanged: a `NullArray` owns no buffers, so the encoding, the dictionary decoding and the
+// nested-nullability reconciliation all have to pass it through by length alone.
+#[test]
+fn null_type_columns_and_children_survive_remote_shuffle() {
+    let rows = 3;
+    let null_field = |name: &str| Arc::new(Field::new(name, DataType::Null, true));
+    let top_level: ArrayRef = Arc::new(NullArray::new(rows));
+    let list: ArrayRef = Arc::new(ListArray::new(
+        null_field("element"),
+        OffsetBuffer::from_lengths([1, 0, 2]),
+        Arc::new(NullArray::new(3)),
+        None,
+    ));
+    let struct_fields = Fields::from(vec![
+        Field::new("a", DataType::Int64, true),
+        Field::new("n", DataType::Null, true),
+    ]);
+    let structs: ArrayRef = Arc::new(StructArray::new(
+        struct_fields,
+        vec![
+            Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])),
+            Arc::new(NullArray::new(rows)),
+        ],
+        Some(NullBuffer::from(vec![true, false, true])),
+    ));
+    let entry_fields = Fields::from(vec![
+        Field::new("key", DataType::Int64, false),
+        Field::new("value", DataType::Null, true),
+    ]);
+    let entries = StructArray::new(
+        entry_fields.clone(),
+        vec![
+            Arc::new(Int64Array::from(vec![10, 20, 30])),
+            Arc::new(NullArray::new(3)),
+        ],
+        None,
+    );
+    let map: ArrayRef = Arc::new(MapArray::new(
+        Arc::new(Field::new("entries", DataType::Struct(entry_fields), false)),
+        OffsetBuffer::from_lengths([2, 0, 1]),
+        entries,
+        None,
+        false,
+    ));
+    let columns = vec![top_level, list, structs, map];
+    assert_roundtrip(columns.clone(), columns);
+}

native/shuffle/src/spark_unsafe/row.rs

@@ -572,6 +572,11 @@ fn append_nested_struct_fields_field_major(
                     }
                 }
             }
+            // A Null field carries no data: every row is null, whether or not the struct is.
+            DataType::Null => {
+                let field_builder = get_field_builder!(struct_builder, NullBuilder, field_idx);
+                field_builder.append_nulls(num_rows);
+            }
             _ => {
                 unreachable!(
                     "Unsupported data type of struct field: {:?}",
@@ -1020,6 +1025,11 @@ fn append_struct_fields_field_major(
                     }
                 }
             }
+            // A Null field carries no data: every row is null, whether or not the struct is.
+            DataType::Null => {
+                let field_builder = get_field_builder!(struct_builder, NullBuilder, field_idx);
+                field_builder.append_nulls(row_end - row_start);
+            }
             _ => {
                 unreachable!(
                     "Unsupported data type of struct field: {:?}",
@@ -1179,9 +1189,7 @@ fn append_columns(
         }
         DataType::Null => {
             let null_builder = downcast_builder_ref!(NullBuilder, builder);
-            for _ in row_start..row_end {
-                null_builder.append_null();
-            }
+            null_builder.append_nulls(row_end - row_start);
         }
         DataType::Timestamp(TimeUnit::Microsecond, _) => {
             append_column_to_builder!(
@@ -1418,6 +1426,14 @@ pub fn process_sorted_row_partition(
             .collect();
         let batch = make_batch(array_refs?, n)?;
 
+        // A finished `NullBuilder` keeps its length; see `recreate_null_type_builders`.
+        recreate_null_type_builders(
+            &mut data_builders,
+            schema,
+            batch_size,
+            prefer_dictionary_ratio,
+        )?;
+
         frozen.clear();
         let mut cursor = Cursor::new(&mut frozen);
 
@@ -1440,6 +1456,37 @@ pub fn process_sorted_row_partition(
     ))
 }
 
+/// Whether `dt` is `Null` or nests a `Null` anywhere below a list, struct or map.
+fn contains_null_type(dt: &DataType) -> bool {
+    match dt {
+        DataType::Null => true,
+        DataType::List(field) | DataType::LargeList(field) | DataType::Map(field, _) => {
+            contains_null_type(field.data_type())
+        }
+        DataType::Struct(fields) => fields.iter().any(|f| contains_null_type(f.data_type())),
+        _ => false,
+    }
+}
+
+/// Replaces every builder whose type holds a `Null` somewhere, after its batch has been
+/// finished. `NullBuilder::finish` keeps its length (a `NullArray` owns no buffers to hand
+/// over), so such a builder would carry this batch's rows into the next: a top-level Null column
+/// comes out longer than the batch, and a Null struct field longer than its parent panics in
+/// `StructBuilder::finish`. Every other builder resets on finish and is kept.
+fn recreate_null_type_builders(
+    builders: &mut [Box<dyn ArrayBuilder>],
+    schema: &[DataType],
+    batch_size: usize,
+    prefer_dictionary_ratio: f64,
+) -> Result<(), CometError> {
+    for (builder, datatype) in builders.iter_mut().zip(schema.iter()) {
+        if contains_null_type(datatype) {
+            *builder = make_builders(datatype, batch_size, prefer_dictionary_ratio)?;
+        }
+    }
+    Ok(())
+}
+
 fn builder_to_array(
     builder: &mut Box<dyn ArrayBuilder>,
     datatype: &DataType,
@@ -1503,10 +1550,98 @@ fn make_batch(arrays: Vec<ArrayRef>, row_count: usize) -> Result<RecordBatch, Ar
 
 #[cfg(test)]
 mod test {
+    use arrow::array::StructArray;
     use arrow::datatypes::Fields;
 
     use super::*;
 
+    // The batch loop in `process_sorted_row_partition` reuses its builders. A `NullBuilder`
+    // keeps its length across `finish`, so without `recreate_null_type_builders` the second
+    // batch's Null struct field is twice as long as its parent (a panic in `StructBuilder::
+    // finish`) and a top-level Null column is twice as long as the batch.
+    #[test]
+    fn null_type_builders_start_every_batch_empty() {
+        let struct_type = DataType::Struct(Fields::from(vec![
+            Field::new("v", DataType::Int64, true),
+            Field::new("n", DataType::Null, true),
+        ]));
+        let nested_type = DataType::Struct(Fields::from(vec![Field::new(
+            "s",
+            struct_type.clone(),
+            true,
+        )]));
+        let schema = vec![struct_type.clone(), nested_type.clone(), DataType::Null];
+        assert!(schema.iter().all(contains_null_type));
+        assert!(!contains_null_type(&DataType::List(Arc::new(Field::new(
+            "item",
+            DataType::Int64,
+            true
+        )))));
+
+        let batch_size = 2;
+        let mut builders: Vec<Box<dyn ArrayBuilder>> = schema
+            .iter()
+            .map(|dt| make_builders(dt, batch_size, 1.0).unwrap())
+            .collect();
+
+        for batch in 0..2 {
+            let struct_builder = builders[0]
+                .as_any_mut()
+                .downcast_mut::<StructBuilder>()
+                .unwrap();
+            for row in 0..batch_size {
+                struct_builder
+                    .field_builder::<Int64Builder>(0)
+                    .unwrap()
+                    .append_value(row as i64);
+                struct_builder
+                    .field_builder::<NullBuilder>(1)
+                    .unwrap()
+                    .append_null();
+                struct_builder.append(true);
+            }
+            let nested_builder = builders[1]
+                .as_any_mut()
+                .downcast_mut::<StructBuilder>()
+                .unwrap();
+            for _ in 0..batch_size {
+                let inner = nested_builder.field_builder::<StructBuilder>(0).unwrap();
+                inner
+                    .field_builder::<Int64Builder>(0)
+                    .unwrap()
+                    .append_null();
+                inner.field_builder::<NullBuilder>(1).unwrap().append_null();
+                inner.append(true);
+                nested_builder.append(true);
+            }
+            builders[2]
+                .as_any_mut()
+                .downcast_mut::<NullBuilder>()
+                .unwrap()
+                .append_nulls(batch_size);
+
+            let arrays: Vec<ArrayRef> = builders
+                .iter_mut()
+                .zip(schema.iter())
+                .map(|(builder, dt)| builder_to_array(builder, dt, 1.0).unwrap())
+                .collect();
+            for (array, dt) in arrays.iter().zip(schema.iter()) {
+                assert_eq!(array.len(), batch_size, "batch {batch} of {dt}");
+            }
+            let outer = arrays[0].as_any().downcast_ref::<StructArray>().unwrap();
+            assert_eq!(
+                outer.column(1).len(),
+                batch_size,
+                "batch {batch} Null field"
+            );
+
+            recreate_null_type_builders(&mut builders, &schema, batch_size, 1.0).unwrap();
+            for (builder, dt) in builders.iter().zip(schema.iter()) {
+                assert_eq!(builder.len(), 0, "builder for {dt} after batch {batch}");
+            }
+        }
+    }
+
     #[test]
     fn test_append_null_row_to_struct_builder() {
         let data_type = DataType::Struct(Fields::from(vec![

native/spark-expr/src/array_funcs/size.rs

@@ -130,11 +130,10 @@ fn spark_size_list_like(array: &ArrayRef) -> Result<ArrayRef, DataFusionError> {
         }
     };
 
-    // Fast path for the production shape: `CometSize.convert` wraps size() in a
-    // `CASE WHEN isnotnull(child)` that filters null rows out before the THEN
-    // branch runs, so this function only ever sees a null-free array in a real
-    // Comet plan. Return the length kernel output as-is; skip the downcast and
-    // `Int32Array::clone` that `rewrite_nulls_to_minus_one` would otherwise pay.
+    // Fast path for a null-free array: return the length kernel output as-is and
+    // skip the downcast and `Int32Array::clone` that `rewrite_nulls_to_minus_one`
+    // would otherwise pay. `CometSize.convert` sends nullable input here only in
+    // legacy mode, where -1 is the answer for a null collection.
     if array.null_count() == 0 {
         return Ok(lengths);
     }