feat: support Map for CreateArray literal

native/core/src/execution/planner.rs

@@ -48,7 +48,9 @@ use crate::execution::{
 };
 use crate::jvm_bridge::{jni_call, JVMClasses};
 use arrow::compute::CastOptions;
-use arrow::datatypes::{DataType, Field, FieldRef, Schema, TimeUnit, DECIMAL128_MAX_PRECISION};
+use arrow::datatypes::{
+    DataType, Field, FieldRef, Fields, Schema, TimeUnit, DECIMAL128_MAX_PRECISION,
+};
 use arrow::ffi_stream::FFI_ArrowArrayStream;
 use datafusion::functions_aggregate::bit_and_or_xor::{bit_and_udaf, bit_or_udaf, bit_xor_udaf};
 use datafusion::functions_aggregate::count::count_udaf;
@@ -216,6 +218,39 @@ fn make_all_fields_nullable(data_type: &DataType) -> DataType {
     }
 }
 
+/// Return a copy of a `Map` type with only the outer entries `value` field marked nullable, keeping
+/// the key field non-nullable and every nested key/value type byte-for-byte unchanged. Any non-`Map`
+/// type is returned unchanged.
+///
+/// This is the shallow counterpart of `make_all_fields_nullable`, used for `map_entries`.
+/// `map_entries` reuses the input map's entries array as its output list values but declares that
+/// element's `value` field nullable (`Struct(key non-null, value nullable)`), deriving both nested
+/// types verbatim from the input. Arrow's `GenericListArray::try_new` compares the declared element
+/// field's type against the reused values array's type in full, so the ONLY field that can mismatch
+/// is the entries `value` field's own `nullable` flag; widening just that field is sufficient.
+/// Recursing into nested types (as `make_all_fields_nullable` does) would additionally flip, say, a
+/// nested `Map`'s `valueContainsNull`, which then diverges from the Spark-serialized `return_type`
+/// and makes a downstream `make_array` see unequal map types and panic.
+fn widen_map_entry_value_nullable(data_type: &DataType) -> DataType {
+    match data_type {
+        DataType::Map(entries, sorted) => match entries.data_type() {
+            DataType::Struct(kv) if kv.len() == 2 && !kv[1].is_nullable() => {
+                let new_value = Arc::new(kv[1].as_ref().clone().with_nullable(true));
+                let new_kv: Fields = vec![Arc::clone(&kv[0]), new_value].into();
+                let new_entries = Arc::new(
+                    entries
+                        .as_ref()
+                        .clone()
+                        .with_data_type(DataType::Struct(new_kv)),
+                );
+                DataType::Map(new_entries, *sorted)
+            }
+            _ => data_type.clone(),
+        },
+        other => other.clone(),
+    }
+}
+
 /// If `expr` evaluates to `Timestamp(_, Some(_))` against `schema`, wrap it in a
 /// metadata-only cast to `Timestamp(_, None)`. This is required because
 /// DataFusion's `SortMergeJoinExec` comparator only supports timezone-less
@@ -2897,13 +2932,13 @@ impl PhysicalPlanner {
             }
             AggExprStruct::CollectSet(expr) => {
                 let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?;
-                let child = Self::coerce_collect_child_nullability(child, &schema)?;
+                let child = Self::coerce_child_fields_nullable(child, &schema)?;
                 let func = AggregateUDF::new_from_impl(SparkCollectSet::new());
                 Self::create_aggr_func_expr("collect_set", schema, vec![child], func)
             }
             AggExprStruct::CollectList(expr) => {
                 let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?;
-                let child = Self::coerce_collect_child_nullability(child, &schema)?;
+                let child = Self::coerce_child_fields_nullable(child, &schema)?;
                 let func = AggregateUDF::new_from_impl(SparkCollectList::new());
                 Self::create_aggr_func_expr("collect_list", schema, vec![child], func)
             }
