diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index a019e9a721f..0930d93356b 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -499,6 +499,7 @@ jobs:
org.apache.comet.serde.CometScalarFunctionSuite
org.apache.comet.serde.CometLiteralSuite
org.apache.comet.CometFallbackInvarianceSuite
+ org.apache.comet.CometNullTypeCompositionSuite
fail-fast: false
name: ${{ matrix.profile.name }} [${{ matrix.suite.name }}]
runs-on: ubuntu-24.04
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index 66a1b62ba3d..63cb7023523 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -253,6 +253,7 @@ jobs:
org.apache.comet.serde.CometScalarFunctionSuite
org.apache.comet.serde.CometLiteralSuite
org.apache.comet.CometFallbackInvarianceSuite
+ org.apache.comet.CometNullTypeCompositionSuite
fail-fast: false
name: ${{ matrix.os }}/${{ matrix.profile.name }} [${{ matrix.suite.name }}]
diff --git a/dev/scalastyle-config.xml b/dev/scalastyle-config.xml
index 9de6df51ef7..edd3de30730 100644
--- a/dev/scalastyle-config.xml
+++ b/dev/scalastyle-config.xml
@@ -193,6 +193,19 @@ This file is divided into 3 sections:
]]>
+
+ new ArrowStreamWriter
+
+
+
(\.toUpperCase|\.toLowerCase)(?!(\(|\(Locale.ROOT\)))
, values: Array>` 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, values: Array>` 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.
+- Struct types with duplicate field names (`named_struct('x', 1, 'x', 2)`) anywhere in an argument or return type. Arrow's `StructVector` addresses children by name, so the two fields would collapse into one; 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.
diff --git a/native/common/src/struct_nulls.rs b/native/common/src/struct_nulls.rs
index f3390707702..8def9081964 100644
--- a/native/common/src/struct_nulls.rs
+++ b/native/common/src/struct_nulls.rs
@@ -28,8 +28,9 @@
//! whatever sits in the child slot
//! ().
-use arrow::array::{make_array, Array, ArrayRef, StructArray};
+use arrow::array::{make_array, Array, ArrayRef, NullArray, StructArray};
use arrow::buffer::NullBuffer;
+use arrow::datatypes::DataType;
use arrow::error::ArrowError;
use std::sync::Arc;
@@ -48,6 +49,14 @@ pub fn child_with_parent_nulls(
ordinal: usize,
) -> Result {
let child = struct_array.column(ordinal);
+ // A Null-typed child is all-null whatever the parent's mask says, and a `NullArray` may not
+ // carry a validity bitmap ("Arrays of type Null cannot contain a null bitmask"). Rebuild it
+ // rather than reuse the child: a kernel that grew the struct through `MutableArrayData`
+ // (`element_at` on an out-of-range index) can hand over a child that already carries such a
+ // bitmap, which only fails validation once it is projected out.
+ if child.data_type() == &DataType::Null {
+ return Ok(Arc::new(NullArray::new(child.len())));
+ }
match struct_array.nulls() {
Some(parent) if parent.null_count() > 0 => {
let combined = NullBuffer::union(Some(parent), child.nulls());
@@ -85,6 +94,22 @@ mod tests {
StructArray::new(fields, vec![child], nulls)
}
+ /// A Null-typed child cannot carry a validity bitmap, so it is rebuilt all-null instead of
+ /// receiving the parent's mask; a `MutableArrayData`-grown struct can even hand over a child
+ /// that already carries one, which only fails validation once projected out.
+ #[test]
+ fn null_typed_child_is_rebuilt_without_a_bitmap() {
+ let fields: Fields = vec![Arc::new(Field::new("n", DataType::Null, true))].into();
+ let parent = NullBuffer::from(vec![true, false, true]);
+ let struct_array =
+ StructArray::new(fields, vec![Arc::new(NullArray::new(3))], Some(parent));
+ let out = child_with_parent_nulls(&struct_array, 0).unwrap();
+ assert_eq!(out.data_type(), &DataType::Null);
+ assert_eq!(out.len(), 3);
+ assert!(out.nulls().is_none(), "a NullArray must not carry a validity bitmap");
+ out.to_data().validate_full().unwrap();
+ }
+
/// The case both bugs came from: the parent is null where the child still holds a value.
#[test]
fn hidden_child_value_under_a_null_parent_becomes_null() {
diff --git a/native/shuffle/src/remote_schema_tests.rs b/native/shuffle/src/remote_schema_tests.rs
index 52fafa0cf6c..ae2089e2379 100644
--- a/native/shuffle/src/remote_schema_tests.rs
+++ b/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);
+}
diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs
index a3f3b3d36fc..263b7a33662 100644
--- a/native/shuffle/src/spark_unsafe/row.rs
+++ b/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],
+ 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,
datatype: &DataType,
@@ -1503,10 +1550,98 @@ fn make_batch(arrays: Vec, row_count: usize) -> Result> = 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::()
+ .unwrap();
+ for row in 0..batch_size {
+ struct_builder
+ .field_builder::(0)
+ .unwrap()
+ .append_value(row as i64);
+ struct_builder
+ .field_builder::(1)
+ .unwrap()
+ .append_null();
+ struct_builder.append(true);
+ }
+ let nested_builder = builders[1]
+ .as_any_mut()
+ .downcast_mut::()
+ .unwrap();
+ for _ in 0..batch_size {
+ let inner = nested_builder.field_builder::(0).unwrap();
+ inner
+ .field_builder::(0)
+ .unwrap()
+ .append_null();
+ inner.field_builder::(1).unwrap().append_null();
+ inner.append(true);
+ nested_builder.append(true);
+ }
+ builders[2]
+ .as_any_mut()
+ .downcast_mut::()
+ .unwrap()
+ .append_nulls(batch_size);
+
+ let arrays: Vec = 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::().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![
diff --git a/native/spark-expr/src/array_funcs/size.rs b/native/spark-expr/src/array_funcs/size.rs
index 2311da5fa9e..a1776d0f308 100644
--- a/native/spark-expr/src/array_funcs/size.rs
+++ b/native/spark-expr/src/array_funcs/size.rs
@@ -130,11 +130,10 @@ fn spark_size_list_like(array: &ArrayRef) -> Result {
}
};
- // 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);
}
diff --git a/native/spark-expr/src/csv_funcs/to_csv.rs b/native/spark-expr/src/csv_funcs/to_csv.rs
index 01fdc901cb7..7e74594aea4 100644
--- a/native/spark-expr/src/csv_funcs/to_csv.rs
+++ b/native/spark-expr/src/csv_funcs/to_csv.rs
@@ -80,8 +80,10 @@ impl PhysicalExpr for ToCsv {
Ok(DataType::Utf8)
}
- fn nullable(&self, input_schema: &Schema) -> Result {
- self.expr.nullable(input_schema)
+ fn nullable(&self, _: &Schema) -> Result {
+ // Spark's `StructsToCsv.nullable` is unconditionally true: a row that renders to an empty
+ // string (a lone null field) yields NULL, whatever the input struct's nullability.
+ Ok(true)
}
fn evaluate(&self, batch: &RecordBatch) -> Result {
@@ -201,7 +203,15 @@ pub fn to_csv_inner(
}
}
}
- builder.append_value(&csv_string);
+ // Spark renders a row through univocity's `writeRowToString`, whose `skipEmptyLines`
+ // setting (always on for Spark's CSV writer) turns an empty rendered row into `null`:
+ // a struct with a single null field (with the default empty `nullValue`) is NULL,
+ // not "". A row with two null fields renders as the delimiter and is kept.
+ if csv_string.is_empty() {
+ builder.append_null();
+ } else {
+ builder.append_value(&csv_string);
+ }
}
}
Ok(Arc::new(builder.finish()))
@@ -216,3 +226,74 @@ fn escape_value(value: &str, quote_char: char, escape_char: char, output: &mut S
output.push(ch);
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::Int64Array;
+ use arrow::datatypes::{Field, Fields};
+
+ fn default_write_options() -> CsvWriteOptions {
+ CsvWriteOptions::new(
+ ",".to_string(),
+ "\"".to_string(),
+ "\\".to_string(),
+ "".to_string(),
+ false,
+ true,
+ true,
+ )
+ }
+
+ fn to_csv_strings(array: &StructArray, write_options: &CsvWriteOptions) -> Vec> {
+ let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false);
+ cast_options.null_string = write_options.null_value.clone();
+ let out = to_csv_inner(array, &cast_options, write_options).unwrap();
+ as_string_array(&out)
+ .iter()
+ .map(|v| v.map(str::to_string))
+ .collect()
+ }
+
+ // Spark's `to_csv` hands the row to univocity's `writeRowToString` with `skipEmptyLines`, so a
+ // row whose rendering is empty comes back NULL rather than "". Only a lone null field (or a
+ // lone empty `nullValue`) renders empty; two null fields render as the delimiter.
+ #[test]
+ fn empty_rendered_row_is_null() {
+ let one_field: Fields = Fields::from(vec![Field::new("a", DataType::Int64, true)]);
+ let single = StructArray::new(
+ one_field,
+ vec![Arc::new(Int64Array::from(vec![Some(7_i64), None]))],
+ None,
+ );
+ assert_eq!(
+ to_csv_strings(&single, &default_write_options()),
+ vec![Some("7".to_string()), None]
+ );
+
+ let two_fields: Fields = Fields::from(vec![
+ Field::new("a", DataType::Int64, true),
+ Field::new("n", DataType::Null, true),
+ ]);
+ let pair = StructArray::new(
+ two_fields,
+ vec![
+ Arc::new(Int64Array::from(vec![Some(7_i64), None])),
+ Arc::new(arrow::array::NullArray::new(2)),
+ ],
+ None,
+ );
+ assert_eq!(
+ to_csv_strings(&pair, &default_write_options()),
+ vec![Some("7,".to_string()), Some(",".to_string())]
+ );
+
+ // A non-empty `nullValue` renders, so the row is kept.
+ let mut named_null = default_write_options();
+ named_null.null_value = "N".to_string();
+ assert_eq!(
+ to_csv_strings(&single, &named_null),
+ vec![Some("7".to_string()), Some("N".to_string())]
+ );
+ }
+}
diff --git a/native/spark-expr/src/struct_funcs/get_struct_field.rs b/native/spark-expr/src/struct_funcs/get_struct_field.rs
index b815b4ed9a0..a5eee75b914 100644
--- a/native/spark-expr/src/struct_funcs/get_struct_field.rs
+++ b/native/spark-expr/src/struct_funcs/get_struct_field.rs
@@ -99,9 +99,17 @@ impl PhysicalExpr for GetStructField {
self.ordinal,
)?))
}
- ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) => Ok(ColumnarValue::Array(
- child_with_parent_nulls(&struct_array, self.ordinal)?,
- )),
+ // A scalar struct (a NULL branch of CASE, a scalar subquery) holds one row, so its
+ // field is a scalar too. Returning the projected one-row array instead makes the
+ // consumer believe it has a full column: a CASE result builder then slices row
+ // ranges out of it and panics ("range end index 2 out of range for slice of length
+ // 1").
+ ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) => {
+ let projected = child_with_parent_nulls(&struct_array, self.ordinal)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(
+ &projected, 0,
+ )?))
+ }
value => Err(DataFusionError::Execution(format!(
"Expected a struct array, got {value:?}"
))),
@@ -136,16 +144,58 @@ impl Display for GetStructField {
#[cfg(test)]
mod tests {
use super::*;
- use arrow::array::{ArrayRef, Int64Array};
+ use arrow::array::{ArrayRef, Int64Array, NullArray};
use arrow::buffer::NullBuffer;
use arrow::datatypes::Fields;
- use datafusion::physical_expr::expressions::Column;
+ use datafusion::physical_expr::expressions::{Column, Literal};
// 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.
+ // A scalar struct input yields a scalar field, so a consumer that broadcasts it over the batch
+ // (a CASE branch, a projection beside a column) sees one value per row rather than a one-row
+ // array. A null struct scalar yields a null scalar, for a typed field and a Null-typed one.
+ #[test]
+ fn field_of_scalar_struct_is_scalar() {
+ let fields: Fields = Fields::from(vec![
+ Field::new("a", DataType::Int64, true),
+ Field::new("n", DataType::Null, true),
+ ]);
+ let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(Int64Array::from(vec![0_i64, 1, 2, 3]))],
+ )
+ .unwrap();
+
+ let present = ScalarValue::Struct(Arc::new(StructArray::new(
+ fields.clone(),
+ vec![
+ Arc::new(Int64Array::from(vec![7_i64])),
+ Arc::new(NullArray::new(1)),
+ ],
+ None,
+ )));
+ let absent = ScalarValue::Struct(Arc::new(StructArray::new_null(fields, 1)));
+
+ for (scalar, ordinal, expected) in [
+ (present.clone(), 0, ScalarValue::Int64(Some(7))),
+ (present, 1, ScalarValue::Null),
+ (absent.clone(), 0, ScalarValue::Int64(None)),
+ (absent, 1, ScalarValue::Null),
+ ] {
+ let expr = GetStructField::new(Arc::new(Literal::new(scalar)), ordinal);
+ match expr.evaluate(&batch).unwrap() {
+ ColumnarValue::Scalar(value) => assert_eq!(value, expected),
+ ColumnarValue::Array(array) => {
+ panic!("expected a scalar for ordinal {ordinal}, got array {array:?}")
+ }
+ }
+ }
+ }
+
#[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.
diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala
index 83fbca6b635..57de9105d47 100644
--- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala
+++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala
@@ -82,13 +82,17 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
* Type surface the kernel covers on both input and output sides. Recursive: complex types are
* supported when their children are.
*
- * Duplicate struct field names are excluded, an output-side rule: Arrow addresses a
- * `StructVector`'s children by name, so `named_struct('x', 10, 'x', 20)` collapses into a
- * single child and the generated writer NPEs on the missing ordinal-1 vector.
- * `CometCreateNamedStruct` declines them on the native path for the same reason, but a struct
- * nested inside a dispatcher-built value (a `CreateMap` value) never reaches that check.
+ * Duplicate struct field names are a separate rule, [[hasDuplicateStructFieldNames]], so that
+ * [[canHandle]] can name them in its reason.
+ *
+ * `NullType` is output-only: [[CometBatchKernelCodegenOutput]] can write an all-null Arrow
+ * `NullVector`, but `CometScalaUDFCodegen.specFor` cannot build an [[ArrowColumnSpec]] for one,
+ * so a `NullType` input (nested or not) has to keep falling back to Spark.
*/
- def isSupportedDataType(dt: DataType): Boolean = dt match {
+ def isSupportedDataType(dt: DataType): Boolean = isSupportedDataType(dt, allowNullType = false)
+
+ private def isSupportedDataType(dt: DataType, allowNullType: Boolean): Boolean = dt match {
+ case NullType => allowNullType
case BooleanType | ByteType | ShortType | IntegerType | LongType => true
case FloatType | DoubleType => true
case _: DecimalType => true
@@ -96,16 +100,43 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
case DateType | TimestampType | TimestampNTZType => true
case dt if isTimeType(dt) => true
case _: YearMonthIntervalType | _: DayTimeIntervalType | CalendarIntervalType => true
- case ArrayType(inner, _) => isSupportedDataType(inner)
+ case ArrayType(inner, _) => isSupportedDataType(inner, allowNullType)
+ case st: StructType => st.fields.forall(f => isSupportedDataType(f.dataType, allowNullType))
+ case mt: MapType =>
+ isSupportedDataType(mt.keyType, allowNullType) &&
+ isSupportedDataType(mt.valueType, allowNullType)
+ case _ => false
+ }
+
+ /**
+ * Spark keeps duplicate struct field names as distinct positional fields, but Arrow's
+ * `StructVector` keys its children by name (`ConflictPolicy.CONFLICT_REPLACE`), so
+ * `named_struct('x', 10, 'x', 20)` collapses into a single child and the generated
+ * ordinal-based child casts hit a missing or differently typed vector. `CometCreateNamedStruct`
+ * refuses the same shape on the native path, but whole-expression dispatch never consults that
+ * rule for a `named_struct` nested inside e.g. a `transform` lambda or a `CreateMap` value.
+ */
+ private def hasDuplicateStructFieldNames(dt: DataType): Boolean = dt match {
case st: StructType =>
// `fieldNames` rebuilds an array on each call, so read it once.
val names = st.fieldNames
- names.distinct.length == names.length &&
- st.fields.forall(f => isSupportedDataType(f.dataType))
- case mt: MapType => isSupportedDataType(mt.keyType) && isSupportedDataType(mt.valueType)
+ names.distinct.length != names.length ||
+ st.fields.exists(f => hasDuplicateStructFieldNames(f.dataType))
+ case ArrayType(inner, _) => hasDuplicateStructFieldNames(inner)
+ case MapType(k, v, _) => hasDuplicateStructFieldNames(k) || hasDuplicateStructFieldNames(v)
case _ => false
}
+ /** Why `dt` cannot cross the kernel boundary as `side` ("output" or "input"), if it cannot. */
+ private def typeRejection(dt: DataType, allowNullType: Boolean, side: String): Option[String] =
+ if (!isSupportedDataType(dt, allowNullType)) {
+ Some(s"codegen dispatch: unsupported $side type $dt")
+ } else if (hasDuplicateStructFieldNames(dt)) {
+ Some(s"codegen dispatch: unsupported $side type $dt (duplicate struct field name)")
+ } else {
+ None
+ }
+
/**
* Mirrors `WholeStageCodegenExec.numOfNestedFields` so [[canHandle]] can reuse
* `spark.sql.codegen.maxFields`.
@@ -123,12 +154,13 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
* back cleanly rather than crashing the Janino compile at execute time.
*
* Checks every `BoundReference`'s data type and the root `expr.dataType` against
- * [[isSupportedDataType]], rejects aggregates / generators / `Unevaluable`, and gates total
- * nested-field count on `spark.sql.codegen.maxFields`.
+ * [[isSupportedDataType]] and [[hasDuplicateStructFieldNames]], rejects aggregates, generators
+ * and `Unevaluable`, and gates total nested-field count on `spark.sql.codegen.maxFields`.
*/
def canHandle(boundExpr: Expression): Option[String] = {
- if (!isSupportedDataType(boundExpr.dataType)) {
- return Some(s"codegen dispatch: unsupported output type ${boundExpr.dataType}")
+ typeRejection(boundExpr.dataType, allowNullType = true, "output") match {
+ case Some(reason) => return Some(reason)
+ case None =>
}
// Mirror WSCG's `spark.sql.codegen.maxFields` gate. Wide schemas blow the generated class's
// typed input field count, the typed-getter switch, and the constant pool. Refuse here so the
@@ -178,12 +210,15 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
"(aggregate, generator, or unevaluable)")
case None =>
}
- val badRef = boundExpr.collectFirst {
- case b: BoundReference if !isSupportedDataType(b.dataType) =>
- b
- }
- badRef.map(b =>
- s"codegen dispatch: unsupported input type ${b.dataType} at ordinal ${b.ordinal}")
+ boundExpr.collectFirst(Function.unlift(inputRejection))
+ }
+
+ /** Why `expr` cannot be read as a codegen input, if it cannot. */
+ private def inputRejection(expr: Expression): Option[String] = expr match {
+ case b: BoundReference =>
+ typeRejection(b.dataType, allowNullType = false, "input")
+ .map(reason => s"$reason at ordinal ${b.ordinal}")
+ case _ => None
}
/**
diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala
index 33e6c0c0355..36902ebee18 100644
--- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala
+++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala
@@ -138,9 +138,22 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
override def getField: Field = exportField
}
+ /**
+ * The superclass constructor builds a `NullableStructWriter` over `getField.getChildren`, and
+ * that writer has no arm for a Null child (`UnsupportedOperationException: Unknown type: NULL`
+ * for a `struct<..., x: null>` output). So the superclass is given a childless copy of the
+ * field, `getField` answers with it until construction completes (`constructed` is still the
+ * JVM default `false` while the superclass constructor runs, since this class's own fields are
+ * initialized afterwards), and the children come from `initializeChildrenFromFields` in
+ * [[allocateOutput]], the same way a struct nested under a list gets them.
+ */
private final class RenamedStructVector(exportField: Field, allocator: BufferAllocator)
- extends StructVector(exportField, allocator, null) {
- override def getField: Field = exportField
+ extends StructVector(
+ new Field(exportField.getName, exportField.getFieldType, null),
+ allocator,
+ null) {
+ private val constructed: Boolean = true
+ override def getField: Field = if (constructed) exportField else super.getField
}
/**
@@ -179,6 +192,7 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
case _: ArrayType => classOf[ListVector].getName
case _: StructType => classOf[StructVector].getName
case _: MapType => classOf[MapVector].getName
+ case NullType => classOf[NullVector].getName
case other =>
throw new UnsupportedOperationException(
s"CometBatchKernelCodegen.outputVectorClass: unsupported output type $other")
@@ -209,6 +223,11 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
dataType: DataType,
ctx: CodegenContext,
nested: Boolean = false): OutputEmit = dataType match {
+ case NullType =>
+ // A NullType value is null by definition: nothing to read from `source`, and `NullVector`
+ // has no data buffer. `setNull` is a no-op; the all-null semantics come from
+ // `CometScalaUDFCodegen.evaluate`'s post-`process` `setValueCount`.
+ OutputEmit("", s"$targetVec.setNull($idx);")
case BooleanType =>
val set = if (nested) "setSafe" else "set"
OutputEmit("", s"$targetVec.$set($idx, $source ? 1 : 0);")
@@ -407,6 +426,10 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
*/
private def emitSpecializedGetterExpr(target: String, idx: String, elemType: DataType): String =
elemType match {
+ // Computed eagerly for every child, then dropped by [[emitWrite]]'s NullType branch, which
+ // emits `setNull` instead. Dead in the generated Java but not here: without this case an
+ // `array` / `struct<.., null>` output throws instead of dispatching.
+ case NullType => "null"
case BooleanType => s"$target.getBoolean($idx)"
case ByteType => s"$target.getByte($idx)"
case ShortType => s"$target.getShort($idx)"
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index be4bc9c3412..3c03d15318c 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -31,6 +31,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate._
import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke}
import org.apache.spark.sql.catalyst.expressions.xml.{XPathBoolean, XPathDouble, XPathFloat, XPathInt, XPathList, XPathLong, XPathShort, XPathString}
import org.apache.spark.sql.comet.DecimalPrecision
+import org.apache.spark.sql.comet.util.Utils
import org.apache.spark.sql.execution.{ScalarSubquery, SparkPlan}
import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils
import org.apache.spark.sql.internal.SQLConf
@@ -72,7 +73,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
classOf[ArrayMin] -> CometArrayMin,
classOf[ArrayPosition] -> CometArrayPosition,
classOf[ArrayRemove] -> CometArrayRemove,
- classOf[ArrayRepeat] -> CometScalarFunction("array_repeat"),
+ classOf[ArrayRepeat] -> CometArrayRepeat,
classOf[Slice] -> CometSlice,
classOf[SortArray] -> CometSortArray,
classOf[ArraysOverlap] -> CometArraysOverlap,
@@ -645,7 +646,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
val info = DataTypeInfo.newBuilder()
val list = ListInfo.newBuilder()
list.setElementType(elementType.get)
- list.setContainsNull(a.containsNull)
+ list.setContainsNull(Utils.declaredChildNullability(a.elementType, a.containsNull))
nestedParquetFieldId(parentField, elementPath, includeFieldIds)
.foreach(list.setElementFieldId)
@@ -669,7 +670,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
val map = MapInfo.newBuilder()
map.setKeyType(keyType.get)
map.setValueType(valueType.get)
- map.setValueContainsNull(m.valueContainsNull)
+ map.setValueContainsNull(Utils.declaredChildNullability(m.valueType, m.valueContainsNull))
nestedParquetFieldId(parentField, keyPath, includeFieldIds).foreach(map.setKeyFieldId)
nestedParquetFieldId(parentField, valuePath, includeFieldIds).foreach(map.setValueFieldId)
@@ -686,7 +687,8 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
val nestedFieldPath = if (nestedParentField.isDefined) Seq(field.name) else Seq.empty
serializeDataType(field.dataType, nestedParentField, nestedFieldPath, includeFieldIds)
}
- val fieldNullable = s.map(f => Boolean.box(f.nullable)).asJava
+ val fieldNullable =
+ s.map(f => Boolean.box(Utils.declaredChildNullability(f.dataType, f.nullable))).asJava
if (fieldDatatypes.exists(_.isEmpty)) {
return None
diff --git a/spark/src/main/scala/org/apache/comet/serde/SupportLevel.scala b/spark/src/main/scala/org/apache/comet/serde/SupportLevel.scala
index 60f2ff79345..1c64896227a 100644
--- a/spark/src/main/scala/org/apache/comet/serde/SupportLevel.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/SupportLevel.scala
@@ -19,6 +19,7 @@
package org.apache.comet.serde
+import org.apache.spark.sql.catalyst.expressions.Expression
import org.apache.spark.sql.types._
import org.apache.comet.CometConf
@@ -103,3 +104,35 @@ object SupportLevel {
}
}
}
+
+/**
+ * Serdes that wrap a kernel in `CASE WHEN guarded IS NOT NULL THEN kernel(children) ELSE NULL`
+ * serialize the guarded children twice, and a stateful child (one built over
+ * monotonically_increasing_id(), say) then produces values Spark's single evaluation never would,
+ * for either of two reasons:
+ *
+ * - A child evaluated natively gets an independent instance per copy, but native CASE evaluates
+ * the THEN branch on the rows the predicate selected, so whenever the guard filters, the THEN
+ * copy's counter runs over a different row sequence than Spark's.
+ * - A child evaluated by the JVM codegen dispatcher (any lambda function) is keyed in the
+ * kernel cache by its serialized bytes, so both copies run one kernel instance and share its
+ * state: the predicate copy consumes the counter for the whole batch and the THEN copy
+ * continues from there, even when nothing filters.
+ *
+ * Neither the stateful child nor the nullable one needs to be the same child, and the stateful
+ * child need not be nullable at all; a serde that builds no guard for a non-nullable child
+ * (`CometSize`) needs no gate for it either.
+ */
+object NullGuard {
+ val reason =
+ "non-deterministic child under a null guard is evaluated on different rows than Spark's"
+
+ def doubleEvaluationReason(evaluated: Seq[Expression]): Option[String] =
+ if (evaluated.exists(!_.deterministic)) Some(reason) else None
+
+ /** [[Unsupported]] when a child inside the guard is non-deterministic, else [[Compatible]]. */
+ def supportLevel(evaluated: Expression*): SupportLevel =
+ doubleEvaluationReason(evaluated)
+ .map(reason => Unsupported(Some(reason)))
+ .getOrElse(Compatible())
+}
diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
index aef649debc7..487dac1d30d 100644
--- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
@@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Expression, L
import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, ApproximatePercentile, Average, BitAndAgg, BitOrAgg, BitXorAgg, BloomFilterAggregate, CentralMomentAgg, CollectList, CollectSet, Corr, Count, Covariance, CovPopulation, CovSample, First, HyperLogLogPlusPlus, Last, Max, MaxBy, MaxMinBy, Min, MinBy, Mode, Percentile, RegrIntercept, RegrR2, RegrReplacement, RegrSlope, RegrSXY, StddevPop, StddevSamp, Sum, VariancePop, VarianceSamp}
import org.apache.spark.sql.catalyst.util.ArrayData
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, NumericType, ShortType, StringType, TimestampNTZType, TimestampType}
+import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, NullType, NumericType, ShortType, StringType, TimestampNTZType, TimestampType}
import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT
import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus, isSpark42Plus, withFallbackReason}
@@ -1053,13 +1053,15 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] {
if (!CometCollectShim.ignoreNulls(expr)) {
Unsupported(Some("collect_set with RESPECT NULLS (ignoreNulls = false) is not supported"))
} else {
- SupportLevel
- .strictFloatingPointReason(
- expr.children.head.dataType,
- "collect_set on floating-point types " +
- "(Comet deduplicates NaN values while Spark treats each NaN as distinct)")
- .map(reason => Incompatible(Some(reason)))
- .getOrElse(Compatible())
+ CometCollectAggregate.nestedNullLevel(expr.children.head.dataType).getOrElse {
+ SupportLevel
+ .strictFloatingPointReason(
+ expr.children.head.dataType,
+ "collect_set on floating-point types " +
+ "(Comet deduplicates NaN values while Spark treats each NaN as distinct)")
+ .map(reason => Incompatible(Some(reason)))
+ .getOrElse(Compatible())
+ }
}
}
@@ -1092,6 +1094,29 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] {
}
}
+/**
+ * Shared gate for collect_list / collect_set. Native collect accumulates the input's own child
+ * array, so every nested element keeps the non-nullability its producer gave it (`map_entries`
+ * yields a non-null entry struct with a non-null key), while the declared output type is built
+ * from Spark's `dataType` widened by `Utils.declaredChildNullability`. Where the two disagree
+ * DataFusion rejects the batch at runtime ("column types must match schema types"). Gated on any
+ * NullType-bearing input rather than the exact failing shapes; the same mismatch exists for
+ * non-Null nesting (`collect_list(array(array(id)))` fails on main) and is a native-aggregate
+ * bug.
+ */
+object CometCollectAggregate {
+
+ def nestedNullLevel(dt: DataType): Option[SupportLevel] =
+ if (SupportLevel.containsType(dt, classOf[NullType])) {
+ Some(
+ Unsupported(
+ Some("native collect_list/collect_set rebuilds a NullType-bearing element with a " +
+ "nullability that does not match the declared output type")))
+ } else {
+ None
+ }
+}
+
object CometCollectList extends CometAggregateExpressionSerde[CollectList] {
override def getSupportLevel(expr: CollectList): SupportLevel = {
@@ -1104,7 +1129,7 @@ object CometCollectList extends CometAggregateExpressionSerde[CollectList] {
if (!CometCollectShim.ignoreNulls(expr)) {
Unsupported(Some("collect_list with RESPECT NULLS (ignoreNulls = false) is not supported"))
} else {
- Compatible()
+ CometCollectAggregate.nestedNullLevel(expr.children.head.dataType).getOrElse(Compatible())
}
}
diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala
index cca9f63f8bf..4cd81eef331 100644
--- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala
@@ -22,7 +22,7 @@ package org.apache.comet.serde
import scala.annotation.tailrec
import scala.jdk.CollectionConverters._
-import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, BoundReference, Cast, CreateArray, ElementAt, EmptyRow, Expression, Flatten, GetArrayItem, IsNotNull, IsNull, LambdaFunction, Literal, NamedLambdaVariable, Reverse, Sequence, Size, Slice, SortArray, ZipWith}
+import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArrayRepeat, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, BoundReference, Cast, CreateArray, ElementAt, EmptyRow, Expression, Flatten, GetArrayItem, IsNotNull, IsNull, LambdaFunction, Literal, NamedLambdaVariable, Reverse, Sequence, Size, Slice, SortArray, ZipWith}
import org.apache.spark.sql.catalyst.util.GenericArrayData
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
@@ -53,6 +53,9 @@ object CometArrayRemove
object CometArrayAppend extends CometExpressionSerde[ArrayAppend] with ArraysBase {
+ override def getSupportLevel(expr: ArrayAppend): SupportLevel =
+ NullGuard.supportLevel(expr.children: _*)
+
override def convert(
expr: ArrayAppend,
inputs: Seq[Attribute],
@@ -214,13 +217,28 @@ object CometArrayIntersect
override def getIncompatibleReasons(): Seq[String] = Seq(incompatReason, collationReason)
+ private val nullElementReason: String =
+ "native array_intersect returns the other side's entries for a NullType-element array"
+
+ override def getUnsupportedReasons(): Seq[String] = Seq(nullElementReason)
+
override def getSupportLevel(expr: ArrayIntersect): SupportLevel = {
- // The native array_intersect dedups by raw bytes, which is wrong under non-default collations,
- // so report Incompatible rather than Unsupported: the JVM codegen dispatcher (Spark's own
- // doGenCode) performs collation-aware set membership and keeps execution native, matching
- // Spark. Only the output elements' collation metadata is dropped, consistent with CometReverse
- // and CometArrayJoin.
- if (hasNonDefaultStringCollation(expr.dataType)) {
+ // The native array_intersect dedups by raw bytes, which is wrong under non-default collations.
+ // That is Incompatible rather than Unsupported because there is something real to opt into: a
+ // user who does not need collation-aware set membership can take the native kernel. Under the
+ // default config the JVM codegen dispatcher (Spark's own doGenCode) runs it instead, keeping
+ // execution native and matching Spark; only the output elements' collation metadata is
+ // dropped, consistent with CometReverse and CometArrayJoin.
+ //
+ // A NullType-element side is Unsupported rather than Incompatible: the kernel's short-circuit
+ // is wrong for an intersection at every setting, so there is nothing to opt into. Unsupported
+ // reaches the same dispatcher, and unlike Incompatible it does so whatever `allowIncompatible`
+ // says. Reporting Incompatible would make EXPLAIN advertise the opt-in and then punish it:
+ // that branch calls `convert` with no dispatcher fallback, so declining there drops the whole
+ // projection back to Spark exactly when the user asked for more native execution.
+ if (NullElementSetOp.hasNullElementSide(expr)) {
+ Unsupported(Some(nullElementReason))
+ } else if (hasNonDefaultStringCollation(expr.dataType)) {
Incompatible(Some(collationReason))
} else {
Incompatible(Some(incompatReason))
@@ -231,8 +249,16 @@ object CometArrayIntersect
expr: ArrayIntersect,
inputs: Seq[Attribute],
binding: Boolean): Option[ExprOuterClass.Expr] = {
- val leftArrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding)
- val rightArrayExprProto = exprToProtoInternal(expr.children(1), inputs, binding)
+ // Defensive: `getSupportLevel` reports Unsupported for this shape, which never calls
+ // `convert`. Kept so a future support-level change cannot silently reach the kernel's
+ // short-circuit.
+ if (NullElementSetOp.hasNullElementSide(expr)) {
+ withFallbackReason(expr, nullElementReason)
+ return None
+ }
+ val Seq(left, right) = SetOpArguments.unify(expr.children)
+ val leftArrayExprProto = exprToProtoInternal(left, inputs, binding)
+ val rightArrayExprProto = exprToProtoInternal(right, inputs, binding)
val arraysIntersectScalarExpr =
scalarFunctionExprToProto("array_intersect", leftArrayExprProto, rightArrayExprProto)
@@ -304,18 +330,26 @@ object CometArrayExcept
private val incompatReason = "Null handling and ordering may differ from Spark"
+ private val elementTypeReason =
+ "native array_except supports only boolean, integral, floating-point, decimal, date, " +
+ "timestamp and string elements, or arrays of those"
+
override def getIncompatibleReasons(): Seq[String] = Seq(incompatReason)
+ override def getUnsupportedReasons(): Seq[String] = Seq(elementTypeReason)
+
override def getSupportLevel(expr: ArrayExcept): SupportLevel = {
- // Surface the native element-type restriction in EXPLAIN. We report Incompatible (not
- // Unsupported) for these types so the JVM codegen dispatcher still evaluates them natively
- // under the default config; the convert-time guard below is only reached under
- // allowIncompatible=true, where the native array_except cannot handle them.
- val reason = expr.children.map(_.dataType).find(dt => !isTypeSupported(dt)) match {
- case Some(dt) => s"native array_except does not support element type $dt: $incompatReason"
- case None => incompatReason
+ // Surface the native element-type restriction in EXPLAIN. Unsupported rather than
+ // Incompatible for these types: the JVM codegen dispatcher evaluates them natively and does
+ // so whatever `allowIncompatible` says, whereas the Incompatible branch calls `convert` with
+ // no dispatcher fallback, so the guard below would drop the projection back to Spark for a
+ // user who opted in. Only the element-type restriction is Unsupported; the ordering and null
+ // handling differences below remain a genuine opt-in.
+ expr.children.map(_.dataType).find(dt => !isTypeSupported(dt)) match {
+ case Some(dt) =>
+ Unsupported(Some(s"$elementTypeReason ($dt)"))
+ case None => Incompatible(Some(incompatReason))
}
- Incompatible(Some(reason))
}
@tailrec
@@ -337,17 +371,18 @@ object CometArrayExcept
expr: ArrayExcept,
inputs: Seq[Attribute],
binding: Boolean): Option[ExprOuterClass.Expr] = {
- // Defensive: only reached under allowIncompatible=true (the default-config Incompatible path
- // routes through the codegen dispatcher before convert). Native array_except cannot handle
- // these element types, so decline and let Spark evaluate.
+ // Defensive: `getSupportLevel` reports Unsupported for these element types, which never
+ // calls `convert`. Kept so a future support-level change cannot reach a kernel that would
+ // raise on them.
expr.children.map(_.dataType).find(dt => !isTypeSupported(dt)) match {
case Some(dt) =>
withFallbackReason(expr, s"data type not supported: $dt")
return None
case None =>
}
- val leftArrayExprProto = exprToProtoInternal(expr.left, inputs, binding)
- val rightArrayExprProto = exprToProtoInternal(expr.right, inputs, binding)
+ val Seq(left, right) = SetOpArguments.unify(expr.children)
+ val leftArrayExprProto = exprToProtoInternal(left, inputs, binding)
+ val rightArrayExprProto = exprToProtoInternal(right, inputs, binding)
val arrayExceptScalarExpr =
scalarFunctionExprToProto("array_except", leftArrayExprProto, rightArrayExprProto)
@@ -494,13 +529,67 @@ object CometSlice extends CometExpressionSerde[Slice] {
}
}
-object CometArrayUnion extends CometExpressionSerde[ArrayUnion] {
+/**
+ * DataFusion's set-op kernel (`array_union` / `array_intersect` / `array_except`) asserts that
+ * both sides carry the identical element type, nested field nullability included, and native
+ * kernels widen nested nullability differently from what Spark planned (a lambda variable over a
+ * list arrives nullable, a `monotonically_increasing_id()` struct field does not). Casting both
+ * sides to a deeply-nullable element type only widens metadata and never changes values, so the
+ * kernel always sees one type; a primitive element type needs nothing, since only nested fields
+ * take part in the comparison.
+ */
+private[serde] object SetOpArguments {
+ def unify(children: Seq[Expression]): Seq[Expression] =
+ children.map { c =>
+ c.dataType match {
+ case ArrayType(et, _) if isComplexType(et) =>
+ val unified = ArrayType(deepNullable(et), containsNull = true)
+ if (c.dataType == unified) c else Cast(c, unified)
+ case _ => c
+ }
+ }
+}
+
+/**
+ * DataFusion's set-op kernel (`general_set_op`, shared by array_union and array_intersect)
+ * short-circuits when either side's element type is Null: it returns distinct(other side). For a
+ * union that drops the NULL entries the Null-typed list actually holds (Spark keeps one NULL);
+ * for an intersection it returns the other side's entries although nothing can be common to both
+ * (Spark returns an empty list). Only a bare `array` takes the branch; a nested Null
+ * (`array>`) goes through the row converter, which handles it.
+ */
+private[serde] object NullElementSetOp {
+ def hasNullElementSide(expr: Expression): Boolean =
+ expr.children.exists(_.dataType match {
+ case ArrayType(NullType, _) => true
+ case _ => false
+ })
+}
+
+object CometArrayUnion extends CometExpressionSerde[ArrayUnion] with CodegenDispatchFallback {
+
+ private val nullElementReason =
+ "native array_union drops the entries of a NullType-element array"
+
+ override def getUnsupportedReasons(): Seq[String] = Seq(nullElementReason)
+
+ // A NullType-element side runs through the JVM codegen dispatcher (`CodegenDispatchFallback`)
+ // rather than the kernel that drops its entries; see `NullElementSetOp`.
+ override def getSupportLevel(expr: ArrayUnion): SupportLevel = {
+ if (NullElementSetOp.hasNullElementSide(expr)) {
+ Unsupported(Some(nullElementReason))
+ } else {
+ Compatible()
+ }
+ }
+
override def convert(
expr: ArrayUnion,
inputs: Seq[Attribute],
binding: Boolean): Option[ExprOuterClass.Expr] = {
- val leftArrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding)
- val rightArrayExprProto = exprToProtoInternal(expr.children(1), inputs, binding)
+ val Seq(left, right) = SetOpArguments.unify(expr.children)
+ val leftArrayExprProto = exprToProtoInternal(left, inputs, binding)
+ val rightArrayExprProto = exprToProtoInternal(right, inputs, binding)
val arraysUnionScalarExpr =
scalarFunctionExprToProto("array_union", leftArrayExprProto, rightArrayExprProto)
@@ -508,7 +597,32 @@ object CometArrayUnion extends CometExpressionSerde[ArrayUnion] {
}
}
-object CometCreateArray extends CometExpressionSerde[CreateArray] with ArraysBase {
+object CometCreateArray
+ extends CometExpressionSerde[CreateArray]
+ with ArraysBase
+ with CodegenDispatchFallback {
+
+ private val nullTypeBatchReason = "native make_array builds a single row from a NullType batch"
+
+ override def getUnsupportedReasons(): Seq[String] = Seq(nullTypeBatchReason)
+
+ override def getSupportLevel(expr: CreateArray): SupportLevel = {
+ // DataFusion's make_array funnels an argument list that is entirely Null-typed into
+ // SingleRowListArrayBuilder, producing ONE list row regardless of the input row count. Spark
+ // coerces CreateArray's children to a common type, so a NullType child means every child is
+ // NullType and that branch is the one taken. A non-scalar NullType argument (e.g.
+ // `aggregate(arr, NULL, (acc, x) -> NULL)` admitted by the JVM codegen dispatcher) therefore
+ // fails the scalar-function row-count check for batches with more than one row. All-literal
+ // NULL arguments arrive as scalars and broadcast correctly, and empty `array()` never reaches
+ // make_array (`convert` emits a literal). `CodegenDispatchFallback` keeps the non-scalar
+ // case in the Comet pipeline through the JVM codegen dispatcher.
+ if (expr.children.exists(c => c.dataType == NullType && !c.foldable)) {
+ Unsupported(Some(nullTypeBatchReason))
+ } else {
+ Compatible()
+ }
+ }
+
override def convert(
expr: CreateArray,
inputs: Seq[Attribute],
@@ -548,6 +662,45 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] with ArraysBas
}
}
+object CometArrayRepeat extends CometExpressionSerde[ArrayRepeat] with CodegenDispatchFallback {
+
+ private val nonNullableItemReason =
+ "native array_repeat rebuilds a non-nullable list item as nullable"
+
+ override def getUnsupportedReasons(): Seq[String] = Seq(nonNullableItemReason)
+
+ override def getSupportLevel(expr: ArrayRepeat): SupportLevel = {
+ expr.left match {
+ // Native `make_array` always emits a nullable item, whatever Spark's containsNull says, so
+ // `array(c)` over a non-nullable `c` repeats natively.
+ case _: CreateArray => Compatible()
+ case left =>
+ left.dataType match {
+ // DataFusion's list repeat rebuilds the repeated list's item field as nullable, and
+ // fails when the input's item is not ("ListArray expected data type List(non-null
+ // Struct(..)) got List(Struct(..))"). A non-NullType item reaches native with Spark's
+ // containsNull (`map_entries` makes a non-nullable entry struct; the JVM codegen
+ // dispatcher declares its list items with Spark's flag), so containsNull = false is
+ // the plan-time proxy for that input. `CodegenDispatchFallback` runs those through
+ // the dispatcher instead. A NullType item is declared nullable on the FFI boundary
+ // (Utils.declaredChildNullability) and matches the rebuild, so plain array
+ // stays native.
+ case ArrayType(elementType, false) if elementType != NullType =>
+ Unsupported(Some(nonNullableItemReason))
+ case _ => Compatible()
+ }
+ }
+ }
+
+ override def convert(
+ expr: ArrayRepeat,
+ inputs: Seq[Attribute],
+ binding: Boolean): Option[ExprOuterClass.Expr] = {
+ val childExprs = expr.children.map(exprToProtoInternal(_, inputs, binding))
+ scalarFunctionExprToProto("array_repeat", childExprs: _*)
+ }
+}
+
object CometGetArrayItem extends CometExpressionSerde[GetArrayItem] {
override def convert(
@@ -768,8 +921,15 @@ object CometSize extends CometExpressionSerde[Size] {
override def getSupportLevel(expr: Size): SupportLevel = {
expr.child.dataType match {
- case _: ArrayType => Compatible()
- case _: MapType => Compatible()
+ case _: ArrayType | _: MapType =>
+ // The null guard is only built for a nullable child in non-legacy mode; see `convert`.
+ // A non-nullable stateful child (`shuffle(arr)`, `filter(arr, x -> x < rand())`) is
+ // serialized once and evaluated once, so it needs no gate.
+ if (expr.legacySizeOfNull || !expr.child.nullable) {
+ Compatible()
+ } else {
+ NullGuard.supportLevel(expr.child)
+ }
case other =>
Unsupported(Some(s"Unsupported child data type: $other"))
}
@@ -780,21 +940,30 @@ object CometSize extends CometExpressionSerde[Size] {
inputs: Seq[Attribute],
binding: Boolean): Option[ExprOuterClass.Expr] = {
val arrayExprProto = exprToProtoInternal(expr.child, inputs, binding)
- for {
- isNotNullExprProto <- createIsNotNullExprProto(expr, inputs, binding)
- sizeScalarExprProto <- scalarFunctionExprToProto("size", arrayExprProto)
- emptyLiteralExprProto <- createLiteralExprProto(expr.legacySizeOfNull)
- } yield {
- val caseWhenExpr = ExprOuterClass.CaseWhen
- .newBuilder()
- .addWhen(isNotNullExprProto)
- .addThen(sizeScalarExprProto)
- .setElseExpr(emptyLiteralExprProto)
- .build()
- ExprOuterClass.Expr
- .newBuilder()
- .setCaseWhen(caseWhenExpr)
- .build()
+ // Native size already returns -1 for a null collection, which is the legacy answer; only
+ // the non-legacy NULL answer needs the null guard, and only when the child can be null.
+ if (expr.legacySizeOfNull || !expr.child.nullable) {
+ scalarFunctionExprToProto("size", arrayExprProto)
+ } else {
+ for {
+ isNotNullExprProto <- createIsNotNullExprProto(expr, inputs, binding)
+ sizeScalarExprProto <- scalarFunctionExprToProto("size", arrayExprProto)
+ nullLiteralExprProto <- exprToProtoInternal(
+ Literal(null, IntegerType),
+ Seq.empty,
+ binding = true)
+ } yield {
+ val caseWhenExpr = ExprOuterClass.CaseWhen
+ .newBuilder()
+ .addWhen(isNotNullExprProto)
+ .addThen(sizeScalarExprProto)
+ .setElseExpr(nullLiteralExprProto)
+ .build()
+ ExprOuterClass.Expr
+ .newBuilder()
+ .setCaseWhen(caseWhenExpr)
+ .build()
+ }
}
}
@@ -809,12 +978,6 @@ object CometSize extends CometExpressionSerde[Size] {
binding,
(builder, unaryExpr) => builder.setIsNotNull(unaryExpr))
}
-
- private def createLiteralExprProto(legacySizeOfNull: Boolean): Option[ExprOuterClass.Expr] = {
- val value = if (legacySizeOfNull) -1 else null
- exprToProtoInternal(Literal(value, IntegerType), Seq.empty, binding = true)
- }
-
}
object CometArrayPosition extends CometExpressionSerde[ArrayPosition] with ArraysBase {
@@ -846,7 +1009,8 @@ object CometArrayPosition extends CometExpressionSerde[ArrayPosition] with Array
object CometArraysZip extends CometExpressionSerde[ArraysZip] {
override def getUnsupportedReasons(): Seq[String] = Seq(
- "Not all input data types are supported; falls back to Spark for unsupported types")
+ "Not all input data types are supported; falls back to Spark for unsupported types",
+ NullGuard.reason)
private def isTypeSupported(dt: DataType): Boolean = {
import DataTypes._
@@ -868,7 +1032,10 @@ object CometArraysZip extends CometExpressionSerde[ArraysZip] {
return Unsupported(Some(s"Unsupported child data type: $dt"))
}
}
- Compatible()
+ // `convert` puts every child in the null-check predicate as well as in the zip itself;
+ // Spark's ArraysZip evaluates all of its children rather than short-circuiting, so a single
+ // stateful child is enough to diverge.
+ NullGuard.supportLevel(expr.children: _*)
}
override def convert(
diff --git a/spark/src/main/scala/org/apache/comet/serde/conditional.scala b/spark/src/main/scala/org/apache/comet/serde/conditional.scala
index df94be8c93c..7878e560757 100644
--- a/spark/src/main/scala/org/apache/comet/serde/conditional.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/conditional.scala
@@ -22,10 +22,36 @@ package org.apache.comet.serde
import scala.jdk.CollectionConverters._
import org.apache.spark.sql.catalyst.expressions.{Attribute, CaseWhen, Coalesce, Expression, If, IsNotNull}
+import org.apache.spark.sql.types.NullType
import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal
-object CometIf extends CometExpressionSerde[If] {
+/**
+ * Native CASE merges the rows each branch produced with Arrow's `merge_n`, which builds the
+ * result through a `MutableArrayData` that carries a validity bitmap; a `NullArray` cannot hold
+ * one ("Arrays of type Null cannot contain a null bitmask"), so a CASE whose result type is
+ * `NullType` fails whenever more than one branch contributes rows. Spark normally folds such an
+ * expression away (`IF(c, NULL, NULL)`), but a `NullType`-typed non-foldable branch keeps it. The
+ * three serdes below mix in `CodegenDispatchFallback`, so that shape runs through the JVM codegen
+ * dispatcher and the projection stays in the Comet pipeline.
+ */
+private[serde] object NullTypeBranches {
+ val reason = "native CASE cannot merge NullType branches"
+
+ def supportLevel(expr: Expression): SupportLevel =
+ if (expr.dataType == NullType) {
+ Unsupported(Some(reason))
+ } else {
+ Compatible()
+ }
+}
+
+object CometIf extends CometExpressionSerde[If] with CodegenDispatchFallback {
+
+ override def getUnsupportedReasons(): Seq[String] = Seq(NullTypeBranches.reason)
+
+ override def getSupportLevel(expr: If): SupportLevel = NullTypeBranches.supportLevel(expr)
+
override def convert(
expr: If,
inputs: Seq[Attribute],
@@ -49,7 +75,13 @@ object CometIf extends CometExpressionSerde[If] {
}
}
-object CometCaseWhen extends CometExpressionSerde[CaseWhen] {
+object CometCaseWhen extends CometExpressionSerde[CaseWhen] with CodegenDispatchFallback {
+
+ override def getUnsupportedReasons(): Seq[String] = Seq(NullTypeBranches.reason)
+
+ override def getSupportLevel(expr: CaseWhen): SupportLevel =
+ NullTypeBranches.supportLevel(expr)
+
override def convert(
expr: CaseWhen,
inputs: Seq[Attribute],
@@ -88,7 +120,19 @@ object CometCaseWhen extends CometExpressionSerde[CaseWhen] {
}
}
-object CometCoalesce extends CometExpressionSerde[Coalesce] {
+object CometCoalesce extends CometExpressionSerde[Coalesce] with CodegenDispatchFallback {
+
+ override def getUnsupportedReasons(): Seq[String] =
+ Seq(NullTypeBranches.reason, NullGuard.reason)
+
+ // Every child but the last is a guard; the last one is the ELSE, evaluated on the rows the
+ // guards left over. The result is a native CASE, so it shares that serde's NullType rule.
+ override def getSupportLevel(expr: Coalesce): SupportLevel =
+ NullTypeBranches.supportLevel(expr) match {
+ case _: Compatible => NullGuard.supportLevel(expr.children: _*)
+ case unsupported => unsupported
+ }
+
override def convert(
expr: Coalesce,
inputs: Seq[Attribute],
diff --git a/spark/src/main/scala/org/apache/comet/serde/hash.scala b/spark/src/main/scala/org/apache/comet/serde/hash.scala
index ee3e80059d5..9d955af560b 100644
--- a/spark/src/main/scala/org/apache/comet/serde/hash.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/hash.scala
@@ -20,7 +20,7 @@
package org.apache.comet.serde
import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Murmur3Hash, Sha1, Sha2, XxHash64}
-import org.apache.spark.sql.types.{ArrayType, DataType, DecimalType, IntegerType, LongType, MapType, StringType, StructType}
+import org.apache.spark.sql.types.{ArrayType, DataType, DecimalType, IntegerType, LongType, MapType, NullType, StringType, StructType}
import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, scalarFunctionExprToProtoWithReturnType, serializeDataType, supportedDataType}
@@ -119,9 +119,15 @@ private object HashUtils {
private val unsupportedDecimalReason =
"`DecimalType` with precision > 18 is not supported (Spark hashes via Java `BigDecimal`)"
private val unsupportedTimeTypeReason = "`TimeType` is not supported"
+ private val unsupportedNullTypeReason =
+ "`NullType` is not supported (the native hasher has no arm for it)"
val unsupportedReasons: Seq[String] =
- Seq(unsupportedDecimalReason, unsupportedTimeTypeReason, "Unsupported child data type")
+ Seq(
+ unsupportedDecimalReason,
+ unsupportedTimeTypeReason,
+ unsupportedNullTypeReason,
+ "Unsupported child data type")
def supportLevelForChildren(expr: Expression): SupportLevel = {
expr.children.iterator
@@ -135,6 +141,7 @@ private object HashUtils {
private def unsupportedReasonFor(dt: DataType): Option[String] = dt match {
case d: DecimalType if d.precision > 18 => Some(unsupportedDecimalReason)
+ case _: NullType => Some(unsupportedNullTypeReason)
case s: StructType =>
s.fields.iterator.flatMap(f => unsupportedReasonFor(f.dataType).iterator).toSeq.headOption
case a: ArrayType => unsupportedReasonFor(a.elementType)
diff --git a/spark/src/main/scala/org/apache/comet/serde/maps.scala b/spark/src/main/scala/org/apache/comet/serde/maps.scala
index 51fa428b543..4d089de8850 100644
--- a/spark/src/main/scala/org/apache/comet/serde/maps.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala
@@ -159,11 +159,21 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] {
override def getCompatibleNotes(): Seq[String] =
Seq(MapKeyDedupPolicySupport.nullKeyReason)
+ private val scalarSideReason: String =
+ "native map takes the first row of a scalar list where the other argument is per-row"
+
+ override def getUnsupportedReasons(): Seq[String] = Seq(scalarSideReason, NullGuard.reason)
+
override def getSupportLevel(expr: MapFromArrays): SupportLevel = {
if (MapKeyDedupPolicySupport.isLastWin) {
Incompatible(Some(MapKeyDedupPolicySupport.incompatibleReason))
+ } else if (expr.left.foldable != expr.right.foldable) {
+ // DataFusion's `map` reads a scalar list argument through its first row only, so a literal
+ // array beside a per-row one fails with "map requires key and value lists to have the same
+ // length" as soon as the batch holds more than one row.
+ Unsupported(Some(scalarSideReason))
} else {
- Compatible(None)
+ NullGuard.supportLevel(expr.left, expr.right)
}
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala
index 9ae92845c9c..2de985db5a7 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala
@@ -45,7 +45,7 @@ import org.apache.spark.util.io.ChunkedByteBuffer
import com.google.common.base.Objects
import org.apache.comet.{CometConf, ConfigEntry}
-import org.apache.comet.serde.OperatorOuterClass
+import org.apache.comet.serde.{Compatible, OperatorOuterClass, SupportLevel, Unsupported}
import org.apache.comet.serde.operator.CometSink
import org.apache.comet.shims.ShimCometBroadcastExchangeExec
@@ -264,6 +264,26 @@ object CometBroadcastExchangeExec extends CometSink[BroadcastExchangeExec] {
override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
CometConf.COMET_EXEC_BROADCAST_EXCHANGE_ENABLED)
+ /**
+ * A build side with a `NullType` directly under a struct or map entry stays on Spark's
+ * broadcast. `Utils.coalesceBroadcastBatches` cannot merge such batches (Arrow's appender
+ * cannot grow a `NullVector` under a struct) and would ship them uncoalesced, which costs every
+ * consuming task one IPC stream per broadcast buffer: measurably slower than Spark's broadcast
+ * once the build side spans many batches. See `Utils.hasNullTypeUnderStruct` for lifting this
+ * when Arrow is fixed.
+ */
+ override def getSupportLevel(b: BroadcastExchangeExec): SupportLevel = {
+ val uncoalescable = b.child.output.filter(a => Utils.hasNullTypeUnderStruct(a.dataType))
+ if (uncoalescable.isEmpty) {
+ Compatible(None)
+ } else {
+ Unsupported(
+ Some("NullType directly under a struct or map entry in the broadcast build side " +
+ s"(${uncoalescable.map(a => s"${a.name}: ${a.dataType.simpleString}").mkString(", ")}) " +
+ "cannot be coalesced for broadcast"))
+ }
+ }
+
override def createExec(
nativeOp: OperatorOuterClass.Operator,
b: BroadcastExchangeExec): CometNativeExec = {
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
index e2454f51322..18bff5bf856 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala
@@ -199,6 +199,10 @@ object CometArrowStream extends Logging {
* whose actual buffer carries validity bits must stay nullable even if Spark thought otherwise.
* Taking only `raw.isNullable` here would advertise non-nullable when the next batch does carry
* a null and crash native validation.
+ *
+ * Children come from the vector too, so a `NullType` map key arrives nullable (see
+ * `Utils.withNonNullableMapKeys`); `ArrowReader.getVectorSchemaRoot` would reject it when it
+ * rebuilds the `MapVector` from this schema.
*/
private def actualFieldOf(col: CometVector, expected: Field): Field = {
val raw = col match {
@@ -211,7 +215,7 @@ object CometArrowStream extends Logging {
val nullable = expected.isNullable || raw.isNullable
val fieldType =
new FieldType(nullable, raw.getType, raw.getDictionary, expected.getMetadata)
- new Field(expected.getName, fieldType, raw.getChildren)
+ Utils.withNonNullableMapKeys(new Field(expected.getName, fieldType, raw.getChildren))
}
/**
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
index d70fdab35e6..3d1217053af 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
@@ -21,7 +21,7 @@ package org.apache.spark.sql.comet.util
import java.io.{DataInputStream, DataOutputStream, File}
import java.nio.ByteBuffer
-import java.nio.channels.Channels
+import java.nio.channels.{Channels, WritableByteChannel}
import scala.jdk.CollectionConverters._
@@ -181,6 +181,18 @@ object Utils extends CometTypeShim with Logging {
s"Unsupported data type: [${dt.getClass.getName}] ${dt.catalogString}")
}
+ /**
+ * Nullability to declare for a nested child (array element, struct field, map value) of type
+ * `dataType`. A `NullType` child is always declared nullable, whatever Spark's flag says: every
+ * value is null, so a non-nullable flag is a contradiction, and native kernels that rebuild a
+ * list around the input's actual child compare that child's nullability with the one they
+ * assume and fail on the mismatch. Applied here and in `QueryPlanSerde.serializeDataType` so
+ * the JVM-exported field and the type declared to native agree. Map keys are not children in
+ * this sense: Arrow requires them non-nullable.
+ */
+ def declaredChildNullability(dataType: DataType, nullable: Boolean): Boolean =
+ nullable || dataType == NullType
+
/** Maps field from Spark to Arrow. NOTE: timeZoneId required for TimestampType */
def toArrowField(name: String, dt: DataType, nullable: Boolean, timeZoneId: String): Field = {
dt match {
@@ -189,7 +201,12 @@ object Utils extends CometTypeShim with Logging {
new Field(
name,
fieldType,
- Seq(toArrowField("element", elementType, containsNull, timeZoneId)).asJava)
+ Seq(
+ toArrowField(
+ "element",
+ elementType,
+ declaredChildNullability(elementType, containsNull),
+ timeZoneId)).asJava)
case StructType(fields) =>
val fieldType = new FieldType(nullable, ArrowType.Struct.INSTANCE, null)
new Field(
@@ -197,24 +214,30 @@ object Utils extends CometTypeShim with Logging {
fieldType,
fields
.map { field =>
- toArrowField(field.name, field.dataType, field.nullable, timeZoneId)
+ toArrowField(
+ field.name,
+ field.dataType,
+ declaredChildNullability(field.dataType, field.nullable),
+ timeZoneId)
}
.toSeq
.asJava)
case MapType(keyType, valueType, valueContainsNull) =>
val mapType = new FieldType(nullable, new ArrowType.Map(false), null)
- // Note: Map Type struct can not be null, Struct Type key field can not be null
- new Field(
- name,
- mapType,
+ // Note: Map Type struct can not be null, Struct Type key field can not be null (so the
+ // key is built here rather than through the StructType case and
+ // `declaredChildNullability`)
+ val entries = new Field(
+ MapVector.DATA_VECTOR_NAME,
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
Seq(
+ toArrowField(MapVector.KEY_NAME, keyType, nullable = false, timeZoneId),
toArrowField(
- MapVector.DATA_VECTOR_NAME,
- new StructType()
- .add(MapVector.KEY_NAME, keyType, nullable = false)
- .add(MapVector.VALUE_NAME, valueType, nullable = valueContainsNull),
- nullable = false,
+ MapVector.VALUE_NAME,
+ valueType,
+ declaredChildNullability(valueType, valueContainsNull),
timeZoneId)).asJava)
+ new Field(name, mapType, Seq(entries).asJava)
case dataType =>
val fieldType = new FieldType(nullable, toArrowType(dataType, timeZoneId), null)
new Field(name, fieldType, Seq.empty[Field].asJava)
@@ -230,6 +253,145 @@ object Utils extends CometTypeShim with Logging {
}.asJava)
}
+ /**
+ * Returns `field` with every map key field, at any nesting depth, marked non-nullable.
+ *
+ * Arrow requires map keys to be non-nullable and `MapVector.initializeChildrenFromFields`
+ * enforces it, so an IPC reader rejects a schema whose key field is nullable. Comet always
+ * declares keys non-nullable (`toArrowField`), but Arrow Java's `MinorType.NULL` factory builds
+ * a `NullVector` from the field name alone, so a `NullType` key (e.g. `map()`) reports a
+ * nullable field once a vector for it has been created by `Field.createVector`, i.e. after a
+ * native import, codegen dispatch output, or row-to-Arrow conversion. Apply this to the fields
+ * written to IPC.
+ */
+ def withNonNullableMapKeys(field: Field): Field = {
+ val children = field.getChildren.asScala.toSeq
+ if (children.isEmpty) {
+ return field
+ }
+ val newChildren = field.getType match {
+ case _: ArrowType.Map =>
+ children.map { entries =>
+ entries.getChildren.asScala.toSeq match {
+ case Seq(key, value) =>
+ val nonNullKey = withNonNullableMapKeys(key)
+ val repairedKey = if (nonNullKey.isNullable) {
+ new Field(
+ nonNullKey.getName,
+ new FieldType(
+ false,
+ nonNullKey.getType,
+ nonNullKey.getDictionary,
+ nonNullKey.getMetadata),
+ nonNullKey.getChildren)
+ } else {
+ nonNullKey
+ }
+ new Field(
+ entries.getName,
+ entries.getFieldType,
+ Seq(repairedKey, withNonNullableMapKeys(value)).asJava)
+ case _ => withNonNullableMapKeys(entries)
+ }
+ }
+ case _ => children.map(withNonNullableMapKeys)
+ }
+ if (newChildren == children) {
+ field
+ } else {
+ new Field(field.getName, field.getFieldType, newChildren.asJava)
+ }
+ }
+
+ /**
+ * Returns `root` unchanged when its declared schema already satisfies Arrow's non-nullable
+ * map-key invariant, otherwise a new root that shares the same vectors but advertises repaired
+ * fields.
+ *
+ * The check is against `root.getSchema`, not the live vectors' fields: a caller that built the
+ * root with already-repaired fields keeps its own root object. That matters when the row count
+ * is set after the writer is created (see `CometArrowPythonRunnerBase.startWriter`), because a
+ * replacement root copies the row count at construction and does not track the original.
+ */
+ private def withNonNullableMapKeys(root: VectorSchemaRoot): VectorSchemaRoot = {
+ val declared = root.getSchema.getFields.asScala.toSeq
+ val repaired = declared.map(withNonNullableMapKeys)
+ if (repaired == declared) {
+ root
+ } else {
+ new VectorSchemaRoot(repaired.asJava, root.getFieldVectors, root.getRowCount)
+ }
+ }
+
+ /**
+ * The only supported way to build an `ArrowStreamWriter` in Comet; enforced by the scalastyle
+ * `arrowstreamwriter` rule. Repairs the declared schema with [[withNonNullableMapKeys]], so a
+ * new IPC writer cannot reintroduce the nullable `NullType` map key by forgetting to ask.
+ *
+ * Returns the writer together with the root it is bound to, which is `root` itself unless the
+ * declared schema needed repairing. Callers must use the returned root: a writer serializes the
+ * root it was constructed with, so a row count set on a superseded root is not seen (it
+ * surfaces in the Python worker as "Array length did not match record batch length").
+ */
+ def newArrowStreamWriter(
+ root: VectorSchemaRoot,
+ provider: DictionaryProvider,
+ channel: WritableByteChannel): (VectorSchemaRoot, ArrowStreamWriter) = {
+ val bound = withNonNullableMapKeys(root)
+ // scalastyle:off arrowstreamwriter
+ (bound, new ArrowStreamWriter(bound, provider, channel))
+ // scalastyle:on arrowstreamwriter
+ }
+
+ /**
+ * Whether a `NullType` is a direct child of a struct (map keys and values included, as the
+ * entries struct's children) anywhere in `dataType`. This is the Spark-side twin of
+ * [[hasNullDirectlyUnderStruct]], for deciding at planning time what the broadcast coalescer
+ * would refuse: `CometBroadcastExchangeExec` keeps such build sides on Spark's broadcast, since
+ * shipping them uncoalesced costs every consuming task one IPC stream per broadcast buffer.
+ *
+ * When Comet moves to an Arrow Java release whose `VectorAppender` can grow a `NullVector`
+ * under a struct, drop this gate from `CometBroadcastExchangeExec.getSupportLevel` and the
+ * bypass in [[coalesceBroadcastBatches]]; `UtilsSuite` runs the real appender over every shape
+ * and will show which ones no longer need it.
+ */
+ def hasNullTypeUnderStruct(dataType: DataType): Boolean = {
+ // A list insulates whatever is below it, so `inStruct` resets when descending into one.
+ def check(dt: DataType, inStruct: Boolean): Boolean = dt match {
+ case NullType => inStruct
+ case ArrayType(element, _) => check(element, inStruct = false)
+ case StructType(fields) => fields.exists(f => check(f.dataType, inStruct = true))
+ case MapType(key, value, _) => check(key, inStruct = true) || check(value, inStruct = true)
+ case _ => false
+ }
+ check(dataType, inStruct = false)
+ }
+
+ /**
+ * Whether an Arrow `Null` field is a direct child of a struct (map entries included) anywhere
+ * in `schema`. Arrow's `VectorAppender` cannot grow such a column: a struct's capacity is the
+ * minimum over its *direct* children, a `NullVector`'s capacity equals its value count and its
+ * `reAlloc()` is a no-op, so the struct's capacity loop never terminates.
+ *
+ * A list breaks that chain from both sides, because `ListVector` overrides
+ * `BaseRepeatedValueVector.getValueCapacity` with one that only looks at its own offset and
+ * validity buffers. So a `Null` under a list appends fine (`array(NULL)`), and so does a list
+ * under a struct even when the list holds nulls (`map(k, array(NULL))`) - the struct only sees
+ * the list's own, growable capacity. Top-level `NullVector`s are fine too.
+ *
+ * The planner refuses these schemas before they reach a Comet broadcast (see
+ * [[hasNullTypeUnderStruct]]); this check keeps the coalescer safe should one arrive anyway.
+ */
+ private def hasNullDirectlyUnderStruct(schema: Schema): Boolean = {
+ def check(field: Field): Boolean = {
+ val isStruct = field.getType.isInstanceOf[ArrowType.Struct]
+ field.getChildren.asScala.exists { child =>
+ (isStruct && child.getType.isInstanceOf[ArrowType.Null]) || check(child)
+ }
+ }
+ schema.getFields.asScala.exists(check)
+ }
+
/**
* Build a `StructType` from a sequence of Spark `Attribute`s. Avoids
* `StructType.fromAttributes` (removed in Spark 4) and `DataTypeUtils.fromAttributes` (only on
@@ -264,10 +426,10 @@ object Utils extends CometTypeShim with Logging {
}
val provider = batchProviderOpt.getOrElse(dictionaryProvider)
- val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out))
+ val (writeRoot, writer) = newArrowStreamWriter(root, provider, Channels.newChannel(out))
writer.start()
writer.writeBatch()
- root.clear()
+ writeRoot.clear()
writer.close()
if (out.size() > 0) {
@@ -300,10 +462,10 @@ object Utils extends CometTypeShim with Logging {
val out = new DataOutputStream(codec.compressedOutputStream(cbbos))
val root = new VectorSchemaRoot(Seq(fieldVector).asJava)
- val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out))
+ val (writeRoot, writer) = newArrowStreamWriter(root, provider, Channels.newChannel(out))
writer.start()
writer.writeBatch()
- root.clear()
+ writeRoot.clear()
writer.close()
cbbos.toChunkedByteBuffer
@@ -388,19 +550,23 @@ object Utils extends CometTypeShim with Logging {
val reader =
new ArrowStreamReader(Channels.newChannel(compressedInputStream), allocator)
try {
- // Comet decodes dictionaries during execution, so this shouldn't happen.
- // If it does, fall back to the original uncoalesced buffers because each
- // partition can have a different dictionary, and appending index vectors
- // would silently mix indices from incompatible dictionaries.
- if (!reader.getDictionaryVectors.isEmpty) {
- logWarning(
- "Unexpected dictionary-encoded column during BroadcastExchange coalescing; " +
- "skipping coalesce")
- reader.close()
- if (targetRoot != null) {
- targetRoot.close()
- targetRoot = null
+ // Schemas that cannot be appended fall back to the original uncoalesced buffers:
+ // - Comet decodes dictionaries during execution, so a dictionary-encoded column
+ // shouldn't happen. If it does, each partition can have a different dictionary,
+ // and appending index vectors would silently mix incompatible dictionaries.
+ // - `VectorSchemaRootAppender` cannot grow a `NullVector` directly under a struct
+ // (see `hasNullDirectlyUnderStruct`).
+ val skipReason =
+ if (!reader.getDictionaryVectors.isEmpty) {
+ Some("unexpected dictionary-encoded column")
+ } else if (hasNullDirectlyUnderStruct(reader.getVectorSchemaRoot.getSchema)) {
+ Some("NullType directly under a struct or map entry")
+ } else {
+ None
}
+ if (skipReason.isDefined) {
+ logWarning(
+ s"${skipReason.get} during BroadcastExchange coalescing; skipping coalesce")
return (buffers, 0L, 0L)
}
while (reader.loadNextBatch()) {
@@ -447,8 +613,8 @@ object Utils extends CometTypeShim with Logging {
val outputStream = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate)
val compressedOutputStream =
new DataOutputStream(codec.compressedOutputStream(outputStream))
- val writer =
- new ArrowStreamWriter(targetRoot, null, Channels.newChannel(compressedOutputStream))
+ val (_, writer) =
+ newArrowStreamWriter(targetRoot, null, Channels.newChannel(compressedOutputStream))
try {
writer.start()
writer.writeBatch()
diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
index 651a9f25bc7..5161c03389c 100644
--- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
+++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
@@ -151,8 +151,15 @@ private[python] trait CometArrowPythonRunnerBase
new FieldType(false, ArrowType.Struct.INSTANCE, null),
childFields.asJava)
val structVec = structField.createVector(allocator).asInstanceOf[StructVector]
- writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
- arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ // Declare the root's schema from `structField`, not from the live vector, whose NullType
+ // map key reports nullable (see `Utils.withNonNullableMapKeys`); `newArrowStreamWriter`
+ // returns the root the writer is bound to, which is the one to close.
+ val declaredRoot =
+ new VectorSchemaRoot(Seq(structField).asJava, Seq[FieldVector](structVec).asJava, 0)
+ val (boundRoot, writer) =
+ Utils.newArrowStreamWriter(declaredRoot, null, Channels.newChannel(dataOut))
+ writeRoot = boundRoot
+ arrowWriter = writer
arrowWriter.start()
}
@@ -200,7 +207,9 @@ private[python] trait CometArrowPythonRunnerBase
// identical.
val childNames = inputStructType.fieldNames
streamFields = batchFields.zipWithIndex.map { case (field, i) =>
- renamed(field, childNames(i), forceNullable = true)
+ // `startWriter` rebuilds the struct from these fields, so repair the NullType map key
+ // nullability first (see `Utils.withNonNullableMapKeys`).
+ Utils.withNonNullableMapKeys(renamed(field, childNames(i), forceNullable = true))
}
startWriter(streamFields, dataOut)
}
diff --git a/spark/src/test/resources/pyspark/test_pyarrow_udf.py b/spark/src/test/resources/pyspark/test_pyarrow_udf.py
index 9a54300eed7..15efda928a5 100644
--- a/spark/src/test/resources/pyspark/test_pyarrow_udf.py
+++ b/spark/src/test/resources/pyspark/test_pyarrow_udf.py
@@ -1091,6 +1091,89 @@ def _normalize(row):
assert out == expected
+def test_map_in_arrow_null_typed_map_children(spark, tmp_path, accelerated):
+ """
+ `transform_values(map(), ...)` yields MapType(NullType, LongType) and `map(id, NULL)`
+ yields MapType(LongType, NullType); the accelerated runner repairs the NullType map key
+ (see CometArrowPythonRunnerBase) before building its destination struct.
+
+ pyarrow refuses to *declare* a non-nullable null-typed field (`A null type field may not
+ be non-nullable`), so the NullType-key map can only ever be a UDF input, never part of the
+ return schema; the UDF drops it and returns the NullType-value map unchanged.
+ """
+ src = str(tmp_path / "src.parquet")
+ spark.range(0, 20, 1, 2).write.parquet(src)
+ df = spark.read.parquet(src).selectExpr(
+ "id",
+ "transform_values(map(), (k, v) -> id) AS null_key",
+ "map(id, NULL) AS null_value",
+ )
+ schema_out = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("null_value", T.MapType(T.LongType(), T.NullType())),
+ ]
+ )
+
+ def drop_null_key(iterator):
+ for batch in iterator:
+ null_key_type = batch.schema.field("null_key").type
+ assert pa.types.is_map(null_key_type)
+ assert pa.types.is_null(null_key_type.key_type)
+ yield batch.select(["id", "null_value"])
+
+ result_df = df.mapInArrow(drop_null_key, schema_out)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ out = sorted((r["id"], r["null_value"]) for r in result_df.collect())
+ assert out == [(i, {i: None}) for i in range(20)]
+
+
+def test_map_in_arrow_null_typed_list_and_struct_children(spark, tmp_path, accelerated):
+ """
+ `transform(array(id), x -> NULL)` yields ArrayType(NullType) and `named_struct('a', id, 'b',
+ NULL)` yields a struct with a NullType field. Both reach the runner as Comet vectors with a
+ NullVector child (no buffers, every slot null), so the destination struct root and the IPC
+ schema must describe a null-typed child pyarrow can read. Unlike the NullType map key these
+ can also be returned: pyarrow accepts a *nullable* null-typed field, and Comet declares a
+ NullType child nullable whatever Spark's containsNull / field nullability says.
+ """
+ src = str(tmp_path / "src.parquet")
+ spark.range(0, 20, 1, 2).write.parquet(src)
+ df = spark.read.parquet(src).selectExpr(
+ "id",
+ "transform(array(id), x -> NULL) AS null_list",
+ "named_struct('a', id, 'b', NULL) AS null_struct",
+ )
+ schema_out = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("null_list", T.ArrayType(T.NullType())),
+ T.StructField(
+ "null_struct",
+ T.StructType(
+ [T.StructField("a", T.LongType()), T.StructField("b", T.NullType())]
+ ),
+ ),
+ ]
+ )
+
+ def check_and_pass(iterator):
+ for batch in iterator:
+ assert pa.types.is_null(batch.schema.field("null_list").type.value_type)
+ assert pa.types.is_null(batch.schema.field("null_struct").type.field("b").type)
+ yield batch
+
+ result_df = df.mapInArrow(check_and_pass, schema_out)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ out = sorted(
+ (r["id"], r["null_list"], (r["null_struct"]["a"], r["null_struct"]["b"]))
+ for r in result_df.collect()
+ )
+ assert out == [(i, [None], (i, None)) for i in range(20)]
+
+
def test_map_in_arrow_deeply_nested(spark, tmp_path, accelerated):
"""
Exercises direct source-vector serialization at depth > 1, in every nesting combination:
diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql
index fa1ef356925..0fd189bc2a9 100644
--- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql
@@ -468,3 +468,9 @@ INSERT INTO cl_src_map VALUES
query spark_answer_only
SELECT size(collect_list(m)) FROM cl_src_map
+
+-- Native collect accumulates the input's own child array, so a NullType-bearing element keeps
+-- its producer's nested nullability while the declared output widens it; the gate keeps such
+-- inputs in Spark rather than failing the batch at runtime.
+query expect_fallback(native collect_list/collect_set rebuilds a NullType-bearing element)
+SELECT grp, sort_array(collect_list(named_struct('i', i, 'n', NULL))) FROM cl_src_int GROUP BY grp
diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql
index bf37b8a3b79..0e162ab422e 100644
--- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql
@@ -298,3 +298,7 @@ SELECT grp, sort_array(collect_set(DISTINCT i)) FROM cs_src_int GROUP BY grp ORD
query
SELECT grp, sort_array(collect_set(i))
FROM cs_src_int GROUP BY grp HAVING size(collect_set(i)) > 1 ORDER BY grp
+
+-- Same gate as collect_list: a NullType-bearing element stays in Spark.
+query expect_fallback(native collect_list/collect_set rebuilds a NullType-bearing element)
+SELECT grp, sort_array(collect_set(named_struct('i', i, 'n', NULL))) FROM cs_src_int GROUP BY grp
diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_except.sql b/spark/src/test/resources/sql-tests/expressions/array/array_except.sql
index afe79948e23..011d3699619 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/array_except.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/array_except.sql
@@ -96,3 +96,11 @@ INSERT INTO test_except_flt_negzero VALUES
query ignore(https://issues.apache.org/jira/browse/SPARK-54918)
SELECT a, b, array_except(a, b) FROM test_except_flt_negzero
+
+-- The set-op kernel asserts identical element types, nested nullability included, and the two
+-- sides can arrive with different nested nullability (a literal element is non-nullable, a lambda
+-- variable over a list is not). Both sides are cast to a deeply-nullable element type first.
+-- A nested list rather than a struct, because `isTypeSupported` declines struct elements before
+-- `convert` runs.
+query
+SELECT array_except(transform(a, x -> array(1)), transform(b, x -> array(x))) FROM test_array_except
diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_intersect.sql b/spark/src/test/resources/sql-tests/expressions/array/array_intersect.sql
index fd5ac89a11e..b93c33e1dc8 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/array_intersect.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/array_intersect.sql
@@ -77,6 +77,11 @@ query
SELECT a, array_intersect(a, a) FROM test_intersect_dups
-- empty array combinations
+-- The empty operands are Null-typed literals. The NullType-element gate reports Unsupported for
+-- them, so the JVM codegen dispatcher evaluates them instead: still inside Comet, and matching
+-- Spark. A plain `query` pins both halves at once -- delete the gate and this file's
+-- allowIncompatible=true hands these to the native kernel, whose NullType short-circuit returns
+-- the other side's entries.
query
SELECT array_intersect(array(), array()), array_intersect(array(), array(1, 2)), array_intersect(array(1, 2), array())
@@ -257,3 +262,18 @@ SELECT array_intersect(array(1, NULL, 3), b) FROM test_array_intersect
-- conditional (CASE WHEN) arrays
query
SELECT array_intersect(CASE WHEN a IS NOT NULL THEN a ELSE array(0) END, b) FROM test_array_intersect
+
+-- The set-op kernel asserts identical element types, nested nullability included, and the two
+-- sides can arrive with different nested nullability (a literal field is non-nullable, a lambda
+-- variable over a list is not). Both sides are cast to a deeply-nullable element type first.
+query
+SELECT array_intersect(transform(a, x -> named_struct('i', 1)), transform(b, x -> named_struct('i', x))) FROM test_array_intersect
+
+-- The set-op kernel short-circuits on a NullType-element side and returns the other side's
+-- distinct entries, so `array_intersect(array(), array(NULL))` would come back as [NULL] where
+-- Spark returns []. Only a side that Spark leaves as array reaches it (a typed sibling
+-- makes Spark cast the Null side first). The gate reports Unsupported, which routes these to the
+-- JVM codegen dispatcher at every setting -- including this file's allowIncompatible=true, where
+-- dropping the gate would hand them to the kernel and turn the first case into [NULL].
+query
+SELECT array_intersect(array(), array(NULL)), array_intersect(array(NULL), array()), array_intersect(array(NULL, NULL), array(NULL))
diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql b/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql
index 58272d164f1..7944a471990 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql
@@ -118,3 +118,27 @@ SELECT array_repeat(CAST(NULL AS INT), cnt) FROM test_array_repeat
query
SELECT array_repeat(CAST(NULL AS STRING), cnt) FROM test_array_repeat
+
+-- A NullType element built by the JVM codegen dispatcher (containsNull=false in Spark) reaches
+-- native declared nullable; the kernel's result must still match the planned type
+query
+SELECT array_repeat(filter(array(), x -> true), 2) FROM test_array_repeat
+
+-- map_entries produces a list whose non-NullType item is declared non-nullable; native
+-- array_repeat rebuilds that item as nullable, so the serde declines it and the JVM codegen
+-- dispatcher runs it inside the Comet pipeline (the plan stays fully native).
+query
+SELECT array_repeat(map_entries(map(coalesce(long_v, 0), NULL)), 2) FROM test_array_repeat
+
+-- Same mismatch with no NullType anywhere, so the guard is not a NullType-specific workaround:
+-- any list whose item Spark declares non-nullable hits it. This one fails on main as well.
+query
+SELECT array_repeat(map_entries(map(coalesce(long_v, 0), long_v)), 2) FROM test_array_repeat
+
+-- array(c) over a non-nullable c has containsNull = false too, but native make_array emits a
+-- nullable item whatever Spark says, so the serde keeps this one on the native kernel.
+query
+SELECT array_repeat(array(coalesce(int_v, 0)), 2) FROM test_array_repeat
+
+query
+SELECT array_repeat(array(coalesce(int_v, 0), coalesce(cnt, 0)), cnt) FROM test_array_repeat
diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_sort_comparator.sql b/spark/src/test/resources/sql-tests/expressions/array/array_sort_comparator.sql
index 67905bb74ef..863bcf643fa 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/array_sort_comparator.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/array_sort_comparator.sql
@@ -42,3 +42,7 @@ FROM test_array_sort
-- all literals
query
SELECT array_sort(array(3, 1, 2), (l, r) -> CASE WHEN l < r THEN 1 WHEN l > r THEN -1 ELSE 0 END)
+
+-- an untyped empty array leaves the element type as NullType
+query
+SELECT array_sort(array(), (l, r) -> 0)
diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_union.sql b/spark/src/test/resources/sql-tests/expressions/array/array_union.sql
index 1f3a9dc9dde..2d7815328b1 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/array_union.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/array_union.sql
@@ -47,6 +47,8 @@ query
SELECT a, b, array_union(a, b) FROM test_union_nulls
-- empty array combinations
+-- Both sides are Null-typed empty arrays; the NullType-element gate hands them to the JVM
+-- codegen dispatcher, which keeps the projection in the Comet pipeline.
query
SELECT array_union(array(), array()) FROM test_union_nulls
@@ -246,3 +248,15 @@ SELECT array_union(array(NULL, 99), b) FROM test_array_union
-- conditional (CASE WHEN) arrays
query
SELECT array_union(CASE WHEN a IS NOT NULL THEN a ELSE array(0) END, b) FROM test_array_union
+
+-- DataFusion's set-op kernel treats a Null element type as "return distinct(other side)" and
+-- drops the NULL entries the Null-typed list actually holds, so NullType-element unions run
+-- through the JVM codegen dispatcher instead of the kernel.
+query
+SELECT array_union(transform(a, x -> NULL), array()) FROM test_array_union
+
+-- The set-op kernel asserts identical element types, nested nullability included, and the two
+-- sides can arrive with different nested nullability (a literal field is non-nullable, a lambda
+-- variable over a list is not). Both sides are cast to a deeply-nullable element type first.
+query
+SELECT array_union(transform(a, x -> named_struct('i', 1)), transform(b, x -> named_struct('i', x))) FROM test_array_union
diff --git a/spark/src/test/resources/sql-tests/expressions/array/arrays_zip.sql b/spark/src/test/resources/sql-tests/expressions/array/arrays_zip.sql
index e62e3842e8f..ae77d006b86 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/arrays_zip.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/arrays_zip.sql
@@ -144,6 +144,12 @@ select arrays_zip(a, b) FROM test_arrays_zip
query
SELECT arrays_zip(a, b)['a'] FROM (SELECT array(1, 2, 3) as a, array(3, 4, 5) as b)
+-- The serde guards every argument with CASE WHEN arg IS NOT NULL, and the zip inside the THEN
+-- branch only runs on the rows the guards selected. A stateful argument then advances its
+-- counter on those rows only, even when it is not the nullable one, so the pair stays in Spark.
+query expect_fallback(non-deterministic child under a null guard is evaluated on different rows than Spark's)
+SELECT arrays_zip(transform(a, x -> named_struct('i', monotonically_increasing_id(), 'n', NULL)), IF(size(b) > 1, b, CAST(NULL AS array))) FROM test_arrays_zip
+
query
SELECT arrays_zip(a, b)['b'] FROM (SELECT array(1, 2, 3) as a, array(3, 4, 5) as b)
diff --git a/spark/src/test/resources/sql-tests/expressions/array/create_array.sql b/spark/src/test/resources/sql-tests/expressions/array/create_array.sql
index 33e37b6765f..c9601adea3c 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/create_array.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/create_array.sql
@@ -140,3 +140,9 @@ SELECT array(named_struct('a', 1), named_struct('a', a)) FROM test_create_array
-- Empty array cast to a nested array type: the element type has to survive with no children.
query
SELECT CAST(array() AS ARRAY>)
+
+-- A non-foldable all-NullType argument (built by the JVM codegen dispatcher) would make native
+-- make_array collapse the whole batch into a single list row, so the whole expression runs
+-- through the dispatcher instead.
+query
+SELECT array(aggregate(arr, NULL, (acc, x) -> NULL)) FROM test_create_array_complex
diff --git a/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql b/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql
index 30ab1bba587..a291d275312 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql
@@ -113,3 +113,11 @@ SELECT id,
element_at(IF(monotonically_increasing_id() % 2 = 0, CAST(NULL AS ARRAY), array(1)), 1) AS v1,
element_at(IF(rand(7L) < 2, CAST(NULL AS ARRAY), array(1)), 1 + (id % (id - 2))) AS v2
FROM ansi_element_at_null
+
+-- non-deterministic collection under the ANSI null guard, built by a lambda
+-- The same restriction as above reached through the JVM codegen dispatcher: the guard's two
+-- copies of the dispatched `transform` share one kernel and its state, so it stays in Spark.
+-- ============================================================================
+
+query expect_fallback(nullable nondeterministic array or map operand)
+SELECT element_at(transform(IF(monotonically_increasing_id() % 2 = 0, arr, CAST(NULL AS array)), x -> x + 1), 1) FROM ansi_element_at_oob
diff --git a/spark/src/test/resources/sql-tests/expressions/array/slice.sql b/spark/src/test/resources/sql-tests/expressions/array/slice.sql
index 03a4bec0480..c10ef5e97a6 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/slice.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/slice.sql
@@ -277,3 +277,14 @@ INSERT INTO test_slice_map VALUES (1, 10), (2, NULL), (3, 30)
query
SELECT slice(array(map(k, v), map(k + 1, v)), 1, 1) FROM test_slice_map
+
+-- A NullType element built by the JVM codegen dispatcher (containsNull=false in Spark) reaches
+-- native declared nullable; the kernel's result must still match the planned type
+query
+SELECT slice(filter(array(), x -> true), 1, 1) FROM test_slice
+
+-- map_entries produces a list whose NullType-bearing struct item is declared non-nullable;
+-- native spark_array_slice keeps that item, and since the serde no longer serializes a return
+-- type (the output field comes from the input's), the two agree and it runs natively.
+query
+SELECT slice(map_entries(map(coalesce(start_idx, 0), NULL)), 1, 1) FROM test_slice
diff --git a/spark/src/test/resources/sql-tests/expressions/array/transform.sql b/spark/src/test/resources/sql-tests/expressions/array/transform.sql
index 46490917619..340cf23dc1e 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/transform.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/transform.sql
@@ -80,3 +80,10 @@ SELECT a FROM test_transform WHERE transform(a, x -> array(x)) = nested
-- all literals (constant folding is disabled by the test harness)
query
SELECT transform(array(1, 2, 3), x -> x * x)
+
+-- an untyped empty array leaves the element type as NullType
+query
+SELECT transform(array(), x -> x)
+
+query
+SELECT transform(array(NULL), x -> x)
diff --git a/spark/src/test/resources/sql-tests/expressions/array/zip_with.sql b/spark/src/test/resources/sql-tests/expressions/array/zip_with.sql
index 9345dbdfcde..6a8943a034e 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/zip_with.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/zip_with.sql
@@ -43,3 +43,7 @@ SELECT zip_with(a, b, (x, y) -> struct(x AS l, y AS r)) FROM test_zip_with
-- all literals
query
SELECT zip_with(array(1, 2), array(3, 4), (x, y) -> x * y)
+
+-- untyped empty arrays leave the element type as NullType
+query
+SELECT zip_with(array(), array(), (x, y) -> x)
diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql b/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql
index 508571d728d..b9817746e40 100644
--- a/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql
+++ b/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql
@@ -33,3 +33,9 @@ SELECT CASE WHEN s IS NULL THEN 'null_val' ELSE s END FROM test_case_when
-- literal arguments
query
SELECT CASE WHEN i = 1 THEN s WHEN i = 2 THEN 'fixed' ELSE s END FROM test_case_when
+
+-- A NullType result cannot run natively (native CASE merges the rows of its branches through
+-- Arrow's merge_n, which cannot build a NullArray with a validity bitmap), so the JVM codegen
+-- dispatcher runs it inside the Comet pipeline.
+query
+SELECT CASE WHEN i = 1 THEN aggregate(array(i), NULL, (acc, x) -> NULL) END FROM test_case_when
diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql b/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql
index 6457768aedf..003d69e12cd 100644
--- a/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql
+++ b/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql
@@ -33,3 +33,16 @@ SELECT coalesce(a, 99) FROM test_coalesce
-- literal arguments
query
SELECT coalesce(NULL, NULL, 99), coalesce(1, NULL, 99), coalesce(NULL)
+
+-- The serde guards every argument but the last with CASE WHEN arg IS NOT NULL THEN arg, and the
+-- two copies of a non-deterministic argument advance their state independently: the THEN copy
+-- can answer NULL for a row the guard selected, in a column declared non-nullable. The JVM
+-- codegen dispatcher evaluates Spark's own code once per row, so it runs there instead.
+query
+SELECT coalesce(IF(monotonically_increasing_id() % 2 = 0, a, NULL), b, 0) FROM test_coalesce
+
+-- A NullType result cannot run natively (the serde builds a native CASE, which merges the rows
+-- of its branches through Arrow's merge_n and cannot build a NullArray with a validity bitmap),
+-- so the JVM codegen dispatcher runs it inside the Comet pipeline.
+query
+SELECT coalesce(aggregate(array(a), NULL, (acc, x) -> NULL), aggregate(array(b), NULL, (acc, x) -> NULL)) FROM test_coalesce
diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql b/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql
index 53b908ca793..87d4748593b 100644
--- a/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql
+++ b/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql
@@ -30,3 +30,10 @@ SELECT IF(a > 0, 'positive', 'non-positive') FROM test_if
-- literal arguments
query
SELECT IF(true, 1, 2), IF(false, 1, 2), IF(NULL, 1, 2)
+
+-- A NullType result cannot run natively (native CASE merges the rows of its branches through
+-- Arrow's merge_n, which cannot build a NullArray with a validity bitmap), so the JVM codegen
+-- dispatcher runs it inside the Comet pipeline. Spark folds `IF(c, NULL, NULL)` itself, so the
+-- branch has to be a non-foldable NullType expression.
+query
+SELECT IF(cond, aggregate(array(a), NULL, (acc, x) -> NULL), NULL) FROM test_if
diff --git a/spark/src/test/resources/sql-tests/expressions/hash/hash.sql b/spark/src/test/resources/sql-tests/expressions/hash/hash.sql
index fc93e9d4bf6..651c4fc83fc 100644
--- a/spark/src/test/resources/sql-tests/expressions/hash/hash.sql
+++ b/spark/src/test/resources/sql-tests/expressions/hash/hash.sql
@@ -30,3 +30,7 @@ SELECT md5(col), md5(cast(a as string)), md5(cast(b as string)), hash(col), hash
-- native engine as scalar values rather than being folded away by Spark's optimizer.
query
SELECT md5('Spark SQL'), sha1('test'), sha2('test', 0), sha2('test', 256), sha2('test', 224), sha2('test', 384), sha2('test', 512), sha2('test', 128), sha2('test', -1), sha2(cast(null as string), 256), hash('test'), xxhash64('test')
+
+-- The native hasher has no arm for NullType, which a non-foldable struct or array can carry.
+query expect_fallback(`NullType` is not supported)
+SELECT hash(named_struct('a', a, 'b', NULL)), xxhash64(transform(array(a), x -> NULL)) FROM test
diff --git a/spark/src/test/resources/sql-tests/expressions/map/create_map.sql b/spark/src/test/resources/sql-tests/expressions/map/create_map.sql
index 8e479412453..40ed8a0f3b3 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/create_map.sql
+++ b/spark/src/test/resources/sql-tests/expressions/map/create_map.sql
@@ -39,3 +39,31 @@ SELECT map(1, named_struct('x', 10, 'x', 20))
-- Distinct field names are unaffected.
query
SELECT map(1, named_struct('x', 10, 'y', 20))
+
+-- ===== untyped constructors leave NullType children =====
+-- `map()` is MapType(NullType, NullType) and `map(k, NULL)` is MapType(_, NullType).
+
+query
+SELECT map()
+
+query
+SELECT map('a', NULL)
+
+query
+SELECT map(k, NULL) FROM test_create_map
+
+query
+SELECT size(map()), size(map('a', NULL))
+
+query
+SELECT map_keys(map()), map_values(map('a', NULL))
+
+query
+SELECT array(map()), struct(map(), map('a', NULL))
+
+query
+SELECT map('a', array(NULL)), map('a', map())
+
+-- carry the NullType children through a sort so the vector survives copy/spill paths
+query
+SELECT k, map(), map('a', NULL) FROM test_create_map ORDER BY k
diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_concat.sql b/spark/src/test/resources/sql-tests/expressions/map/map_concat.sql
index 453c76aa94e..130a3fd3594 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/map_concat.sql
+++ b/spark/src/test/resources/sql-tests/expressions/map/map_concat.sql
@@ -49,3 +49,7 @@ SELECT map_concat(map(1, 'x', 2, 'y'), map(3, 'z'))
-- a NULL literal map makes the whole result NULL
query
SELECT map_concat(map('a', 1), CAST(NULL AS map))
+
+-- untyped operands leave the result as MapType(NullType, NullType)
+query
+SELECT map_concat(map(), map())
diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql b/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql
index 2d172f436f8..f6b9ddcc01b 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql
+++ b/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql
@@ -63,3 +63,8 @@ SELECT array(
map_entries(element_at(map(1, map(1, true)), id))[0],
named_struct('key', 1, 'value', id IS NOT NULL)) AS a
FROM test_map_entries_nested
+
+-- A NullType map value built by the JVM codegen dispatcher (valueContainsNull=false in Spark)
+-- must reach DataFusion's map_entries declared nullable, or its ListArray build panics
+query
+SELECT map_entries(map_filter(map(), (k, v) -> true)) FROM test_map_entries
diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_filter.sql b/spark/src/test/resources/sql-tests/expressions/map/map_filter.sql
index edf04847107..1aaf3e4dbdf 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/map_filter.sql
+++ b/spark/src/test/resources/sql-tests/expressions/map/map_filter.sql
@@ -42,3 +42,7 @@ SELECT map_filter(m, (k, v) -> v > threshold) FROM test_map_filter
-- all literals
query
SELECT map_filter(map('a', 1, 'b', 2), (k, v) -> v > 1)
+
+-- an untyped value leaves the value type as NullType
+query
+SELECT map_filter(map('a', NULL), (k, v) -> true)
diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql
index 178c07f432a..202e9b562f4 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql
+++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql
@@ -58,4 +58,19 @@ query
SELECT map_from_arrays(array('a'), NULL)
query
-SELECT map_from_arrays(NULL, NULL)
\ No newline at end of file
+SELECT map_from_arrays(NULL, NULL)
+
+-- empty arrays produce MapType(NullType, NullType)
+query
+SELECT map_from_arrays(array(), array())
+
+-- The serde's null guard serializes both inputs twice, and a non-deterministic input would
+-- advance differently in each copy, so it falls back.
+query expect_fallback(non-deterministic child under a null guard is evaluated on different rows than Spark's)
+SELECT map_from_arrays(IF(monotonically_increasing_id() % 2 = 0, k, NULL), v) FROM test_map_from_arrays
+
+-- A literal array beside a per-row one: DataFusion's map kernel reads the scalar list through its
+-- first row only and fails the length check, so the shape stays in Spark. Independent of NullType
+-- (the typed literal fails the same way); found while probing the NullType flavour.
+query expect_fallback(native map takes the first row of a scalar list where the other argument is per-row)
+SELECT map_from_arrays(array(coalesce(k[0], 'x')), array(1)), map_from_arrays(array(coalesce(k[0], 'x')), array(NULL)) FROM test_map_from_arrays
diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_zip_with.sql b/spark/src/test/resources/sql-tests/expressions/map/map_zip_with.sql
index 181e6a7c159..74dc324ec02 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/map_zip_with.sql
+++ b/spark/src/test/resources/sql-tests/expressions/map/map_zip_with.sql
@@ -39,3 +39,7 @@ SELECT map_zip_with(m, n, (k, v1, v2) -> struct(v1 AS left, v2 AS right)) FROM t
-- all literals
query
SELECT map_zip_with(map('a', 1), map('a', 2, 'b', 3), (k, v1, v2) -> coalesce(v1, 0) + coalesce(v2, 0))
+
+-- untyped operands leave both key and value types as NullType
+query
+SELECT map_zip_with(map(), map(), (k, v1, v2) -> v1)
diff --git a/spark/src/test/resources/sql-tests/expressions/map/transform_keys.sql b/spark/src/test/resources/sql-tests/expressions/map/transform_keys.sql
index cc91075a720..526de0e17e0 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/transform_keys.sql
+++ b/spark/src/test/resources/sql-tests/expressions/map/transform_keys.sql
@@ -38,3 +38,7 @@ SELECT transform_keys(m, (k, v) -> concat(k, suffix)) FROM test_transform_keys
-- all literals
query
SELECT transform_keys(map('a', 1, 'b', 2), (k, v) -> upper(k))
+
+-- an untyped value leaves the value type as NullType
+query
+SELECT transform_keys(map('a', NULL), (k, v) -> k)
diff --git a/spark/src/test/resources/sql-tests/expressions/map/transform_values.sql b/spark/src/test/resources/sql-tests/expressions/map/transform_values.sql
index 50e49d8959a..23efb5ed0e5 100644
--- a/spark/src/test/resources/sql-tests/expressions/map/transform_values.sql
+++ b/spark/src/test/resources/sql-tests/expressions/map/transform_values.sql
@@ -42,3 +42,7 @@ SELECT transform_values(m, (k, v) -> v + delta) FROM test_transform_values
-- all literals
query
SELECT transform_values(map('a', 1, 'b', 2), (k, v) -> v * 100)
+
+-- a lambda returning NULL leaves the value type as NullType
+query
+SELECT transform_values(map('a', NULL), (k, v) -> v)
diff --git a/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql b/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql
index a3ad834c3ff..ce1253a876a 100644
--- a/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql
+++ b/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql
@@ -33,3 +33,11 @@ SELECT named_struct('x', 1, 'y', 'hello', 'z', 3.14)
query
SELECT named_struct('x', a, 'y', 'fixed_val', 'z', c) FROM test_named_struct
+
+-- Spark keeps duplicate field names as distinct positional fields, but Arrow's StructVector
+-- keys children by name, so both the native serde and the codegen dispatcher must fall back.
+query expect_fallback(duplicate field names)
+SELECT named_struct('x', a, 'x', b) FROM test_named_struct
+
+query expect_fallback(duplicate struct field name)
+SELECT transform(array(a), v -> named_struct('x', v, 'x', v + 1)) FROM test_named_struct
diff --git a/spark/src/test/resources/sql-tests/expressions/struct/get_struct_field.sql b/spark/src/test/resources/sql-tests/expressions/struct/get_struct_field.sql
index c5d9e64a634..3afaf8b32bc 100644
--- a/spark/src/test/resources/sql-tests/expressions/struct/get_struct_field.sql
+++ b/spark/src/test/resources/sql-tests/expressions/struct/get_struct_field.sql
@@ -26,3 +26,9 @@ SELECT s.name, s.age, s.score FROM test_struct
query
SELECT s.name, s.age + 1, s.score * 2 FROM test_struct
+
+-- A struct row that a kernel grew through MutableArrayData (element_at on an out-of-range index)
+-- can carry a Null-typed child with a validity bitmap; projecting that child rebuilds a clean
+-- NullArray instead of failing validation. The typed sibling is the control.
+query
+SELECT try_element_at(array(named_struct('n', s.name, 'z', NULL)), CAST(s.age % 2 + 1 AS INT)).z, try_element_at(array(named_struct('n', s.name, 'z', NULL)), CAST(s.age % 2 + 1 AS INT)).n FROM test_struct
diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala
index ad86dd15bc1..c2b50644108 100644
--- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala
@@ -1099,6 +1099,33 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp
}
}
+ test("size - non-deterministic child under the null guard") {
+ withParquetTable((0 until 16).map(i => Tuple1(i.toLong)), "t", withDictionary = false) {
+ // Non-legacy size wraps a nullable child in `CASE WHEN child IS NOT NULL`, which would
+ // evaluate a stateful child twice, so that shape stays in Spark; legacy mode builds no
+ // guard and keeps it native.
+ val nullableStateful =
+ "SELECT _1, size(IF(monotonically_increasing_id() % 2 = 0, array(_1), NULL)) FROM t"
+ withSQLConf(SQLConf.LEGACY_SIZE_OF_NULL.key -> "false") {
+ checkSparkAnswerAndFallbackReason(
+ nullableStateful,
+ "non-deterministic child under a null guard is evaluated on different rows than Spark's")
+ }
+ withSQLConf(
+ SQLConf.LEGACY_SIZE_OF_NULL.key -> "true",
+ SQLConf.ANSI_ENABLED.key -> "false") {
+ checkSparkAnswerAndOperator(nullableStateful)
+ }
+ // A non-nullable child gets no guard, so a stateful one whose length depends on the
+ // counter is evaluated once and matches Spark. The lambda runs through the JVM codegen
+ // dispatcher, where a guard's two copies would share one kernel and its counter.
+ withSQLConf(SQLConf.LEGACY_SIZE_OF_NULL.key -> "false") {
+ checkSparkAnswerAndOperator(
+ "SELECT _1, size(filter(array(_1, 1, 2), x -> x < monotonically_increasing_id())) FROM t")
+ }
+ }
+ }
+
// https://github.com/apache/datafusion-comet/issues/4560
test("array_size returns null for null input") {
val table = "t1"
diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
index 400ec6fcd2f..dce8e01d6bb 100644
--- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
@@ -19,18 +19,21 @@
package org.apache.comet
+import scala.jdk.CollectionConverters._
+
import org.scalatest.funsuite.AnyFunSuite
import org.apache.spark.SparkConf
import org.apache.spark.serializer.JavaSerializer
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Add, AddMonths, BoundReference, Cast, Coalesce, Concat, CreateArray, CreateMap, DateFormatClass, ElementAt, Expression, GetStructField, IntegralDivide, LeafExpression, Length, Literal, MakeTimestamp, MicrosToTimestamp, MillisToTimestamp, MonthsBetween, Nondeterministic, Rand, Size, Substring, ToUnixTimestamp, Unevaluable, UnixMicros, UnixMillis, UnixSeconds, Upper}
+import org.apache.spark.sql.catalyst.expressions.{Add, AddMonths, BoundReference, Cast, Coalesce, Concat, CreateArray, CreateMap, CreateNamedStruct, DateFormatClass, ElementAt, Expression, GetStructField, IntegralDivide, IsNull, LeafExpression, Length, Literal, MakeTimestamp, MicrosToTimestamp, MillisToTimestamp, MonthsBetween, Nondeterministic, Rand, Size, Substring, ToUnixTimestamp, Unevaluable, UnixMicros, UnixMillis, UnixSeconds, Upper}
import org.apache.spark.sql.catalyst.expressions.codegen.{CodeFormatter, CodegenContext, CodegenFallback, ExprCode}
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String
import org.apache.comet.codegen.CometBatchKernelCodegen
import org.apache.comet.codegen.CometBatchKernelCodegen.{ArrayColumnSpec, ArrowColumnSpec, MapColumnSpec, ScalarColumnSpec, StructColumnSpec, StructFieldSpec}
+import org.apache.comet.serde.QueryPlanSerde
import org.apache.comet.udf.codegen.CometScalaUDFCodegen
// Resolve Arrow vector classes through the codegen object so tests see the same `Class` objects
@@ -287,6 +290,203 @@ class CometCodegenSourceSuite extends AnyFunSuite {
s"expected reason to name the rejected expression class; got: ${reason.get}")
}
+ test("canHandle accepts NullType outputs, including nested in complex types") {
+ // Spark leaves the children of untyped constructors as NullType: `map()` is
+ // MapType(NullType, NullType), `map('a', NULL)` is MapType(StringType, NullType) and
+ // `array()` is ArrayType(NullType). Rejecting those forced the whole projection back to
+ // Spark even though a NullType output only ever has to write nulls.
+ Seq(
+ Literal(null, NullType),
+ Literal.create(Map.empty[Any, Any], MapType(NullType, NullType, valueContainsNull = false)),
+ Literal.create(Map("a" -> null), MapType(StringType, NullType, valueContainsNull = true)),
+ Literal.create(Seq(null), ArrayType(NullType)),
+ Literal.create(
+ Seq(Map.empty[Any, Any]),
+ ArrayType(MapType(NullType, NullType, valueContainsNull = false)))).foreach { expr =>
+ assert(
+ CometBatchKernelCodegen.canHandle(expr).isEmpty,
+ s"expected canHandle to accept output type ${expr.dataType}")
+ }
+ }
+
+ test("canHandle rejects NullType inputs, including nested in complex types") {
+ // `CometScalaUDFCodegen.specFor` has no ArrowColumnSpec for a NullVector, so a NullType input
+ // must keep falling back at plan time rather than throwing once the plan has committed to a
+ // kernel.
+ Seq(
+ NullType,
+ ArrayType(NullType),
+ MapType(StringType, NullType),
+ StructType(Seq(StructField("f", NullType)))).foreach { dt =>
+ val reason = CometBatchKernelCodegen.canHandle(BoundReference(0, dt, nullable = true))
+ assert(reason.isDefined, s"expected canHandle to reject input type $dt")
+ assert(
+ reason.get.contains("unsupported"),
+ s"expected an unsupported-type reason for $dt; got: ${reason.get}")
+ }
+ }
+
+ test("canHandle rejects duplicate struct field names, in outputs and inputs") {
+ // See `CometBatchKernelCodegen.hasDuplicateStructFieldNames`; the gate must also catch a
+ // struct buried inside a larger expression dispatched as a whole (e.g. a `transform` body).
+ val x = BoundReference(0, IntegerType, nullable = false)
+ val dupNull = CreateNamedStruct(Seq(Literal("a"), x, Literal("a"), Literal(null, NullType)))
+ val dupInt = CreateNamedStruct(Seq(Literal("a"), x, Literal("a"), Add(x, Literal(1))))
+ val dupInput = IsNull(
+ BoundReference(
+ 0,
+ StructType(Seq(StructField("a", IntegerType), StructField("a", StringType))),
+ nullable = true))
+ Seq(dupNull, dupInt, CreateArray(Seq(dupInt)), dupInput).foreach { expr =>
+ val reason = CometBatchKernelCodegen.canHandle(expr)
+ assert(
+ reason.exists(_.contains("duplicate struct field name")),
+ s"expected canHandle to reject $expr; got: $reason")
+ }
+ val distinct =
+ CreateNamedStruct(Seq(Literal("a"), x, Literal("b"), Literal(null, NullType)))
+ assert(CometBatchKernelCodegen.canHandle(distinct).isEmpty)
+ }
+
+ test("NullType children are declared nullable whatever the Spark flags say") {
+ // Spark leaves containsNull / valueContainsNull / field nullability false on some untyped
+ // shapes (`filter(array(), ...)`, `map_filter(map(), ...)`), but a non-nullable Null child is a
+ // contradiction, and native kernels that rebuild a list around the input's actual child
+ // (map_entries, array_repeat, spark_array_slice) fail on the nullability mismatch. Non-Null
+ // children keep their flags; map keys stay non-nullable.
+ val exact = ArrayType(
+ StructType(
+ Seq(
+ StructField("n", NullType, nullable = false),
+ StructField("i", IntegerType, nullable = false),
+ StructField(
+ "m",
+ MapType(NullType, NullType, valueContainsNull = false),
+ nullable = false))),
+ containsNull = false)
+ val field = CometBatchKernelCodegen.toFfiArrowField("out", exact, nullable = false)
+ assert(!field.isNullable, "top-level nullability is the caller's")
+ val item = field.getChildren.get(0)
+ assert(!item.isNullable, "a non-Null element keeps containsNull = false")
+ val Seq(n, i, m) = item.getChildren.asScala.toSeq
+ assert(n.isNullable && !i.isNullable && !m.isNullable)
+ val Seq(key, value) = m.getChildren.get(0).getChildren.asScala.toSeq
+ assert(!key.isNullable, "map keys stay non-nullable")
+ assert(value.isNullable)
+
+ val nullElements = CometBatchKernelCodegen
+ .toFfiArrowField("out", ArrayType(NullType, containsNull = false), nullable = true)
+ assert(nullElements.getChildren.get(0).isNullable)
+
+ // The serde declares the same nullability to native.
+ val proto = QueryPlanSerde.serializeDataType(exact).get.getTypeInfo
+ assert(!proto.getList.getContainsNull)
+ val struct = proto.getList.getElementType.getTypeInfo.getStruct
+ assert(
+ struct.getFieldNullable(0) && !struct.getFieldNullable(1) && !struct.getFieldNullable(2))
+ assert(struct.getFieldDatatypes(2).getTypeInfo.getMap.getValueContainsNull)
+ assert(
+ QueryPlanSerde
+ .serializeDataType(ArrayType(NullType, containsNull = false))
+ .get
+ .getTypeInfo
+ .getList
+ .getContainsNull)
+ }
+
+ test("top-level struct output with a NullType field allocates") {
+ // `StructVector`'s constructor builds a `NullableStructWriter` over the field's children,
+ // which throws `Unknown type: NULL` for a Null child; the children must be added after
+ // construction. Reachable from the JVM codegen dispatcher through e.g.
+ // `coalesce(named_struct('i', monotonically_increasing_id(), 'n', NULL), ...)`.
+ val dataType = StructType(
+ Seq(
+ StructField("i", LongType, nullable = false),
+ StructField("n", NullType),
+ StructField("l", ArrayType(IntegerType))))
+ val field = CometBatchKernelCodegen.toFfiArrowField("out", dataType, nullable = true)
+ val vector = CometBatchKernelCodegen.allocateOutput(field, 4, 0)
+ try {
+ val struct = vector.asInstanceOf[org.apache.arrow.vector.complex.StructVector]
+ assert(struct.getChild("n").isInstanceOf[org.apache.arrow.vector.NullVector])
+ assert(struct.getChildrenFromFields.size == 3)
+ // The export names survive: the list child is labelled `item`, not Arrow Java's `$data$`.
+ assert(struct.getField == field)
+ assert(struct.getField.getChildren.get(2).getChildren.get(0).getName == "item")
+ } finally {
+ vector.close()
+ }
+ }
+
+ test("nested NullType output casts the child vector and writes setNull into it") {
+ // A scalar NullType output cannot distinguish `emitWrite`'s NullType branch from
+ // `defaultBody`'s own `ev.isNull -> output.setNull(i)` short-circuit, which emits the same
+ // text. A NullType *value* child inside a map is only reachable through `emitWrite`, so
+ // assert on that: the child vector must be cast to NullVector and written through `setNull`.
+ val src = CometBatchKernelCodegen
+ .generateSource(
+ Literal.create(Map("a" -> null), MapType(StringType, NullType, valueContainsNull = true)),
+ IndexedSeq(nullableString))
+ .body
+ val childCast =
+ """org\.apache\.arrow\.vector\.NullVector\s+(\w+)\s*=\s*\(org\.apache\.arrow\.vector\.NullVector\)""".r
+ val childVar = childCast
+ .findFirstMatchIn(src)
+ .map(_.group(1))
+ .getOrElse(fail(s"expected a NullVector child-vector cast; got:\n$src"))
+ assert(
+ src.contains(s"$childVar.setNull("),
+ s"expected a setNull write into the NullVector child `$childVar`; got:\n$src")
+ }
+
+ test("gate and output emitters agree across the whole accepted type surface") {
+ // If `canHandle` greenlights an output type, generating the kernel for it must not throw.
+ val leaves: Seq[DataType] = Seq(
+ NullType,
+ BooleanType,
+ ByteType,
+ ShortType,
+ IntegerType,
+ LongType,
+ FloatType,
+ DoubleType,
+ DecimalType(10, 2),
+ DecimalType(38, 18),
+ StringType,
+ BinaryType,
+ DateType,
+ TimestampType,
+ TimestampNTZType,
+ YearMonthIntervalType(),
+ DayTimeIntervalType(),
+ CalendarIntervalType)
+
+ val candidates: Seq[DataType] = leaves.flatMap { dt =>
+ Seq(
+ dt,
+ ArrayType(dt),
+ ArrayType(ArrayType(dt)),
+ MapType(StringType, dt),
+ StructType(Seq(StructField("f", dt))),
+ ArrayType(StructType(Seq(StructField("f", dt)))),
+ MapType(StringType, ArrayType(dt)))
+ }
+
+ val accepted =
+ candidates.filter(dt => CometBatchKernelCodegen.canHandle(Literal.create(null, dt)).isEmpty)
+ assert(
+ accepted.size > leaves.size,
+ s"expected the gate to accept nested shapes too; only got ${accepted.size}")
+
+ accepted.foreach { dt =>
+ withClue(s"canHandle accepted output type $dt, so the emitters must handle it: ") {
+ CometBatchKernelCodegen.generateSource(
+ Literal.create(null, dt),
+ IndexedSeq(nullableString))
+ }
+ }
+ }
+
test("CSE collapses a repeated subtree to one evaluation in the generated body") {
// `Add(Length(Upper(c0)), Length(Upper(c0)))` has `Length(Upper(c0))` as a common subtree.
// Length.doGenCode emits `$value.numChars()` on every Spark version the project targets,
diff --git a/spark/src/test/scala/org/apache/comet/CometCsvExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometCsvExpressionSuite.scala
index 6e50e7cb9ae..d3d69cf2e5d 100644
--- a/spark/src/test/scala/org/apache/comet/CometCsvExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCsvExpressionSuite.scala
@@ -23,12 +23,14 @@ import scala.jdk.CollectionConverters._
import scala.util.Random
import org.apache.hadoop.fs.Path
-import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.{CometTestBase, Row}
import org.apache.spark.sql.catalyst.expressions.StructsToCsv
+import org.apache.spark.sql.comet.CometProjectExec
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.StringType
+import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus
import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions}
class CometCsvExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
@@ -65,6 +67,28 @@ class CometCsvExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper
}
}
+ test("to_csv - a row that renders empty is NULL") {
+ // Spark hands the row to univocity's `writeRowToString` with `skipEmptyLines`, which turns an
+ // empty rendering (a lone null field under the default empty `nullValue`) into NULL rather
+ // than "". Spark 3.5+ crashes on that NULL in its own generated code
+ // (`nullSafeCodeGen` never marks the result null), so Spark's answer is only comparable on
+ // 3.4; the native answer is pinned on every version.
+ withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[StructsToCsv]) -> "true") {
+ val df = spark
+ .range(0, 4, 1, 1)
+ .select(
+ to_csv(struct(when(col("id") > 1, col("id")).as("c"))).as("lone"),
+ to_csv(struct(when(col("id") > 1, col("id")).as("c"), lit(null).as("n"))).as("pair"))
+ if (!isSpark35Plus) {
+ checkSparkAnswerAndOperator(df)
+ } else {
+ checkAnswer(df, Seq(Row(null, ","), Row(null, ","), Row("2", "2,"), Row("3", "3,")))
+ assert(collect(df.queryExecution.executedPlan) { case p: CometProjectExec => p }.nonEmpty)
+ checkSparkAnswerAndOperator(df.select(col("pair")))
+ }
+ }
+ }
+
test("to_csv - with configurable formatting options") {
val table = "t1"
withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[StructsToCsv]) -> "true") {
diff --git a/spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala b/spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala
new file mode 100644
index 00000000000..aaff7cee8d3
--- /dev/null
+++ b/spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala
@@ -0,0 +1,838 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet
+
+import scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
+import scala.util.{Failure, Success, Try}
+
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.catalyst.expressions.{ArrayExcept, ArrayIntersect, ArrayRepeat, ArraysZip, ArrayUnion, CaseWhen, Coalesce, CreateArray, Expression, If, JsonToStructs, MapFromArrays, RuntimeReplaceable, Sequence, StringToMap}
+import org.apache.spark.sql.catalyst.expressions.aggregate._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.comet.CometProjectExec
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.NullType
+
+import org.apache.comet.serde.{CodegenDispatchFallback, CometExpressionSerde, Compatible, QueryPlanSerde, SupportLevel, Unsupported}
+
+/**
+ * Cross-product sweep of the `NullType` shapes the JVM codegen dispatcher admits against the
+ * expressions that consume them.
+ *
+ * Three sweeps share one driver, `sweep`: producers under consumers, producers under operators,
+ * and producers nested in a container and then put under a serializing operator.
+ */
+class CometNullTypeCompositionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
+
+ /**
+ * Non-foldable `NullType`-bearing expressions over a live column. Each is a shape that the
+ * widened gate lets reach native execution; the constant-folded forms are literals and take a
+ * different path, so every producer here references `id`.
+ */
+ private val arrayOfNull = Seq(
+ "transform(array(id), x -> NULL)",
+ "filter(array(CAST(NULL AS int)), x -> id IS NOT NULL)")
+
+ private val mapWithNullValue = Seq("map(id, NULL)")
+
+ private val mapWithNullKey = Seq("transform_values(map(), (k, v) -> id)")
+
+ private val arrayOfStructWithNull = Seq("map_entries(map(id, NULL))")
+
+ private val structWithNull = Seq("named_struct('a', id, 'b', NULL)")
+
+ private val scalarNull = Seq("aggregate(array(id), NULL, (acc, x) -> NULL)")
+
+ /** Consumers valid for a value of any type, written as templates over the producer `%s`. */
+ private val anyTypeConsumers = Seq(
+ "to_json(struct(%s AS c))",
+ "to_csv(struct(%s AS c))",
+ "hash(%s)",
+ "xxhash64(%s)",
+ "CAST(%s AS string)",
+ "CASE WHEN id > 2 THEN %s END",
+ "IF(id > 2, %s, NULL)",
+ "coalesce(%s, %s)",
+ "%s IS NULL",
+ "%s = %s",
+ "%s <=> %s")
+
+ /** Consumers valid for any `array`. */
+ private val arrayConsumers = Seq(
+ "size(%s)",
+ "reverse(%s)",
+ "array_distinct(%s)",
+ "sort_array(%s)",
+ "array_sort(%s)",
+ "element_at(%s, 1)",
+ "slice(%s, 1, 1)",
+ "array_repeat(%s, 2)",
+ "array_union(%s, %s)",
+ "array_except(%s, %s)",
+ "array_intersect(%s, %s)",
+ // Set ops only diverge when a Null-typed side meets a side that actually holds entries, so
+ // the same-producer templates above cannot reach that branch on their own.
+ "array_union(%s, array(1))",
+ "array_union(array(1), %s)",
+ "array_except(%s, array(1))",
+ "array_intersect(%s, array(1))",
+ "arrays_overlap(%s, %s)",
+ "arrays_zip(%s, %s)",
+ "concat(%s, %s)",
+ "flatten(array(%s))",
+ "array_position(%s, NULL)",
+ "array_contains(%s, NULL)",
+ // Spark rejects a NullType needle, so the typed needle is what reaches the serde.
+ "array_contains(%s, %s[0])",
+ "array_remove(%s, NULL)",
+ "array_append(%s, NULL)",
+ "array_insert(%s, 1, NULL)",
+ "exists(%s, x -> x IS NULL)",
+ "forall(%s, x -> x IS NULL)",
+ "filter(%s, x -> x IS NULL)",
+ "transform(%s, x -> x)",
+ "aggregate(%s, 0, (acc, x) -> acc)",
+ "zip_with(%s, %s, (x, y) -> x)",
+ "map_from_entries(arrays_zip(%s, %s))",
+ "%s[0]",
+ "array_compact(%s)",
+ "array_max(%s)",
+ "array_min(%s)",
+ "array_join(%s, ',')",
+ // shuffle is non-deterministic, so only the sorted result compares.
+ "sort_array(shuffle(%s))",
+ "map_from_arrays(%s, array(id))") ++ anyTypeConsumers
+
+ /** Consumers valid for any `map`. */
+ private val mapConsumers = Seq(
+ "size(%s)",
+ "map_keys(%s)",
+ "map_values(%s)",
+ "map_entries(%s)",
+ "map_concat(%s, %s)",
+ "map_filter(%s, (k, v) -> k IS NOT NULL)",
+ "transform_values(%s, (k, v) -> v)",
+ "transform_keys(%s, (k, v) -> k)",
+ "map_zip_with(%s, %s, (k, v1, v2) -> v1)",
+ "map_from_entries(map_entries(%s))",
+ "element_at(%s, id)",
+ "%s[id]") ++ anyTypeConsumers
+
+ private val structConsumers = Seq(
+ "%s.a",
+ "%s.b",
+ "struct(%s)",
+ "array(%s)",
+ "array(%s).b",
+ "size(array(%s))") ++ anyTypeConsumers
+
+ private val scalarConsumers = Seq(
+ "array(%s)",
+ "map(id, %s)",
+ "named_struct('a', %s)",
+ "coalesce(%s, NULL)",
+ "abs(%s)",
+ // Spark's analyzer rejects a void scalar in an array position (no implicit cast), so these
+ // stay skipped; they are here so the sweep notices if that ever changes.
+ "array_union(%s, array(1))",
+ "array_distinct(%s)",
+ "array_contains(%s, 1)",
+ "CAST(%s AS array)") ++ anyTypeConsumers
+
+ private def cases(wrap: String => String = identity): Seq[(String, String)] =
+ Seq(
+ arrayOfNull -> arrayConsumers,
+ arrayOfStructWithNull -> arrayConsumers,
+ mapWithNullValue -> mapConsumers,
+ mapWithNullKey -> mapConsumers,
+ structWithNull -> structConsumers,
+ scalarNull -> scalarConsumers).flatMap { case (producers, consumers) =>
+ for (p <- producers; c <- consumers)
+ yield {
+ val wrapped = wrap(p)
+ (wrapped, substitute(c, wrapped, p))
+ }
+ }
+
+ /**
+ * Fills a consumer template, putting `first` in the leading placeholder and `rest` in any
+ * others. Wrapping only the first argument keeps a stateful producer to a single occurrence:
+ * two of them desynchronize inside Spark itself, whereas the first child is one that every
+ * consumer evaluates, so a divergence there belongs to Comet.
+ */
+ private def substitute(template: String, first: String, rest: String): String = {
+ val at = template.indexOf("%s")
+ template.substring(0, at) + first + template.substring(at + 2).replace("%s", rest)
+ }
+
+ /**
+ * Makes a producer nullable and non-deterministic, the only input that can tell apart the two
+ * copies a serde null guard serializes (see `NullGuard.doubleEvaluationReason`).
+ *
+ * Applied to the first argument only, so multi-argument consumers stay in the sweep: a serde
+ * that null-guards every child, as `CometArraysZip` does, needs just one stateful argument to
+ * diverge, and restricting the sweep to single-placeholder consumers hid exactly that case.
+ */
+ private def nullableNondeterministic(producer: String): String =
+ s"IF(monotonically_increasing_id() % 2 = 0, $producer, NULL)"
+
+ /** Makes a producer nullable while keeping it deterministic. */
+ private def nullableDeterministic(producer: String): String =
+ s"IF(id % 2 = 0, $producer, NULL)"
+
+ /**
+ * Non-nullable, non-deterministic, NullType-bearing producers that record the counter, either
+ * in their value or, for the `filter` one, in their length. Under a null guard whose other
+ * argument is nullable, the THEN branch evaluates them on the guard's filtered rows only, so
+ * the counter sequence differs from Spark's even though the stateful child itself is never
+ * null; and a lambda producer runs through the JVM codegen dispatcher, whose kernel cache makes
+ * the guard's two copies share one counter, so it diverges under a single-argument guard too
+ * (`size(%s)` on a head that guarded non-nullable children). Paired with a deterministic
+ * producer of the same type for the sibling slot.
+ */
+ private val statefulProducers: Seq[(String, String, Seq[String])] = Seq(
+ (
+ "transform(array(id), x -> named_struct('i', monotonically_increasing_id(), 'n', NULL))",
+ "transform(array(id), x -> named_struct('i', x, 'n', NULL))",
+ arrayConsumers),
+ (
+ "filter(transform(array(id, 1, 2), x -> named_struct('i', x, 'n', NULL)), " +
+ "s -> s.i < monotonically_increasing_id())",
+ "filter(transform(array(id, 1, 2), x -> named_struct('i', x, 'n', NULL)), s -> s.i < id)",
+ arrayConsumers),
+ ("map(monotonically_increasing_id(), NULL)", "map(id, NULL)", mapConsumers),
+ (
+ "named_struct('i', monotonically_increasing_id(), 'n', NULL)",
+ "named_struct('i', id, 'n', NULL)",
+ structConsumers))
+
+ /**
+ * Each consumer with the stateful producer in the first slot and either a nullable
+ * deterministic sibling or a plain one in the others. The stateful producer stays first because
+ * Spark's `BinaryExpression.eval` returns NULL on a null left operand without evaluating the
+ * right one, so a stateful right operand's counter is an artifact of Spark's short-circuit and
+ * not a contract Comet can match; the left operand is evaluated on every row by both engines.
+ */
+ private def crossInputCases: Seq[(String, String)] =
+ for {
+ (stateful, deterministic, consumers) <- statefulProducers
+ c <- consumers
+ rest <- Seq(nullableDeterministic(deterministic), deterministic)
+ } yield (stateful, substitute(c, stateful, rest))
+
+ /**
+ * Query templates that put a producer under a different physical operator. The consumer sweep
+ * holds the operator fixed at a projection, so it never sees the paths that serialize and
+ * re-read a value rather than compute over it.
+ */
+ private val operatorTemplates = Seq(
+ "project" -> "SELECT %s AS c FROM t",
+ "filter" -> "SELECT %s AS c FROM t WHERE id > 2",
+ "sort-by-id" -> "SELECT %s AS c FROM t ORDER BY id",
+ "sort-by-value" -> "SELECT %s AS c FROM t ORDER BY c",
+ "limit" -> "SELECT %s AS c FROM t LIMIT 3",
+ "take-ordered" -> "SELECT %s AS c FROM t ORDER BY id LIMIT 3",
+ "groupby-key" -> "SELECT %s AS c, count(*) FROM t GROUP BY c",
+ "groupby-value" -> "SELECT id, first(%s) AS c FROM t GROUP BY id",
+ "groupby-last" -> "SELECT id, last(%s) AS c FROM t GROUP BY id",
+ "collect-list" -> "SELECT collect_list(%s) AS c FROM t",
+ "collect-set" -> "SELECT collect_set(%s) AS c FROM t",
+ "max" -> "SELECT max(%s) AS c FROM t",
+ "min" -> "SELECT min(%s) AS c FROM t",
+ "count-distinct" -> "SELECT count(DISTINCT %s) AS c FROM t",
+ "distinct" -> "SELECT DISTINCT %s AS c FROM t",
+ "union-all" -> "SELECT %s AS c FROM t UNION ALL SELECT %s AS c FROM t",
+ "union-distinct" -> "SELECT %s AS c FROM t UNION SELECT %s AS c FROM t",
+ "window-order" -> "SELECT %s AS c, row_number() OVER (ORDER BY id) AS r FROM t",
+ "window-partition" -> "SELECT id, count(*) OVER (PARTITION BY %s) AS n FROM t",
+ // The producer is computed on the build side, so the value itself crosses the exchange
+ // (shuffle or broadcast) rather than being projected after the join.
+ "join-shuffle" ->
+ "SELECT a.id, b.c FROM t a JOIN (SELECT id AS bid, %s AS c FROM t) b ON a.id = b.bid",
+ "join-broadcast" ->
+ ("SELECT /*+ BROADCAST(b) */ a.id, b.c FROM t a " +
+ "JOIN (SELECT id AS bid, %s AS c FROM t) b ON a.id = b.bid"),
+ "join-nested-loop" ->
+ ("SELECT /*+ BROADCAST(b) */ a.id, b.c FROM t a " +
+ "JOIN (SELECT id AS bid, %s AS c FROM t) b ON a.id > b.bid"),
+ "expand-cube" -> "SELECT id, count(%s) AS n FROM t GROUP BY CUBE(id)",
+ "explode" -> "SELECT explode(%s) AS c FROM t",
+ "repartition" -> "SELECT /*+ REPARTITION(3) */ %s AS c FROM t",
+ // Above `spark.shuffle.sort.bypassMergeThreshold`, so the JVM shuffle takes its sort-based
+ // writer, which hands a whole destination partition to one native call.
+ "repartition-many" -> "SELECT /*+ REPARTITION(300) */ %s AS c FROM t",
+ "coalesce-partitions" -> "SELECT /*+ COALESCE(1) */ %s AS c FROM t",
+ // The value itself is the partitioning key, the join key and a scalar subquery result, so
+ // the hash partitioner, the join's key comparison and the subquery's scalar conversion see
+ // a NullType rather than only carrying one past.
+ "distribute-by-value" -> "SELECT %s AS c FROM t DISTRIBUTE BY c",
+ "join-on-value" ->
+ ("SELECT a.id, b.id FROM (SELECT id, %s AS c FROM t) a " +
+ "JOIN (SELECT id, %s AS c FROM t) b ON a.c <=> b.c AND a.id = b.id"),
+ "scalar-subquery" -> "SELECT id, (SELECT max(x) FROM (SELECT %s AS x FROM t)) AS c FROM t",
+ "subquery" -> "SELECT id FROM t WHERE id IN (SELECT id FROM t WHERE %s IS NOT NULL)",
+ "nested-project" -> "SELECT c FROM (SELECT %s AS c, id FROM t ORDER BY id) x")
+
+ /**
+ * The subset of `operatorTemplates` that serializes the value rather than only computing over
+ * it.
+ */
+ private val serializingOperators = {
+ val names = Set(
+ "project",
+ "sort-by-id",
+ "groupby-value",
+ "collect-list",
+ "collect-set",
+ "repartition",
+ "repartition-many",
+ "distribute-by-value",
+ "join-on-value",
+ "union-all",
+ "join-shuffle")
+ require(names.subsetOf(operatorTemplates.map(_._1).toSet), "unknown operator name")
+ operatorTemplates.filter { case (name, _) => names(name) }
+ }
+
+ private val arrayOfStructWithNullField = Seq("array(named_struct('a', id, 'b', NULL))")
+
+ private def allProducers: Seq[String] =
+ arrayOfNull ++ mapWithNullValue ++ mapWithNullKey ++ arrayOfStructWithNull ++
+ structWithNull ++ scalarNull ++ arrayOfStructWithNullField
+
+ /**
+ * Containers to wrap a producer in before putting it under a serializing operator. The nested
+ * nullability mismatches only show where a container nests the value and an operator re-reads
+ * the nesting: `collect_list(array(map(k, NULL)))` fails while both `array(map(k, NULL))` and
+ * `collect_list(map(k, NULL))` pass.
+ */
+ private val nestingWrappers = Seq(
+ "array(%s)",
+ "array_repeat(%s, 2)",
+ "element_at(%s, 1)",
+ "slice(%s, 1, 1)",
+ "map(id, %s)",
+ "named_struct('s', %s)")
+
+ private val excludedOptimizerRules = Seq(
+ "ConstantFolding",
+ "NullPropagation",
+ "SimplifyBinaryComparison",
+ "SimplifyConditionals").map("org.apache.spark.sql.catalyst.optimizer." + _).mkString(",")
+
+ /**
+ * How rows are cut into batches and which path an exchange takes. The expression sweeps vary
+ * what is computed; these vary the machinery it runs through, because the defaults hide a whole
+ * class of defects: 16 rows in one batch never reuse a builder, `REPARTITION(3)` never leaves
+ * the JVM shuffle's bypass writer (one native call per batch), and a `spark.comet.batchSize` of
+ * 8192 never puts a batch boundary inside a stateful expression.
+ *
+ * `rows` is the size of `t`; the shuffle profiles need enough rows per destination partition to
+ * span several `spark.comet.shuffle.jvm.batchSize` batches inside one native call.
+ */
+ private case class Profile(name: String, rows: Int, confs: Seq[(String, String)])
+
+ private val defaultProfile = Profile("default", 16, Seq.empty)
+
+ /** Native batches of two rows, so every operator and stateful expression crosses a boundary. */
+ private val smallBatchProfile = Profile(
+ "small-batches",
+ 200,
+ Seq(
+ CometConf.COMET_BATCH_SIZE.key -> "2",
+ CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE.key -> "2",
+ CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true"))
+
+ /**
+ * Every registered array, map, struct and any-type aggregate serde opted into its native
+ * kernel. An `Incompatible` serde routes through the JVM codegen dispatcher by default, so
+ * without this profile its native kernel never runs in the sweep; `allowIncompatible` is what
+ * lets a user reach it, and the kernel's NullType handling has to hold there too.
+ */
+ private lazy val allowIncompatibleProfile = Profile(
+ "allow-incompatible",
+ 16,
+ registeredSerdes.toSeq
+ .map(cls => CometConf.getExprAllowIncompatConfigKey(cls) -> "true")
+ .sortBy(_._1))
+
+ /**
+ * Profiles for the sweeps whose queries end in a projection: batching and the kernel choice are
+ * what vary.
+ */
+ private lazy val kernelProfiles =
+ Seq(defaultProfile, smallBatchProfile, allowIncompatibleProfile)
+
+ /**
+ * Profiles for the sweeps that put a value through an operator. Each shuffle path appears once:
+ * JVM shuffle through the bypass writer (partition count under
+ * `spark.shuffle.sort.bypassMergeThreshold`), through the sort-based writer (above it, one
+ * whole partition per native call) with and without forced spills, and native shuffle. Across
+ * them AQE, native columnar-to-row and the two-row batch size each take both values, and every
+ * pair of those settings appears together at least once except native columnar-to-row off with
+ * AQE off, whose two mechanisms do not interact. Five profiles rather than one per setting,
+ * because each one runs the whole operator and nesting sweeps.
+ */
+ private val physicalProfiles = Seq(
+ defaultProfile,
+ smallBatchProfile.copy(confs = smallBatchProfile.confs ++ Seq(
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")),
+ Profile(
+ "jvm-sort-writer",
+ 900,
+ Seq(
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "300",
+ CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE.key -> "2",
+ CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> "100000",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true")),
+ Profile(
+ "jvm-sort-writer-spills",
+ 900,
+ Seq(
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "300",
+ CometConf.COMET_BATCH_SIZE.key -> "2",
+ CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE.key -> "2",
+ CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> "10",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true")),
+ Profile(
+ "native-shuffle",
+ 900,
+ Seq(
+ CometConf.COMET_SHUFFLE_MODE.key -> "native",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "300",
+ CometConf.COMET_BATCH_SIZE.key -> "2",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "false")))
+
+ private def rowsOf(
+ query: String,
+ cometEnabled: Boolean,
+ ansi: Boolean,
+ profile: Profile): (Seq[String], Boolean) = {
+ val confs = Seq(
+ CometConf.COMET_ENABLED.key -> cometEnabled.toString,
+ CometConf.COMET_EXEC_ENABLED.key -> cometEnabled.toString,
+ SQLConf.ANSI_ENABLED.key -> ansi.toString,
+ // Keep the producers and their consumers out of the optimizer's hands so they are
+ // evaluated per row by the engine under test rather than folded at plan time: the
+ // producers are deterministic and non-nullable, so without this `p IS NULL` becomes
+ // `false`, `p = p` and `p <=> p` become `true`, and `coalesce(p, p)` / `IF(c, p, p)`
+ // collapse to `p`, leaving a projection of literals that proves nothing.
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excludedOptimizerRules) ++ profile.confs
+ // withSQLConf's body is typed `=> Unit` on the older supported Spark versions and generic
+ // only on the newer ones, so the result is captured out of band to stay portable.
+ var result: (Seq[String], Boolean) = (Seq.empty, false)
+ withSQLConf(confs: _*) {
+ val df = spark.sql(query)
+ val rows = df.collect().map(_.toString()).sorted.toSeq
+ // The sweep's queries put the NullType expression in a projection, so a native
+ // CometProjectExec is what distinguishes a case that actually exercised a native kernel
+ // from one that merely fell back. Without this, a sweep where everything falls back would
+ // still be green while proving nothing.
+ val nativeProject =
+ collectFirst(df.queryExecution.executedPlan) { case _: CometProjectExec =>
+ ()
+ }.isDefined
+ result = (rows, nativeProject)
+ }
+ result
+ }
+
+ // ANSI is a dimension of the consumer sweeps because it changes which serdes wrap their child
+ // in a null guard and which kernels raise on out-of-range access; the operators below carry
+ // no ANSI semantics of their own.
+ // Each consumer sweep runs under the kernel profiles that can change its outcome and no
+ // others, since every (ansi, profile) pair is a full pass over the producer x consumer product.
+ for (ansi <- Seq(false, true)) {
+ def tag(profile: Profile): String = s"(ansi=$ansi ${profile.name})"
+
+ // The plain sweep is the one that exercises every kernel, so it takes every kernel profile.
+ for (profile <- kernelProfiles) {
+ test(s"NullType producers survive every consumer that Spark accepts ${tag(profile)}") {
+ sweep(
+ "consumer",
+ cases().map { case (producer, expr) => (producer, s"SELECT $expr FROM t") },
+ comparedFloor = 120,
+ nativeFloor = 100,
+ ansi = ansi,
+ profile = profile)
+ }
+ }
+
+ // The plain producers are non-nullable, so this is the only sweep where a serde's null guard
+ // runs natively over a NullType-bearing column that is NULL on some rows and takes its ELSE
+ // branch; the non-deterministic sweep below makes the same serdes leave the kernel instead.
+ // The guard is per row, so batching changes nothing here; which kernel sits under the guard
+ // does, so the allow-incompatible profile stays.
+ for (profile <- Seq(defaultProfile, allowIncompatibleProfile)) {
+ test(s"nullable deterministic NullType producers survive every consumer ${tag(profile)}") {
+ sweep(
+ "nullable",
+ cases(nullableDeterministic).map { case (producer, expr) =>
+ (producer, s"SELECT $expr FROM t")
+ },
+ comparedFloor = 140,
+ nativeFloor = 110,
+ ansi = ansi,
+ profile = profile)
+ }
+ }
+
+ // A stateful producer only diverges where a batch boundary falls inside it, so this sweep
+ // takes the two-row batches as well as the default; the kernel choice never sees the state.
+ for (profile <- Seq(defaultProfile, smallBatchProfile)) {
+ test(s"stateful NullType producers survive guards filtered by a sibling ${tag(profile)}") {
+ sweep(
+ "cross-input",
+ crossInputCases.map { case (producer, expr) => (producer, s"SELECT $expr FROM t") },
+ comparedFloor = 100,
+ nativeFloor = 60,
+ ansi = ansi,
+ profile = profile)
+ }
+ }
+
+ // A non-deterministic child makes every guarding serde decline its native kernel, so neither
+ // batching nor the kernel choice reaches the producer; ANSI still decides which serdes guard.
+ test(
+ s"nullable non-deterministic NullType producers survive every consumer ${tag(defaultProfile)}") {
+ sweep(
+ "non-deterministic",
+ cases(nullableNondeterministic).map { case (producer, expr) =>
+ (producer, s"SELECT $expr FROM t")
+ },
+ comparedFloor = 140,
+ nativeFloor = 110,
+ ansi = ansi,
+ profile = defaultProfile)
+ }
+ }
+
+ /**
+ * The gates the producers above trip. Each must publish its reason for the compatibility guide
+ * (`getUnsupportedReasons` is what `GenerateDocs` reads; the runtime note can add detail after
+ * it), and the ones whose whole shape Spark's own `doGenCode` can evaluate must enroll in the
+ * JVM codegen dispatcher so the projection stays in the Comet pipeline instead of falling back.
+ * Every query here is a shape from this suite's sweeps or the SQL file tests.
+ */
+ test("NullType gates enroll in the JVM codegen dispatcher and publish their reasons") {
+ val nullArray = "transform(array(id), x -> NULL)"
+ val nullScalar = "aggregate(array(id), NULL, (acc, x) -> NULL)"
+ val dispatched: Seq[(String, Class[_ <: Expression])] = Seq(
+ "array_repeat(map_entries(map(id, NULL)), 2)" -> classOf[ArrayRepeat],
+ // Beside a typed side Spark's coercion casts the NullType one away, so both sides are NullType.
+ s"array_union($nullArray, array())" -> classOf[ArrayUnion],
+ s"array($nullScalar)" -> classOf[CreateArray],
+ s"IF(id > 2, $nullScalar, NULL)" -> classOf[If],
+ s"CASE WHEN id > 2 THEN $nullScalar END" -> classOf[CaseWhen],
+ s"coalesce($nullScalar, $nullScalar)" -> classOf[Coalesce],
+ "coalesce(IF(monotonically_increasing_id() % 2 = 0, id, NULL), id)" -> classOf[Coalesce])
+ // Declined for a reason the dispatcher shares (a NullType input it cannot read, a kernel
+ // that only runs when the user opts in) or on a serde outside this PR's enrollment; these
+ // only have to publish their reason.
+ val documentedOnly: Seq[(String, Class[_ <: Expression])] = Seq(
+ s"array_intersect($nullArray, array())" -> classOf[ArrayIntersect],
+ "array_except(array(named_struct('a', id)), array(named_struct('a', id)))" ->
+ classOf[ArrayExcept],
+ "arrays_zip(array(monotonically_increasing_id()), array(id))" -> classOf[ArraysZip],
+ "map_from_arrays(array(id), array(1))" -> classOf[MapFromArrays])
+ // Shapes a gate must leave on the native kernel: the gate's plan-time proxy would match, but
+ // the native producer is known to be safe.
+ val nativeShapes: Seq[(String, Class[_ <: Expression])] = Seq(
+ // Native make_array emits a nullable item whatever Spark's containsNull says.
+ "array_repeat(array(id), 2)" -> classOf[ArrayRepeat])
+
+ withTempView("t") {
+ spark.range(0, 8).createOrReplaceTempView("t")
+ def serdeAndNode(
+ expr: String,
+ cls: Class[_ <: Expression]): (CometExpressionSerde[Expression], Expression) = {
+ val query = s"SELECT $expr AS c FROM t"
+ val node = spark
+ .sql(query)
+ .queryExecution
+ .analyzed
+ .flatMap(_.expressions)
+ .flatMap(_.collect { case e if e.getClass == cls => e })
+ .headOption
+ .getOrElse(fail(s"$query has no ${cls.getSimpleName} in its analyzed plan"))
+ (QueryPlanSerde.exprSerdeMap(cls).asInstanceOf[CometExpressionSerde[Expression]], node)
+ }
+ for ((expr, cls) <- dispatched ++ documentedOnly) {
+ val (serde, node) = serdeAndNode(expr, cls)
+ val reason = serde.getSupportLevel(node) match {
+ case Unsupported(Some(reason)) => reason
+ case other =>
+ fail(s"${cls.getSimpleName} reports $other for $expr, expected Unsupported")
+ }
+ assert(
+ serde.getUnsupportedReasons().exists(reason.startsWith),
+ s"${cls.getSimpleName}'s reason for $expr is not in getUnsupportedReasons: $reason")
+ }
+ for ((expr, cls) <- nativeShapes) {
+ val (serde, node) = serdeAndNode(expr, cls)
+ assert(
+ serde.getSupportLevel(node).isInstanceOf[Compatible],
+ s"${cls.getSimpleName} must keep $expr on the native kernel")
+ }
+ for ((expr, cls) <- dispatched ++ nativeShapes) {
+ val query = s"SELECT $expr AS c FROM t"
+ assert(
+ QueryPlanSerde.exprSerdeMap(cls).isInstanceOf[CodegenDispatchFallback],
+ s"${cls.getSimpleName} must mix in CodegenDispatchFallback")
+ val (sparkRows, _) = rowsOf(query, cometEnabled = false, ansi = false, defaultProfile)
+ val (cometRows, nativeProject) =
+ rowsOf(query, cometEnabled = true, ansi = false, defaultProfile)
+ assert(
+ cometRows == sparkRows,
+ s"$query: comet ${preview(cometRows)}, spark ${preview(sparkRows)}")
+ assert(nativeProject, s"$query left the Comet pipeline instead of dispatching")
+ }
+ }
+ }
+
+ // Every physical profile, because the operator sweep is the one whose queries actually cut
+ // rows into batches and push them through an exchange.
+ for (profile <- physicalProfiles) {
+ test(s"NullType producers survive every operator that Spark accepts (${profile.name})") {
+ sweep(
+ "operator",
+ for (producer <- allProducers; (op, template) <- operatorTemplates)
+ yield (op, template.replace("%s", producer)),
+ comparedFloor = 160,
+ nativeFloor = 100,
+ profile = profile)
+ }
+ }
+
+ /**
+ * Registered array, map and struct serdes whose nested-typed argument no `NullType` producer
+ * can occupy, so no consumer template can reach them: `Sequence` takes integral or temporal
+ * bounds, and `StringToMap` and `JsonToStructs` take a string.
+ */
+ private lazy val serdesWithoutNestedInput: Set[Class[_ <: Expression]] =
+ Set(classOf[Sequence], classOf[StringToMap], classOf[JsonToStructs])
+
+ /** Aggregates whose serde accepts any input type, so a `NullType` shape can reach them. */
+ /** The serdes the sweep must reach: those with a nested-typed argument plus the aggregates. */
+ private lazy val registeredSerdes: Set[Class[_]] =
+ (QueryPlanSerde.arrayExpressions.keySet ++
+ QueryPlanSerde.mapExpressions.keySet ++
+ QueryPlanSerde.structExpressions.keySet).toSet[Class[_]] --
+ serdesWithoutNestedInput ++ anyTypeAggregates
+
+ private lazy val anyTypeAggregates: Set[Class[_]] = Set(
+ classOf[CollectList],
+ classOf[CollectSet],
+ classOf[Count],
+ classOf[First],
+ classOf[Last],
+ classOf[Max],
+ classOf[Min])
+
+ /**
+ * Every expression class in the analyzed and optimized plans of `query`, or none if Spark
+ * rejects it. Both plans, because the optimizer is what inserts some expressions (`MapSort`
+ * under a map grouping key) and what expands `RuntimeReplaceable` ones.
+ */
+ private def expressionClasses(query: String): Set[Class[_]] =
+ Try {
+ val qe = spark.sql(query).queryExecution
+ Seq(qe.analyzed, qe.optimizedPlan)
+ }.toOption.toSeq.flatten
+ .toSet[LogicalPlan]
+ .flatMap { plan =>
+ plan.flatMap(_.expressions).flatMap { root =>
+ root.collect { case e: Expression => e }.flatMap {
+ case r: RuntimeReplaceable => r +: r.replacement.collect { case e: Expression => e }
+ case e => Seq(e)
+ }
+ }
+ }
+ .map(_.getClass)
+
+ // The consumer and operator lists are written by hand, so this is the check that a serde added
+ // to the registry later, or one forgotten now, does not silently stay outside the sweep.
+ test("the sweep reaches every registered array, map, struct and any-type aggregate serde") {
+ withTempView("t") {
+ spark.range(0, 8).createOrReplaceTempView("t")
+ // `registeredSerdes` is derived from QueryPlanSerde's registries and feeds two protections:
+ // this staleness check and `allowIncompatibleProfile`'s conf list. Were those registries
+ // renamed or emptied, both would degrade silently in the same direction -- `missing` becomes
+ // trivially empty and the profile collapses into a duplicate of `defaultProfile`. Pin the
+ // size so that cannot pass unnoticed.
+ assert(
+ registeredSerdes.size >= 40,
+ s"only ${registeredSerdes.size} serdes discovered; QueryPlanSerde's registries have moved " +
+ "and both this check and the allow-incompatible profile are now vacuous")
+ assert(
+ allowIncompatibleProfile.confs.size == registeredSerdes.size,
+ "the allow-incompatible profile must opt every serde the sweep reaches into its kernel")
+ val reached =
+ (cases().map { case (_, expr) =>
+ s"SELECT $expr FROM t"
+ } ++
+ (for (producer <- allProducers; (_, template) <- operatorTemplates)
+ yield template.replace("%s", producer)))
+ .flatMap(expressionClasses)
+ .toSet
+ val missing = registeredSerdes -- reached
+ assert(
+ missing.isEmpty,
+ s"registered serdes no sweep template reaches: ${missing
+ .map(_.getSimpleName)
+ .toSeq
+ .sorted
+ .mkString(", ")}")
+ }
+ }
+
+ for (profile <- physicalProfiles) {
+ test(s"nested NullType values survive the operators that serialize them (${profile.name})") {
+ // Select producers by the type they actually have: `filter(array(CAST(NULL AS int)), ...)`
+ // is `array`, and nesting it only exercises a pre-existing nested-container limitation.
+ val nullBearingProducers = allProducers.filter { p =>
+ Try(spark.range(0, 8).selectExpr(s"$p AS c").schema.head.dataType).toOption
+ .exists(SupportLevel.containsType(_, classOf[NullType]))
+ }
+ assert(
+ nullBearingProducers.size >= 5,
+ s"only ${nullBearingProducers.size} producers still carry a NullType; the producer list " +
+ "has drifted away from what this sweep is meant to cover")
+ val nested =
+ for (producer <- nullBearingProducers; wrapper <- nestingWrappers)
+ yield wrapper.replace("%s", producer)
+ sweep(
+ "nesting",
+ for (value <- nested.distinct; (op, template) <- serializingOperators)
+ yield (op, template.replace("%s", value)),
+ comparedFloor = 200,
+ nativeFloor = 100,
+ profile = profile)
+ }
+ }
+
+ /**
+ * Runs every `(label, query)` twice, Comet off and on, and reports all divergences together. A
+ * query Spark itself rejects is skipped; one whose Comet arm throws or disagrees is a failure.
+ *
+ * The floors are floors rather than equalities because which cases Spark accepts varies across
+ * the supported Spark versions. `comparedFloor` fails a sweep whose templates have gone stale;
+ * `nativeFloor` fails one where Comet fell back almost everywhere, since falling back is a pass
+ * and such a sweep would prove nothing about the native kernels.
+ */
+ /**
+ * Spark's answer to a query depends on the ANSI flag and the row count but on no Comet setting,
+ * so profiles that share both reuse it rather than run the reference side once per profile.
+ * `None` records a query Spark rejects.
+ */
+ private val sparkRowsCache = mutable.HashMap.empty[(String, Boolean, Int), Option[Seq[String]]]
+
+ private def sweep(
+ name: String,
+ queries: Seq[(String, String)],
+ comparedFloor: Int,
+ nativeFloor: Int,
+ ansi: Boolean = false,
+ profile: Profile = defaultProfile): Unit = {
+ var native = 0
+ withTempPath { dir =>
+ // One partition, so every batch holds several rows: a null guard over a non-deterministic
+ // child only diverges where the CASE sees both matching and non-matching rows in one batch.
+ spark.range(0, profile.rows, 1, 1).write.parquet(dir.getAbsolutePath)
+ withTempView("t") {
+ spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("t")
+
+ val failures = ArrayBuffer.empty[String]
+ var compared = 0
+ var skipped = 0
+
+ for ((label, query) <- queries) {
+ val sparkRows = sparkRowsCache.getOrElseUpdate(
+ (query, ansi, profile.rows),
+ Try(rowsOf(query, cometEnabled = false, ansi, profile)).toOption.map(_._1))
+ sparkRows match {
+ case None =>
+ skipped += 1
+ case Some(sparkRows) =>
+ compared += 1
+ Try(rowsOf(query, cometEnabled = true, ansi, profile)) match {
+ case Failure(e) =>
+ failures += s"[threw:$label] $query\n " +
+ s"${e.getClass.getSimpleName}: ${firstLine(causeText(e))}"
+ case Success((cometRows, _)) if cometRows != sparkRows =>
+ failures += s"[mismatch:$label] $query\n" +
+ s" spark: ${preview(sparkRows)}\n" +
+ s" comet: ${preview(cometRows)}"
+ case Success((_, nativeProject)) =>
+ if (nativeProject) native += 1
+ }
+ }
+ }
+
+ logWarning(
+ s"NullType $name sweep (ansi=$ansi, profile=${profile.name}): " +
+ s"compared=$compared native=$native skipped=$skipped total=${queries.length}")
+ assert(
+ compared >= comparedFloor,
+ s"$name sweep only compared $compared of ${queries.length} cases; " +
+ "its templates have gone stale")
+ assert(
+ native >= nativeFloor,
+ s"$name sweep only executed $native cases natively; Comet is falling back almost " +
+ "everywhere, so the comparison proves nothing about the native kernels")
+ if (failures.nonEmpty) {
+ // Name the profile: under `allow-incompatible` a mismatch can be a serde's *documented*
+ // difference (array_intersect's element order, array_except's null handling) rather than
+ // a NullType defect, and the two have to be told apart.
+ fail(
+ s"${failures.length} of $compared NullType $name cases diverge from Spark " +
+ s"(ansi=$ansi, profile=${profile.name}, $skipped invalid in Spark):\n" +
+ failures.mkString("\n"))
+ }
+ }
+ }
+ }
+
+ private def causeText(e: Throwable): String = {
+ val builder = new StringBuilder
+ var current: Throwable = e
+ var depth = 0
+ while (current != null && depth < 10) {
+ builder.append(Option(current.getMessage).getOrElse("")).append('\n')
+ current = current.getCause
+ depth += 1
+ }
+ builder.toString()
+ }
+
+ private def firstLine(message: String): String =
+ Option(message).map(_.linesIterator.next()).getOrElse("")
+
+ private def preview(rows: Seq[String]): String =
+ rows.take(3).mkString(", ") + (if (rows.length > 3) s", ... (${rows.length} rows)" else "")
+}
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
index bf354356386..70a045901d8 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
@@ -96,12 +96,57 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
checkShuffleAnswer(shuffled, 1)
}
+ test("columnar shuffle spanning several native writer batches with NullType columns") {
+ // The sort-based writer (`CometUnsafeShuffleWriter`, chosen once the partition count exceeds
+ // `spark.shuffle.sort.bypassMergeThreshold`) hands a whole destination partition to
+ // `process_sorted_row_partition`, which cuts it into `spark.comet.shuffle.jvm.batchSize`
+ // batches while reusing its builders. A `NullBuilder` keeps its length across `finish`, so
+ // every Null-bearing shape used to fail or miscount from the second batch on. The bypass
+ // writer never sends more than one batch per call, hence the partition count here, and the
+ // spill threshold is lifted so the writer sees several rows per partition.
+ withSQLConf(
+ CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE.key -> "2",
+ CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> "100000",
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.key -> "false") {
+ withParquetTable((0L until 2000L).map(Tuple1(_)), "tbl") {
+ val producers = Seq(
+ "named_struct('v', _1, 'n', NULL)",
+ "named_struct('s', named_struct('v', _1, 'n', NULL))",
+ "array(named_struct('v', _1, 'n', NULL))",
+ "element_at(transform(array(_1), x -> named_struct('v', x, 'n', NULL)), 1)",
+ "NULL",
+ "map(_1, NULL)",
+ "transform(array(_1), x -> NULL)")
+ producers.foreach { producer =>
+ val df = sql(s"SELECT /*+ REPARTITION(300) */ _1, $producer AS c FROM tbl")
+ checkShuffleAnswer(df, 1)
+ }
+ }
+ }
+ }
+
test("columnar shuffle with Map[_, NullType] column") {
val df = sql("SELECT id, map(id, null) AS m FROM VALUES (1), (2), (3) AS t(id)")
val shuffled = df.repartition(2, $"id")
checkShuffleAnswer(shuffled, 1)
}
+ test("columnar shuffle with Map[NullType, _] column") {
+ // map() leaves a NullType key, which Arrow requires to be non-nullable in the IPC schema;
+ // transform_values keeps a second copy non-foldable so it is built by the child rather than
+ // constant-folded into a literal. `map()` is MapType(NullType, NullType) only while the
+ // legacy flag is off.
+ withSQLConf(SQLConf.LEGACY_CREATE_EMPTY_COLLECTION_USING_STRING_TYPE.key -> "false") {
+ val df = sql(
+ "SELECT id, map() AS m1, transform_values(map(), (k, v) -> id) AS m2 " +
+ "FROM VALUES (1), (2), (3) AS t(id)")
+ assert(df.schema("m1").dataType.asInstanceOf[MapType].keyType === NullType)
+ assert(df.schema("m2").dataType.asInstanceOf[MapType].keyType === NullType)
+ val shuffled = df.repartition(2, $"id")
+ checkShuffleAnswer(shuffled, 1)
+ }
+ }
+
test("columnar shuffle on nested struct including nulls") {
Seq(10, 201).foreach { numPartitions =>
Seq("1.0", "10.0").foreach { ratio =>
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
index e72f4d10f74..36c667ae9b0 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -1007,6 +1007,39 @@ class CometInMemoryCacheSuite extends CometTestBase {
}
}
+ test("Comet in-memory cache round-trips NullType columns and children") {
+ withNativeCache {
+ // Non-foldable NullType shapes (each references `id`), so the cached batches hold Arrow
+ // NullVectors at the top level and under a list, a map value and a struct field rather
+ // than literals folded away at plan time. The nullable IF shape puts a NULL row around a
+ // NullType-bearing struct.
+ val query =
+ """
+ SELECT
+ id AS l,
+ aggregate(array(id), NULL, (acc, x) -> NULL) AS n,
+ transform(array(id), x -> NULL) AS an,
+ map(id, NULL) AS mn,
+ named_struct('a', id, 'b', NULL) AS sn,
+ IF(id % 2 = 0, named_struct('a', id, 'b', NULL), NULL) AS sno
+ FROM range(100)
+ """
+ val expected = spark.sql(query).orderBy("l").collect()
+
+ spark.sql(query).createOrReplaceTempView("null_types_cache")
+ spark.catalog.cacheTable("null_types_cache")
+ spark.table("null_types_cache").count()
+
+ assert(
+ cachedBatchTypes("null_types_cache").sameElements(
+ Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")))
+
+ val df = spark.sql("SELECT * FROM null_types_cache").orderBy("l")
+ assert(df.collect() === expected)
+ assert(df.queryExecution.executedPlan.toString().contains("CometInMemoryTableScan"))
+ }
+ }
+
test("Comet in-memory cache prunes only on columns that have bounds") {
assume(isSpark40Plus, "collated string types require Spark 4.0+")
withNativeCache {
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala
index 5af9a37c267..baaba71238c 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala
@@ -32,9 +32,10 @@ import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, IsNot
import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec}
import org.apache.spark.sql.execution.SparkPlan
-import org.apache.spark.sql.execution.adaptive.AQEShuffleReadExec
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec}
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType}
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType, MetadataBuilder, NullType, StructField, StructType}
import org.apache.comet.CometConf
@@ -1382,4 +1383,50 @@ class CometJoinSuite extends CometTestBase {
}
}
}
+
+ test("Broadcast HashJoin with NullType map columns on the build side") {
+ // map(k, NULL) and transform_values(map(), ...) leave NullType children in the map type
+ // (a bare map() would be constant-folded into a literal the optimizer hoists above the
+ // join). Such a build side cannot be coalesced for a Comet broadcast (see #5525 and
+ // `Utils.hasNullTypeUnderStruct`), so the exchange and the join stay on Spark with a
+ // reason that names the columns.
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ CometConf.COMET_BATCH_SIZE.key -> "100",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ // `map()` is MapType(NullType, NullType) only while this legacy flag is off.
+ SQLConf.LEGACY_CREATE_EMPTY_COLLECTION_USING_STRING_TYPE.key -> "false") {
+ withParquetTable((0 until 1000).map(i => (i, i % 5)), "tbl_a") {
+ withParquetTable((0 until 300).map(i => (i % 10, i + 2)), "tbl_b") {
+ val df = sql("""
+ |SELECT /*+ BROADCAST(b) */ tbl_a._1, b.m1, b.m2
+ |FROM tbl_a JOIN (
+ | SELECT _1, map(_1, NULL) AS m1, transform_values(map(), (k, v) -> _1) AS m2
+ | FROM tbl_b) b
+ |ON tbl_a._2 = b._1""".stripMargin)
+ // Guard the premise: without a NullType child this test still passes while
+ // covering nothing.
+ assert(df.schema("m1").dataType.asInstanceOf[MapType].valueType === NullType)
+ assert(df.schema("m2").dataType.asInstanceOf[MapType].keyType === NullType)
+ val (_, cometPlan) = checkSparkAnswerAndFallbackReason(
+ df,
+ "NullType directly under a struct or map entry in the broadcast build side " +
+ "(m1: map, m2: map) cannot be coalesced for broadcast")
+ def operators(p: SparkPlan): Seq[SparkPlan] = p +: (p match {
+ case a: AdaptiveSparkPlanExec => operators(a.executedPlan)
+ case s: QueryStageExec => operators(s.plan)
+ case r: ReusedExchangeExec => operators(r.child)
+ case _ => p.children.flatMap(operators)
+ })
+ val all = operators(cometPlan)
+ assert(all.exists(_.isInstanceOf[BroadcastExchangeExec]), cometPlan.treeString)
+ assert(!all.exists(_.isInstanceOf[CometBroadcastExchangeExec]), cometPlan.treeString)
+ }
+ }
+ }
+ }
+ }
}
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala
index 487e4797d12..d605c6d83a9 100644
--- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala
@@ -42,6 +42,8 @@ import org.apache.spark.sql.types.DecimalType
import org.apache.comet.CometConf
import org.apache.comet.CometSparkSessionExtensions
+import CometBenchmarkBase.BenchmarkArm
+
trait CometBenchmarkBase
extends SqlBasedBenchmark
with AdaptiveSparkPlanHelper
@@ -62,6 +64,7 @@ trait CometBenchmarkBase
.set(
"spark.shuffle.manager",
"org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+ extraSparkConf.foreach { case (key, value) => conf.set(key, value) }
val sparkSession = SparkSession
.builder()
@@ -80,6 +83,12 @@ trait CometBenchmarkBase
sparkSession
}
+ /**
+ * Static Spark settings a benchmark needs in its context, applied on top of the defaults above
+ * before the context starts. Session-level SQL configs belong in the cases instead.
+ */
+ protected def extraSparkConf: Map[String, String] = Map.empty
+
def runCometBenchmark(args: Array[String]): Unit
override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
@@ -123,26 +132,25 @@ trait CometBenchmarkBase
* SQL query to benchmark
* @param extraCometConfigs
* Additional configurations to apply for the Comet case (optional)
+ * @param extraArms
+ * Extra Comet cases, each overlaying its configs on the Comet case, e.g. a prior path with a
+ * dispatcher off (see [[CometBenchmarkBase.BenchmarkArm]]).
+ * @param expectNative
+ * Whether the Comet case is checked to be fully Comet native. False for a case that measures
+ * a deliberate fallback.
*/
final def runExpressionBenchmark(
name: String,
cardinality: Long,
query: String,
- extraCometConfigs: Map[String, String] = Map.empty): Unit = {
+ extraCometConfigs: Map[String, String] = Map.empty,
+ extraArms: Seq[BenchmarkArm] = Nil,
+ expectNative: Boolean = true): Unit = {
val benchmark = new Benchmark(name, cardinality, output = output)
- // Constant folding is excluded so that expressions over literal arguments are still evaluated
- // per row. It must be excluded for both arms: if only Comet excludes it, Spark folds the
- // expression away and does no per-row work, and the comparison is meaningless.
- val noConstantFolding =
- SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excludedRulesWith(ConstantFolding.ruleName)
-
- val sparkConfigs = Seq(noConstantFolding, CometConf.COMET_ENABLED.key -> "false")
-
- val cometConfigs = Seq(
- noConstantFolding,
- CometConf.COMET_ENABLED.key -> "true",
- CometConf.COMET_EXEC_ENABLED.key -> "true") ++ extraCometConfigs
+ val arms = expressionBenchmarkArms(extraCometConfigs, extraArms)
+ .map(arm => if (arm.name == "Comet") arm.copy(expectNative = expectNative) else arm)
+ val sparkConfigs = arms.head.configs
// Check that the benchmarked expression survives optimization into the Spark baseline plan.
// The query is not executed: before execution `stripAQEPlan` yields the initial physical
@@ -169,38 +177,64 @@ trait CometBenchmarkBase
}
}
- // Check that the plan is fully Comet native before running the benchmark. Unlike the check
- // above this one does execute the query, because AQE can re-plan a join after the first
- // stage completes and the Comet operators are what we are checking for.
- withSQLConf(cometConfigs: _*) {
- val df = spark.sql(query)
- df.noop()
- val plan = stripAQEPlan(df.queryExecution.executedPlan)
- findFirstNonCometOperator(plan).foreach { op =>
- warn(
- benchmark,
- s"""WARNING: the Comet plan is NOT fully Comet native, so the Comet case below is
- |partly or wholly measuring Spark.
- |First non-Comet operator: ${op.nodeName}
- |Comet plan:""".stripMargin + "\n" + plan.treeString)
+ // Check that the plan of every arm expected to be native is fully Comet native before running
+ // the benchmark. Unlike the check above this one does execute the query, because AQE can
+ // re-plan a join after the first stage completes and the Comet operators are what we are
+ // checking for.
+ arms.filter(_.expectNative).foreach { arm =>
+ withSQLConf(arm.configs: _*) {
+ val df = spark.sql(query)
+ df.noop()
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ findFirstNonCometOperator(plan).foreach { op =>
+ warn(
+ benchmark,
+ s"""WARNING: the ${arm.name} plan is NOT fully Comet native, so that case below is
+ |partly or wholly measuring Spark.
+ |First non-Comet operator: ${op.nodeName}
+ |Comet plan:""".stripMargin + "\n" + plan.treeString)
+ }
}
}
- benchmark.addCase("Spark") { _ =>
- withSQLConf(sparkConfigs: _*) {
- spark.sql(query).noop()
- }
- }
-
- benchmark.addCase("Comet") { _ =>
- withSQLConf(cometConfigs: _*) {
- spark.sql(query).noop()
+ arms.foreach { arm =>
+ benchmark.addCase(arm.name) { _ =>
+ withSQLConf(arm.configs: _*) {
+ spark.sql(query).noop()
+ }
}
}
benchmark.run()
}
+ /**
+ * The arms `runExpressionBenchmark` times, in order: the Spark baseline, the Comet case with
+ * `extraCometConfigs`, then each extra arm overlaid on the Comet case. Exposed so a benchmark
+ * can run its own checks (result equality, memory, plan) over exactly the timed configurations.
+ */
+ protected def expressionBenchmarkArms(
+ extraCometConfigs: Map[String, String],
+ extraArms: Seq[BenchmarkArm]): Seq[BenchmarkArm] = {
+ // Constant folding is excluded so that expressions over literal arguments are still evaluated
+ // per row. It must be excluded for both arms: if only Comet excludes it, Spark folds the
+ // expression away and does no per-row work, and the comparison is meaningless.
+ val noConstantFolding =
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excludedRulesWith(ConstantFolding.ruleName)
+
+ val sparkConfigs = Seq(noConstantFolding, CometConf.COMET_ENABLED.key -> "false")
+
+ val cometConfigs = Seq(
+ noConstantFolding,
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true") ++ extraCometConfigs
+
+ Seq(
+ BenchmarkArm("Spark", sparkConfigs, expectNative = false),
+ BenchmarkArm("Comet", cometConfigs, expectNative = true)) ++
+ extraArms.map(arm => arm.copy(configs = cometConfigs ++ arm.configs))
+ }
+
/**
* Returns the value of `spark.sql.optimizer.excludedRules` with `rule` appended, so that
* benchmark-specific exclusions do not clobber exclusions already configured by the caller.
@@ -382,6 +416,19 @@ trait CometBenchmarkBase
object CometBenchmarkBase {
+ /**
+ * One timed case of `runExpressionBenchmark`.
+ *
+ * @param name
+ * the case name in the results table
+ * @param configs
+ * the SQL configs the case runs under; for an extra arm, these overlay the Comet case
+ * @param expectNative
+ * whether the plan is checked to be fully Comet native before timing. False for an arm that
+ * deliberately measures a fallback path.
+ */
+ case class BenchmarkArm(name: String, configs: Seq[(String, String)], expectNative: Boolean)
+
/**
* SplitMix64 finalizer, used to turn a row id into a well distributed pseudo-random value.
*
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometNullTypeColumnsBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometNullTypeColumnsBenchmark.scala
new file mode 100644
index 00000000000..3be38cf4a8a
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometNullTypeColumnsBenchmark.scala
@@ -0,0 +1,475 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.benchmark
+
+import java.lang.management.{ManagementFactory, MemoryType}
+import java.nio.charset.StandardCharsets
+
+import scala.collection.JavaConverters._
+
+import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted, SparkListenerTaskEnd}
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometPlan}
+import org.apache.spark.sql.execution.{SparkPlan, SQLExecution}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, QueryStageExec}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec}
+import org.apache.spark.util.io.ChunkedByteBuffer
+
+import org.apache.comet.{CometArrowAllocator, CometConf}
+
+import CometBenchmarkBase.BenchmarkArm
+
+/**
+ * Benchmark for NullType-bearing columns on the two paths the NullType work touches: projections
+ * of NullType maps / structs consumed by the JVM codegen dispatcher, and broadcast joins whose
+ * build side carries a NullType column (a plain `array` coalesces through
+ * `Utils.coalesceBroadcastBatches`; a NullType directly under a struct or map entry cannot be
+ * coalesced, so the exchange and join stay on Spark).
+ *
+ * Every case is timed on five arms under the same session, warmup and iteration policy: Spark,
+ * Comet, Comet on the prior path (codegen dispatcher off, so the NullType projection falls back
+ * and takes its operator with it), and the two Comet arms again with a small
+ * `spark.comet.batchSize`. The build side is sized for few and for many broadcast buffers.
+ *
+ * After timing, each case checks that every arm returns the same rows and profiles one execution
+ * per arm: the executed plan (native or where it falls back), JVM allocation, peak and retained
+ * heap, Spark's peak task execution memory, the native operators' own memory metrics, retained
+ * Arrow memory, and for the broadcast cases the number of broadcast buffers, whether they were
+ * coalesced, the number of consuming tasks, and the resulting count of Arrow IPC streams opened
+ * (every consuming task decodes every buffer, see `CometBatchRDD.compute`). To run:
+ * {{{
+ * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometNullTypeColumnsBenchmark
+ * }}}
+ * Results will be written to "spark/benchmarks/CometNullTypeColumnsBenchmark-**results.txt".
+ */
+object CometNullTypeColumnsBenchmark extends CometBenchmarkBase {
+
+ private val smallBatchSize = 1024
+
+ private val dispatchOff = CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false"
+ private val smallBatch = CometConf.COMET_BATCH_SIZE.key -> smallBatchSize.toString
+
+ // The prior path: with the JVM codegen dispatcher off, a NullType-bearing projection falls
+ // back to Spark and takes its enclosing operator with it. Each Comet arm is also run with a
+ // small `spark.comet.batchSize`, the batch size of the native operators and of the JVM-side
+ // conversions. The native Parquet scan emits 8192-row batches whatever this is set to, so the
+ // broadcast buffer count is varied through the build size instead (see the profile columns).
+ private val extraArms: Seq[BenchmarkArm] = Seq(
+ BenchmarkArm("Comet (dispatch off, prior path)", Seq(dispatchOff), expectNative = false),
+ BenchmarkArm(s"Comet (batch $smallBatchSize)", Seq(smallBatch), expectNative = true),
+ BenchmarkArm(
+ s"Comet (dispatch off, prior path, batch $smallBatchSize)",
+ Seq(dispatchOff, smallBatch),
+ expectNative = false))
+
+ /**
+ * Times the case on every arm, then checks result equality and profiles each arm.
+ * `expectNative` is false for a case whose Comet plan is meant to fall back.
+ */
+ private def runNullTypeCase(
+ name: String,
+ cardinality: Long,
+ query: String,
+ expectNative: Boolean = true): Unit = {
+ val arms = if (expectNative) extraArms else extraArms.map(_.copy(expectNative = false))
+ runExpressionBenchmark(
+ name,
+ cardinality,
+ query,
+ extraArms = arms,
+ expectNative = expectNative)
+ verifyAndProfile(name, query, expressionBenchmarkArms(Map.empty, arms))
+ }
+
+ private def projectionBenchmarks(values: Int): Unit = {
+ withTempPath { dir =>
+ withTempTable("parquetV1Table") {
+ prepareTable(dir, spark.sql(s"SELECT value AS c1 FROM $tbl"))
+
+ runBenchmark("NullType projection consumed") {
+ runNullTypeCase(
+ "map(c1, NULL) projected and consumed",
+ values,
+ "SELECT size(map_keys(map(c1, NULL))) FROM parquetV1Table")
+ // Control: a plain list, the shape that stays on every fast path.
+ runNullTypeCase(
+ "transform to array projected and consumed (control)",
+ values,
+ "SELECT size(transform(array(c1), x -> NULL)) FROM parquetV1Table")
+ }
+ }
+ }
+ }
+
+ private def broadcastBuildBenchmarks(): Unit = {
+ val probeRows = 1024 * 1024
+
+ // Few and many broadcast buffers: the native scan emits 8192-row batches and the broadcast
+ // collects one buffer per batch, so the build size sets the buffer count. Files are split and
+ // packed by Spark's defaults; the hint forces the broadcast join.
+ for ((label, buildRows) <- Seq(("2 buffers", 2 * 8192), ("64 buffers", 64 * 8192))) {
+ withTempPath { dir =>
+ withTempTable("probe", "build") {
+ spark
+ .range(probeRows)
+ .selectExpr("id AS k", "id % 100 AS v")
+ .write
+ .parquet(s"${dir.getAbsolutePath}/probe")
+ // Parquet cannot store NullType, so the build table holds a plain column and each case
+ // projects its collection column in the join subquery.
+ spark
+ .range(buildRows)
+ .selectExpr("id AS k")
+ .write
+ .parquet(s"${dir.getAbsolutePath}/build")
+
+ spark.read.parquet(s"${dir.getAbsolutePath}/probe").createOrReplaceTempView("probe")
+ spark.read.parquet(s"${dir.getAbsolutePath}/build").createOrReplaceTempView("build")
+
+ runBenchmark(s"BroadcastHashJoin build with NullType column ($label)") {
+ def joinQuery(buildProjection: String, consumed: String): String =
+ s"""SELECT /*+ BROADCAST(b) */ count(p.v + $consumed)
+ |FROM probe p JOIN (SELECT k, $buildProjection FROM build) b ON p.k = b.k
+ |""".stripMargin
+
+ // Control: a plain list coalesces, since a NullVector under a list is safe to
+ // append.
+ runNullTypeCase(
+ "array column (control, coalesced)",
+ probeRows,
+ joinQuery("transform(array(k), x -> NULL) AS null_list", "size(b.null_list)"))
+ // A NullType directly under a struct or map entry cannot be coalesced (Arrow's
+ // appender cannot grow such a NullVector), so `CometBroadcastExchangeExec` refuses
+ // the build side and the exchange and join stay on Spark.
+ runNullTypeCase(
+ "map column (Spark broadcast)",
+ probeRows,
+ joinQuery("map(k, NULL) AS null_map", "size(map_keys(b.null_map))"),
+ expectNative = false)
+ // Consumed whole: a field access like `b.null_struct.a` is simplified to `b.k` by
+ // SimplifyExtractValueOps and the struct would never reach the broadcast.
+ runNullTypeCase(
+ "struct with NULL field column (Spark broadcast)",
+ probeRows,
+ joinQuery(
+ "named_struct('a', k, 'b', NULL) AS null_struct",
+ "size(array(b.null_struct))"),
+ expectNative = false)
+ }
+ }
+ }
+ }
+ }
+
+ override def runCometBenchmark(mainArgs: Array[String]): Unit = {
+ runBenchmarkWithTable("NullType projections", 1024 * 1024 * 10) { values =>
+ projectionBenchmarks(values)
+ }
+ broadcastBuildBenchmarks()
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Result equality and per-arm profiling
+ // ---------------------------------------------------------------------------------------------
+
+ /** An order-independent digest of a result set: row count and the sum / xor of row hashes. */
+ private case class Digest(rows: Long, sum: Long, xor: Long)
+
+ private case class BroadcastProfile(buffers: Int, coalescedBatches: Long, consumerTasks: Int) {
+ def ipcStreams: Long = buffers.toLong * consumerTasks
+ }
+
+ private case class Profile(
+ arm: String,
+ plan: String,
+ wallMs: Long,
+ jvmAllocatedBytes: Long,
+ heapPeakBytes: Long,
+ heapRetainedBytes: Long,
+ sparkPeakExecutionMemoryBytes: Long,
+ nativeMemoryBytes: Long,
+ arrowRetainedBytes: Long,
+ broadcast: Either[String, BroadcastProfile])
+
+ /** Records task and stage metrics of the jobs run while it is registered. */
+ private class ProfileListener extends SparkListener {
+ var taskPeakExecutionMemory = 0L
+ var broadcastConsumerTasks = 0
+
+ override def onTaskEnd(event: SparkListenerTaskEnd): Unit = synchronized {
+ Option(event.taskMetrics).foreach { metrics =>
+ taskPeakExecutionMemory = math.max(taskPeakExecutionMemory, metrics.peakExecutionMemory)
+ }
+ }
+
+ // A stage that consumes a Comet broadcast has the broadcast's `CometBatchRDD` in its lineage,
+ // and every one of its tasks decodes every broadcast buffer.
+ override def onStageCompleted(event: SparkListenerStageCompleted): Unit = synchronized {
+ if (event.stageInfo.rddInfos.exists(_.name == "CometBatchRDD")) {
+ broadcastConsumerTasks += event.stageInfo.numTasks
+ }
+ }
+ }
+
+ private def verifyAndProfile(name: String, query: String, arms: Seq[BenchmarkArm]): Unit = {
+ val digests = arms.map(arm => arm.name -> underConf(arm.configs)(digest(query)))
+ val (_, reference) = digests.head
+ val disagreeing = digests.filter(_._2 != reference)
+ if (disagreeing.nonEmpty) {
+ val detail = digests.map { case (arm, d) => s" $arm: $d" }.mkString("\n")
+ throw new IllegalStateException(
+ s"Result mismatch in '$name': arms disagree with ${digests.head._1}\n$detail")
+ }
+ emit(s"\n$name: results identical across ${arms.size} arms " +
+ s"(${reference.rows} rows, digest ${reference.sum.toHexString}/${reference.xor.toHexString})")
+
+ val profiles = arms.map(arm => underConf(arm.configs)(profile(arm, query)))
+ emitProfiles(profiles)
+ }
+
+ /** `withSQLConf` returns Unit; this variant returns the block's value. */
+ private def underConf[T](configs: Seq[(String, String)])(f: => T): T = {
+ var result: Option[T] = None
+ withSQLConf(configs: _*) { result = Some(f) }
+ result.get
+ }
+
+ private def digest(query: String): Digest = {
+ val (rows, sum, xor) = spark
+ .sql(query)
+ .rdd
+ .mapPartitions(NullTypeRowDigest.digestPartition)
+ .fold((0L, 0L, 0L)) { case ((r1, s1, x1), (r2, s2, x2)) =>
+ (r1 + r2, s1 + s2, x1 ^ x2)
+ }
+ Digest(rows, sum, xor)
+ }
+
+ private def profile(arm: BenchmarkArm, query: String): Profile = {
+ val heapPools =
+ ManagementFactory.getMemoryPoolMXBeans.asScala.filter(_.getType == MemoryType.HEAP)
+ val threads = ManagementFactory.getThreadMXBean.asInstanceOf[com.sun.management.ThreadMXBean]
+ def allocatedByAllThreads(): Long =
+ threads.getThreadAllocatedBytes(threads.getAllThreadIds).filter(_ > 0).sum
+ def usedHeapAfterGc(): Long = {
+ System.gc()
+ System.gc()
+ val runtime = Runtime.getRuntime
+ runtime.totalMemory - runtime.freeMemory
+ }
+
+ val listener = new ProfileListener
+ spark.sparkContext.addSparkListener(listener)
+ try {
+ val heapBefore = usedHeapAfterGc()
+ heapPools.foreach(_.resetPeakUsage())
+ val arrowBefore = CometArrowAllocator.getAllocatedMemory
+ val allocatedBefore = allocatedByAllThreads()
+ val start = System.nanoTime()
+
+ // Execute the dataset's own physical plan rather than `noop()`: a write plans a separate
+ // query, whose operators would be the ones carrying the metrics and the final AQE plan.
+ val qe = spark.sql(query).queryExecution
+ SQLExecution.withNewExecutionId(qe) {
+ qe.toRdd.foreach(_ => ())
+ }
+
+ val wallMs = (System.nanoTime() - start) / 1000000
+ val allocated = allocatedByAllThreads() - allocatedBefore
+ val heapPeak = heapPools.map(_.getPeakUsage.getUsed).sum
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+ val arrowRetained = CometArrowAllocator.getAllocatedMemory - arrowBefore
+ val heapRetained = usedHeapAfterGc() - heapBefore
+
+ val plan = stripAQEPlan(qe.executedPlan)
+ Profile(
+ arm.name,
+ describePlan(arm, plan),
+ wallMs,
+ allocated,
+ heapPeak,
+ heapRetained,
+ listener.taskPeakExecutionMemory,
+ nativeMemory(plan),
+ arrowRetained,
+ broadcastProfile(plan, listener.broadcastConsumerTasks))
+ } finally {
+ spark.sparkContext.removeSparkListener(listener)
+ }
+ }
+
+ /**
+ * Memory reported by the native operators' own metrics: the hash join's `build_mem_used` and
+ * the aggregate's `peak_mem_used`, summed over the plan and over tasks. Spark's task-level
+ * `peakExecutionMemory` only sees Spark operators, so the two columns are complementary.
+ */
+ private def nativeMemory(plan: SparkPlan): Long =
+ operators(plan)
+ .collect { case p: CometPlan => p }
+ .flatMap { p =>
+ p.metrics.collect {
+ case (name, metric) if name == "build_mem_used" || name == "peak_mem_used" =>
+ metric.value
+ }
+ }
+ .sum
+
+ /**
+ * Every operator of the executed plan in pre-order, descending into AQE query stages and reused
+ * exchanges, whose operators are only reachable through `plan` / `child` after execution. The
+ * original instances are returned (never a transformed copy: a copy carries fresh, unregistered
+ * metrics), each once, since a reused exchange is shared.
+ */
+ private def operators(plan: SparkPlan): Seq[SparkPlan] = {
+ val seen = java.util.Collections.newSetFromMap(
+ new java.util.IdentityHashMap[SparkPlan, java.lang.Boolean])
+ def walk(p: SparkPlan): Seq[SparkPlan] =
+ if (!seen.add(p)) {
+ Nil
+ } else {
+ p match {
+ case s: QueryStageExec => s +: walk(s.plan)
+ case r: ReusedExchangeExec => r +: walk(r.child)
+ case _ => p +: p.children.flatMap(walk)
+ }
+ }
+ walk(plan)
+ }
+
+ private def describePlan(arm: BenchmarkArm, plan: SparkPlan): String =
+ if (arm.name == "Spark") {
+ "Spark"
+ } else {
+ // `findFirstNonCometOperator` walks a subtree and stops at query stage leaves, so ask it
+ // about each operator of the expanded plan in turn; a node is non-Comet when it is the
+ // first one reported for its own subtree. Stage wrappers, reused exchanges and AQE shuffle
+ // reads are plumbing Comet consumes natively.
+ val plumbing =
+ Seq(classOf[QueryStageExec], classOf[ReusedExchangeExec], classOf[AQEShuffleReadExec])
+ operators(plan).find(op =>
+ findFirstNonCometOperator(op, plumbing: _*).exists(_ eq op)) match {
+ case None => "fully Comet native"
+ case Some(op) => s"falls back at ${op.nodeName}"
+ }
+ }
+
+ private def broadcastProfile(
+ plan: SparkPlan,
+ consumerTasks: Int): Either[String, BroadcastProfile] = {
+ val all = operators(plan)
+ all.collect { case b: CometBroadcastExchangeExec => b } match {
+ case Seq(b) =>
+ val buffers = b.executeBroadcast[Array[ChunkedByteBuffer]]().value.length
+ Right(BroadcastProfile(buffers, b.metrics("numCoalescedBatches").value, consumerTasks))
+ case Seq() =>
+ val sparkBroadcast = all.exists(_.isInstanceOf[BroadcastExchangeExec])
+ Left(if (sparkBroadcast) "Spark broadcast" else "no broadcast")
+ case many => Left(s"${many.size} Comet broadcasts")
+ }
+ }
+
+ private def emitProfiles(profiles: Seq[Profile]): Unit = {
+ def mb(bytes: Long): String = f"${bytes / (1024.0 * 1024.0)}%.1f"
+ val header = Seq(
+ "arm",
+ "plan",
+ "wall ms",
+ "JVM alloc MB",
+ "heap peak MB",
+ "heap retained MB",
+ "Spark peak exec MB",
+ "native mem MB",
+ "Arrow retained MB",
+ "bcast buffers",
+ "coalesced",
+ "consumer tasks",
+ "IPC streams")
+ val rows = profiles.map { p =>
+ val (buffers, coalesced, consumers, ipc) = p.broadcast match {
+ case Right(b) =>
+ (
+ b.buffers.toString,
+ b.coalescedBatches.toString,
+ b.consumerTasks.toString,
+ b.ipcStreams.toString)
+ case Left(reason) => (reason, "-", "-", "-")
+ }
+ Seq(
+ p.arm,
+ p.plan,
+ p.wallMs.toString,
+ mb(p.jvmAllocatedBytes),
+ mb(p.heapPeakBytes),
+ mb(p.heapRetainedBytes),
+ mb(p.sparkPeakExecutionMemoryBytes),
+ mb(p.nativeMemoryBytes),
+ mb(p.arrowRetainedBytes),
+ buffers,
+ coalesced,
+ consumers,
+ ipc)
+ }
+ val widths = (header +: rows).transpose.map(_.map(_.length).max)
+ def line(cells: Seq[String]): String =
+ cells
+ .zip(widths)
+ .zipWithIndex
+ .map { case ((cell, w), i) =>
+ if (i < 2) cell.padTo(w, ' ') else cell.reverse.padTo(w, ' ').reverse
+ }
+ .mkString(" ")
+ emit(
+ "Per-arm profile of one execution after warmup (JVM alloc: bytes allocated by all live " +
+ "threads; Spark peak exec: max peakExecutionMemory over tasks, reported by Spark " +
+ "operators only; native mem: the native join build / aggregate peak memory metrics " +
+ "summed over operators and tasks; IPC streams = broadcast buffers x consumer tasks):")
+ emit(line(header))
+ emit("-" * (widths.sum + 2 * (widths.size - 1)))
+ rows.foreach(row => emit(line(row)))
+ emit("")
+ }
+
+ /** Writes a line to the console and to the results file, like the timing tables. */
+ private def emit(text: String): Unit = {
+ // scalastyle:off println
+ println(text)
+ // scalastyle:on println
+ output.foreach(_.write((text + "\n").getBytes(StandardCharsets.UTF_8)))
+ }
+}
+
+/**
+ * Row digest used from executor closures; kept outside the benchmark object so it captures
+ * nothing.
+ */
+private[benchmark] object NullTypeRowDigest {
+ def digestPartition(rows: Iterator[Row]): Iterator[(Long, Long, Long)] = {
+ var count = 0L
+ var sum = 0L
+ var xor = 0L
+ rows.foreach { row =>
+ val h = row.mkString("\u0001").hashCode.toLong
+ count += 1
+ sum += h
+ xor ^= h
+ }
+ Iterator((count, sum, xor))
+ }
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala
index 4510f9d0ac1..b54616bf332 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala
@@ -19,11 +19,26 @@
package org.apache.spark.sql.comet.util
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
+import java.nio.channels.Channels
+
+import scala.concurrent.{Await, ExecutionContext, Future}
+import scala.concurrent.duration.DurationInt
+import scala.jdk.CollectionConverters._
+
import org.apache.arrow.c.CDataDictionaryProvider
+import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{ListVector, MapVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field}
import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData}
+import org.apache.spark.sql.comet.execution.arrow.CometArrowConverters
import org.apache.spark.sql.execution.vectorized.ConstantColumnVector
-import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType, TimestampType}
+import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, MapType, NullType, StringType, StructField, StructType, TimestampType}
import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.util.io.ChunkedByteBuffer
import org.apache.comet.CometArrowAllocator
import org.apache.comet.vector.CometVector
@@ -197,4 +212,198 @@ class UtilsSuite extends CometTestBase {
}
}
}
+
+ /**
+ * One `map` column of `numRows` empty maps, as `map()` produces (a `NullVector`
+ * key).
+ */
+ private def emptyNullKeyMapBatch(numRows: Int): ColumnarBatch = {
+ val field = Utils.toArrowField("m", MapType(NullType, NullType), nullable = true, "UTC")
+ val vector = field.createVector(CometArrowAllocator).asInstanceOf[MapVector]
+ vector.allocateNew()
+ (0 until numRows).foreach { i =>
+ vector.startNewValue(i)
+ vector.endValue(i, 0)
+ }
+ vector.getDataVector.setValueCount(0)
+ vector.setValueCount(numRows)
+ new ColumnarBatch(Array[ColumnVector](CometVector.getVector(vector, null)), numRows)
+ }
+
+ /**
+ * Writes `bound` through `writer` and asserts the stream reads back with a non-nullable key.
+ */
+ private def assertRoundTrip(
+ bound: VectorSchemaRoot,
+ writer: ArrowStreamWriter,
+ out: ByteArrayOutputStream,
+ numRows: Int): Unit = {
+ bound.setRowCount(numRows)
+ writer.start()
+ writer.writeBatch()
+ writer.end()
+ val reader =
+ new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray), CometArrowAllocator)
+ try {
+ assert(reader.loadNextBatch())
+ assert(reader.getVectorSchemaRoot.getRowCount == numRows)
+ assert(!mapKeyField(reader.getVectorSchemaRoot.getSchema.getFields.get(0)).isNullable)
+ } finally {
+ reader.close()
+ }
+ }
+
+ private def mapKeyField(field: Field): Field = field.getChildren.get(0).getChildren.get(0)
+
+ /** One `array` column; row `i` holds `i` nulls. */
+ private def nullListBatch(numRows: Int): ColumnarBatch = {
+ val field = Utils.toArrowField("l", ArrayType(NullType), nullable = true, "UTC")
+ val vector = field.createVector(CometArrowAllocator).asInstanceOf[ListVector]
+ vector.allocateNew()
+ (0 until numRows).foreach { i =>
+ vector.startNewValue(i)
+ vector.endValue(i, i)
+ }
+ vector.setValueCount(numRows)
+ new ColumnarBatch(Array[ColumnVector](CometVector.getVector(vector, null)), numRows)
+ }
+
+ test("withNonNullableMapKeys restores the non-nullable key flag a NullVector drops") {
+ val batch = emptyNullKeyMapBatch(2)
+ val field = batch.column(0).asInstanceOf[CometVector].getValueVector.getField
+ // The live vector reports a nullable key (see `Utils.withNonNullableMapKeys`). If this
+ // assertion starts failing, Arrow fixed that and the repair can go.
+ assert(mapKeyField(field).isNullable)
+
+ val repaired = Utils.withNonNullableMapKeys(field)
+ assert(!mapKeyField(repaired).isNullable)
+ assert(mapKeyField(repaired).getType.isInstanceOf[ArrowType.Null])
+ assert(repaired.getName == field.getName)
+ assert(repaired.getFieldType == field.getFieldType)
+ assert(repaired.getChildren.get(0).getFieldType == field.getChildren.get(0).getFieldType)
+ // Idempotent, and a no-op on fields that already satisfy the invariant.
+ assert(Utils.withNonNullableMapKeys(repaired) eq repaired)
+ batch.close()
+ }
+
+ test("newArrowStreamWriter keeps a root whose declared schema is already valid") {
+ val batch = emptyNullKeyMapBatch(2)
+ val vector =
+ batch.column(0).asInstanceOf[CometVector].getValueVector.asInstanceOf[FieldVector]
+ val declared = Utils.withNonNullableMapKeys(vector.getField)
+ // A root whose declared schema is already valid must be kept, not swapped for a copy.
+ val root = new VectorSchemaRoot(Seq(declared).asJava, Seq(vector).asJava, 0)
+ val out = new ByteArrayOutputStream()
+ val (bound, writer) = Utils.newArrowStreamWriter(root, null, Channels.newChannel(out))
+ assert(bound eq root)
+ assertRoundTrip(bound, writer, out, numRows = 2)
+ batch.close()
+ }
+
+ test("newArrowStreamWriter returns the root a later row count must be set on") {
+ val batch = emptyNullKeyMapBatch(2)
+ val vector =
+ batch.column(0).asInstanceOf[CometVector].getValueVector.asInstanceOf[FieldVector]
+ // Declared from the live vector, so the key is nullable and the root must be swapped.
+ val root = new VectorSchemaRoot(Seq(vector).asJava)
+ val out = new ByteArrayOutputStream()
+ val (bound, writer) = Utils.newArrowStreamWriter(root, null, Channels.newChannel(out))
+ assert(bound ne root)
+ assertRoundTrip(bound, writer, out, numRows = 2)
+ batch.close()
+ }
+
+ test("coalesceBroadcastBatches bypasses exactly the schemas with a NullType under a struct") {
+ // Exhaustive over the shape space of the bypass rule: VectorAppender hangs only when a
+ // NullVector is a *direct* child of a struct (see `Utils.hasNullDirectlyUnderStruct` for
+ // the Arrow mechanics). Each shape runs the real appender under a timeout, so a rule that
+ // is too narrow shows up as a timeout on the hanging shapes instead of a hung build, and
+ // one that is too wide shows up as a needless bypass. Once an Arrow release fixes the
+ // appender, the shapes reported as needless bypasses are the ones whose gate can go.
+ val nullStruct = StructType(Seq(StructField("a", NullType)))
+ val shapes: Seq[(DataType, Any)] = Seq(
+ NullType -> null,
+ ArrayType(NullType) -> new GenericArrayData(Array[Any](null)),
+ ArrayType(ArrayType(NullType)) ->
+ new GenericArrayData(Array[Any](new GenericArrayData(Array[Any](null)))),
+ nullStruct -> InternalRow(null),
+ ArrayType(nullStruct) -> new GenericArrayData(Array[Any](InternalRow(null))),
+ StructType(Seq(StructField("l", ArrayType(NullType)))) ->
+ InternalRow(new GenericArrayData(Array[Any](null))),
+ MapType(IntegerType, NullType) -> ArrayBasedMapData(Array[Any](1), Array[Any](null)),
+ // `map(k, array(NULL))`: the entry struct's direct child is a list, not a NullVector, so
+ // this still coalesces.
+ MapType(IntegerType, ArrayType(NullType)) ->
+ ArrayBasedMapData(Array[Any](1), Array[Any](new GenericArrayData(Array[Any](null)))),
+ MapType(NullType, NullType) -> ArrayBasedMapData(Array.empty[Any], Array.empty[Any]))
+
+ val numRows = 4
+ val numBatches = 3
+ shapes.foreach { case (dataType, value) =>
+ val name = dataType.simpleString
+ val schema = StructType(Seq(StructField("c", dataType)))
+ val batches = (0 until numBatches).map { _ =>
+ CometArrowConverters
+ .rowToArrowBatchIter(
+ Iterator.fill(numRows)(InternalRow(value)),
+ schema,
+ numRows,
+ "UTC",
+ CometArrowAllocator)
+ .next()
+ }
+ // Force serialization eagerly: on Scala 2.12 `toSeq` is a lazy Stream, and the inputs
+ // are closed on the next line.
+ val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toVector
+ batches.foreach(_.close())
+
+ val (result, batchCount, totalRows) = Await.result(
+ Future(Utils.coalesceBroadcastBatches(bufs.iterator))(ExecutionContext.global),
+ 10.seconds)
+
+ // The planner's gate (`CometBroadcastExchangeExec.getSupportLevel`) uses this same
+ // predicate, so this also pins that it refuses exactly the shapes the appender cannot take.
+ val expectBypass = Utils.hasNullTypeUnderStruct(dataType)
+ assert(
+ (batchCount == 0) == expectBypass,
+ s"$name: batchCount=$batchCount but bypass expected=$expectBypass")
+ assert(result.length == (if (expectBypass) numBatches else 1), name)
+ if (expectBypass) assert(totalRows == 0, name)
+ else assert(totalRows == numRows.toLong * numBatches, name)
+
+ val decoded = result.iterator.flatMap(b => Utils.decodeBatches(b, "test")).toSeq
+ assert(decoded.map(_.numRows()).sum == numRows.toLong * numBatches, name)
+ decoded.foreach(_.close())
+ }
+ }
+
+ test("coalesceBroadcastBatches keeps coalescing plain null lists") {
+ // A NullVector directly under a list appends fine (see `Utils.hasNullDirectlyUnderStruct`);
+ // bypassing here would cost every consuming task one IPC stream per original buffer.
+ val numRows = 4
+ val numBatches = 3
+ val batches = (0 until numBatches).map(_ => nullListBatch(numRows))
+ val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toVector
+
+ val (result, batchCount, totalRows) = Utils.coalesceBroadcastBatches(bufs.iterator)
+ assert(batchCount == numBatches)
+ assert(totalRows == numRows.toLong * numBatches)
+ assert(result.length == 1)
+
+ // Every list keeps its length across the append. Read each batch before the stream moves
+ // on, since the reader reclaims a batch's buffers when it loads the next one.
+ def listLengths(bufs: Iterator[ChunkedByteBuffer]): Seq[Int] =
+ bufs
+ .flatMap(b => Utils.decodeBatches(b, "test"))
+ .flatMap { batch =>
+ val lengths =
+ (0 until batch.numRows()).map(r => batch.column(0).getArray(r).numElements())
+ batch.close()
+ lengths
+ }
+ .toSeq
+ val expected = Seq.fill(numBatches)(0 until numRows).flatten
+ assert(listLengths(bufs.iterator) == expected, "uncoalesced input")
+ assert(listLengths(result.iterator) == expected, "coalesced output")
+ }
}