feat: support Map for CreateArray literal

native/core/src/execution/planner.rs

@@ -49,7 +49,9 @@ use crate::execution::{
 };
 use crate::jvm_bridge::{jni_call, JVMClasses, ShufflePartitionPusher};
 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;
@@ -218,6 +220,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
@@ -3490,6 +3525,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()))
@@ -3653,6 +3699,26 @@ impl PhysicalPlanner {
         Ok(Arc::new(CastExpr::new(child, nullable_type, None)))
     }
 
+    /// Casts `child` so its type matches `widen(child_type)`, wrapping it in a `CastExpr` only when
+    /// the widened type differs. Used for the `map_entries` argument widening (with
+    /// `widen_map_entry_value_nullable`). Unlike `coerce_collect_child_nullability`, which casts
+    /// unconditionally as a normalization barrier for aggregate accumulators, this casts only when
+    /// the type actually changes. That is sufficient for the scalar `map_entries`, whose reused
+    /// entries array only needs its declared element type to line up.
+    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 widened = widen(&child_type);
+        if child_type.equals_datatype(&widened) {
+            Ok(child)
+        } else {
+            Ok(Arc::new(CastExpr::new(child, widened, None)))
+        }
+    }
+
     fn create_aggr_func_expr(
         name: &str,
         schema: SchemaRef,

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

@@ -18,7 +18,7 @@
 use arrow::array::StructArray;
 use arrow::datatypes::{DataType, Field, Schema};
 use arrow::record_batch::RecordBatch;
-use datafusion::common::Result as DataFusionResult;
+use datafusion::common::{Result as DataFusionResult, ScalarValue};
 use datafusion::logical_expr::ColumnarValue;
 use datafusion::physical_expr::PhysicalExpr;
 use std::{
@@ -71,13 +71,25 @@ impl PhysicalExpr for CreateNamedStruct {
             .iter()
             .map(|expr| expr.evaluate(batch))
             .collect::<datafusion::common::Result<Vec<_>>>()?;
+        // When every field value is a scalar (e.g. an all-literal `named_struct` that reaches native
+        // as a `CreateNamedStruct` because constant folding is disabled), return a scalar struct
+        // rather than a length-1 `StructArray`. A downstream consumer such as `make_array` then
+        // broadcasts the constant struct to the batch row count instead of failing with a
+        // mixed-length error when a sibling child is a full-length column. This matches what a
+        // constant-folded struct literal produces, and `GetStructField::evaluate` already handles a
+        // scalar struct input.
+        let all_scalar =
+            !values.is_empty() && values.iter().all(|v| matches!(v, ColumnarValue::Scalar(_)));
         let arrays = ColumnarValue::values_to_arrays(&values)?;
         let fields = self.fields(&batch.schema())?;
-        Ok(ColumnarValue::Array(Arc::new(StructArray::new(
-            fields.into(),
-            arrays,
-            None,
-        ))))
+        let struct_array = StructArray::new(fields.into(), arrays, None);
+        if all_scalar {
+            Ok(ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(
+                struct_array,
+            ))))
+        } else {
+            Ok(ColumnarValue::Array(Arc::new(struct_array)))
+        }
     }
 
     fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {

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/DataTypeSupport.scala

@@ -80,6 +80,21 @@ object DataTypeSupport {
     case _ => false
   }
 
+  /**
+   * `dt` with every array/map/struct nullability flag forced to `true` at all nesting levels (map
+   * key fields stay non-null per Arrow's map invariant). Re-derives Spark's `private[spark]`
+   * `DataType.asNullable`, used as a common cast target to unify types whose Comet runtime
+   * nullability exceeds Spark's Catalyst nullability.
+   */
+  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 = deepNullable(f.dataType), nullable = true)))
+    case other => other
+  }
+
   def hasTemporalType(t: DataType): Boolean = t match {
     case DataTypes.DateType | DataTypes.TimestampType | DataTypes.TimestampNTZType =>
       true

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

@@ -531,6 +531,22 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
     }
   }
 
+  /**
+   * Attach a fresh `expr_id` and, when the expression's origin carries one, its `QueryContext` to
+   * `protoExpr`. Native ANSI errors resolve their context by the `expr_id` of the `Expr` that
+   * throws (see `register_query_context` and `ListExtract` in the native planner), so the
+   * metadata has to sit on that `Expr`. The generic serde path applies this to the top-level
+   * `Expr` it returns; a serde that nests a throwing expression inside a wrapper (e.g.
+   * `CometElementAt`'s CASE-WHEN NULL guard) must also call it on the inner `Expr`, or that
+   * expression's error renders without Spark's `== SQL ... ==` query context.
+   */
+  private[serde] def attachExprIdAndContext(expr: Expression, protoExpr: Expr): Expr = {
+    val builder = protoExpr.toBuilder
+    builder.setExprId(nextExprId())
+    extractQueryContext(expr).foreach(builder.setQueryContext)
+    builder.build()
+  }
+
   def supportedDataType(dt: DataType, allowComplex: Boolean = false): Boolean = dt match {
     case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType |
         _: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _: TimestampNTZType |
@@ -1004,12 +1020,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
           withNativeExpr(expr, CometExplainInfo.exprDisplayName(expr))
         }
         // Attach QueryContext and expr_id to the expression
-        val builder = protoExpr.toBuilder
-        builder.setExprId(nextExprId())
-        extractQueryContext(expr).foreach { ctx =>
-          builder.setQueryContext(ctx)
-        }
-        builder.build()
+        attachExprIdAndContext(expr, protoExpr)
       }
   }