@@ -3384,6 +3419,17 @@ impl PhysicalPlanner {
             .collect::<Result<Vec<_>, _>>()?;
 
         let fun_name = &expr.func;
+        // `map_entries` needs its argument's entry `value` field widened to nullable first (only
+        // that outer field). See `widen_map_entry_value_nullable`.
+        let args = if fun_name == "map_entries" {
+            args.into_iter()
+                .map(|arg| {
+                    Self::coerce_child_to(arg, &input_schema, widen_map_entry_value_nullable)
+                })
+                .collect::<Result<Vec<_>, ExecutionError>>()?
+        } else {
+            args
+        };
         let input_expr_types = args
             .iter()
             .map(|x| x.data_type(input_schema.as_ref()))
@@ -3516,27 +3562,41 @@ impl PhysicalPlanner {
         Ok(scalar_expr)
     }
 
-    /// `collect_list` / `collect_set` build their result list with all element fields marked
-    /// nullable, regardless of the input's nullability. However `SparkCollectList::return_type`
-    /// (and `SparkCollectSet`) derive the element type directly from the child, preserving any
-    /// non-nullable nested field. When the child is a nested type with a non-nullable inner field
-    /// (e.g. a struct field), the declared aggregate output disagrees with the array the
-    /// accumulator actually produces, and DataFusion's grouped `AggregateExec` fails validating
-    /// its output batch ("column types must match schema types"). Cast the child to the
-    /// all-nullable variant of its type so the declared and produced types stay consistent.
-    fn coerce_collect_child_nullability(
+    /// Casts `child` so its type matches `widen(child_type)`, wrapping it in a `CastExpr` only when
+    /// the widened type differs. Shared skeleton for `coerce_child_fields_nullable` and the
+    /// `map_entries` argument widening; each caller supplies the widening its consumer needs.
+    fn coerce_child_to(
         child: Arc<dyn PhysicalExpr>,
         schema: &SchemaRef,
+        widen: impl Fn(&DataType) -> DataType,
     ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
         let child_type = child.data_type(schema.as_ref())?;
-        let nullable_type = make_all_fields_nullable(&child_type);
-        if child_type.equals_datatype(&nullable_type) {
+        let widened = widen(&child_type);
+        if child_type.equals_datatype(&widened) {
             Ok(child)
         } else {
-            Ok(Arc::new(CastExpr::new(child, nullable_type, None)))
+            Ok(Arc::new(CastExpr::new(child, widened, None)))
         }
     }
 
+    /// Casts `child` to the all-nullable variant of its own type, or returns it unchanged when it
+    /// already is. Used by `collect_list` / `collect_set`, which build their result list with all
+    /// element fields marked nullable while `SparkCollectList::return_type` (and `SparkCollectSet`)
+    /// derive the element type directly from the child. With a non-nullable nested field (e.g. a
+    /// struct field), DataFusion's grouped `AggregateExec` fails validating its output batch
+    /// ("column types must match schema types").
+    ///
+    /// `make_all_fields_nullable` keeps an Arrow map's key field non-nullable, so widening a map
+    /// stays a legal map type. `map_entries` needs only its outer entry `value` field widened, not
+    /// the deep widening this performs; it passes `widen_map_entry_value_nullable` to
+    /// `coerce_child_to` instead.
+    fn coerce_child_fields_nullable(
+        child: Arc<dyn PhysicalExpr>,
+        schema: &SchemaRef,
+    ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
+        Self::coerce_child_to(child, schema, make_all_fields_nullable)
+    }
+
     fn create_aggr_func_expr(
         name: &str,
         schema: SchemaRef,

spark/src/main/java/org/apache/comet/codegen/CometBatchKernel.java

@@ -31,7 +31,13 @@
  */
 public abstract class CometBatchKernel extends CometInternalRow {
 
-  protected final Object[] references;
+  // `public` (not `protected`) so that the nested helper classes Spark's codegen emits when it
+  // splits a large expression (e.g. a folded map rebuilt as a big `CreateMap`) can read it. Those
+  // helpers are separate classes in the generated package, not subclasses of this one, so a
+  // `protected` field would raise `IllegalAccessError` at runtime under cross-package protected
+  // access rules. Spark's own generated classes avoid this by declaring `references` on the
+  // generated class itself; Comet inherits it here instead, so it must be public.
+  public final Object[] references;
 
   protected CometBatchKernel(Object[] references) {
     this.references = references;

spark/src/main/scala/org/apache/comet/serde/arrays.scala

@@ -472,27 +472,20 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] {
 
     // DataFusion's `make_array` asserts strict element-type equality in
     // `MutableArrayData::with_capacities` and panics on a mismatch. Spark's CreateArray is more
-    // permissive: its type coercion compares element types with `sameType`, which ignores
-    // nullability, so children that share a surface type but differ only in nested field
-    // nullability get no unifying cast. DataFusion tolerates container nullability differences
-    // (an `ArrayType.containsNull` / `MapType.valueContainsNull` mismatch is coerced), but NOT a
-    // struct field's nullability -- `array(struct(a not null), struct(a nullable))` panics inside
-    // `make_array_inner`. Decline only those cases (i.e. children that still differ after
-    // normalizing container nullability) so Spark's evaluator handles them.
-    //
-    // TODO: remove this decline once apache/datafusion#22366 lands; the upstream fix widens the
-    // element type via nullability-OR-merge and casts each child before MutableArrayData.
-    val normalizedTypes = children.map(c => normalizeContainerNullability(c.dataType))
-    if (normalizedTypes.distinct.size > 1) {
-      withFallbackReason(
-        expr,
-        "CreateArray children have mismatched data types: " +
-          children.map(_.dataType).distinct.mkString(", "))
-      return None
+    // permissive: its coercion compares element types with `sameType` (nullability ignored), so
+    // children that share a surface type but differ in nullability reach here as distinct types.
+    // Comet's native runtime types are also frequently MORE nullable than Spark's Catalyst types
+    // (`map_entries` forces the entry `value` field nullable, list elements are nullable, ...), so
+    // casting to Spark's declared element type does not reliably unify them. Cast every child to a
+    // deeply-nullable element type instead (every array/map/struct field nullable at all nesting
+    // levels; the cast only widens metadata and never changes values), so `make_array` always sees
+    // identical Arrow types. A child whose cast is unsupported declines below.
+    val elementType = deepNullable(expr.dataType.asInstanceOf[ArrayType].elementType)
+    val childExprs = children.map { c =>
+      val unified = if (c.dataType == elementType) c else Cast(c, elementType)
+      exprToProtoInternal(unified, inputs, binding)
     }
 
-    val childExprs = children.map(exprToProtoInternal(_, inputs, binding))
-
     if (childExprs.forall(_.isDefined)) {
       scalarFunctionExprToProto("make_array", childExprs: _*)
     } else {
@@ -502,22 +495,20 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] {
   }
 
   /**
-   * Rewrites a type so that container nullability (`ArrayType.containsNull`,
-   * `MapType.valueContainsNull`) is forced to `true` everywhere, while struct field nullability
-   * is left intact. Two CreateArray children whose types differ ONLY in container nullability are
-   * tolerated by DataFusion's `make_array` (coerced), so they normalize equal here; a difference
-   * in a struct field's nullability survives normalization and triggers the decline above.
+   * A copy of `dt` with every array `containsNull`, map `valueContainsNull`, and struct field
+   * nullability forced to `true` at every nesting level (map key fields stay non-null per Arrow's
+   * map invariant). This is Spark's `DataType.asNullable`, re-derived here because that method is
+   * `private[spark]` and unreachable from this package. Used as the common cast target for
+   * `CometCreateArray` so that children whose native runtime types are more nullable than Spark's
+   * Catalyst types (e.g. a `map_entries` entry struct, whose `value` field Comet forces nullable)
+   * still unify for `make_array`.
    */
-  private def normalizeContainerNullability(dt: DataType): DataType = dt match {
-    case ArrayType(elementType, _) =>
-      ArrayType(normalizeContainerNullability(elementType), containsNull = true)
-    case MapType(keyType, valueType, _) =>
-      MapType(
-        normalizeContainerNullability(keyType),
-        normalizeContainerNullability(valueType),
-        valueContainsNull = true)
+  private def deepNullable(dt: DataType): DataType = dt match {
+    case ArrayType(et, _) => ArrayType(deepNullable(et), containsNull = true)
+    case MapType(kt, vt, _) =>
+      MapType(deepNullable(kt), deepNullable(vt), valueContainsNull = true)
     case StructType(fields) =>
-      StructType(fields.map(f => f.copy(dataType = normalizeContainerNullability(f.dataType))))
+      StructType(fields.map(f => f.copy(dataType = deepNullable(f.dataType), nullable = true)))
     case other => other
   }
 }
@@ -593,7 +584,7 @@ object CometElementAt extends CometExpressionSerde[ElementAt] {
   override def getSupportLevel(expr: ElementAt): SupportLevel = {
     expr.left.dataType match {
       case _: ArrayType => Compatible()
-      case _: MapType => Compatible()
+      case MapType(keyType, _, _) => MapKeySupport.keySupport(keyType)
       case _ => Unsupported(Some("Input must be an array or map"))
     }
   }