diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index d109627e825..3536073daa0 100644 --- a/native/core/src/execution/planner.rs +++ b/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::, _>>()?; 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::, 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, schema: &SchemaRef, + widen: impl Fn(&DataType) -> DataType, ) -> Result, 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, + schema: &SchemaRef, + ) -> Result, ExecutionError> { + Self::coerce_child_to(child, schema, make_all_fields_nullable) + } + fn create_aggr_func_expr( name: &str, schema: SchemaRef, diff --git a/spark/src/main/java/org/apache/comet/codegen/CometBatchKernel.java b/spark/src/main/java/org/apache/comet/codegen/CometBatchKernel.java index a515cbe32d6..a7bdfe93cf1 100644 --- a/spark/src/main/java/org/apache/comet/codegen/CometBatchKernel.java +++ b/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; 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 748b1cee231..8800d7512e0 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/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")) } } diff --git a/spark/src/main/scala/org/apache/comet/serde/literals.scala b/spark/src/main/scala/org/apache/comet/serde/literals.scala index bac3631babe..d50a0af5d60 100644 --- a/spark/src/main/scala/org/apache/comet/serde/literals.scala +++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala @@ -22,17 +22,17 @@ package org.apache.comet.serde.literals import java.lang import org.apache.spark.internal.Logging -import org.apache.spark.sql.catalyst.expressions.{Attribute, Literal} -import org.apache.spark.sql.catalyst.util.ArrayData -import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, NullType, ShortType, StringType, TimestampNTZType, TimestampType} +import org.apache.spark.sql.catalyst.expressions.{Attribute, CreateArray, CreateMap, Expression, KnownNullable, Literal} +import org.apache.spark.sql.catalyst.util.{ArrayData, MapData, TypeUtils} +import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, StructType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.UTF8String import com.google.protobuf.ByteString import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.DataTypeSupport.isComplexType -import org.apache.comet.serde.{CometExpressionSerde, Compatible, ExprOuterClass, LiteralOuterClass, SupportLevel, Unsupported} -import org.apache.comet.serde.QueryPlanSerde.{isTimeType, serializeDataType, supportedDataType} +import org.apache.comet.serde.{CometExpressionSerde, Compatible, ExprOuterClass, LiteralOuterClass, MapKeySupport, SupportLevel, Unsupported} +import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, serializeDataType, supportedDataType} import org.apache.comet.serde.Types.ListLiteral object CometLiteral extends CometExpressionSerde[Literal] with Logging { @@ -55,6 +55,9 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { .elementType .isInstanceOf[ArrayType])))) { Compatible(None) + } else if (canExpandComplexLiteral(expr)) { + // Rebuilt as a `CreateArray` / `CreateMap` tree in `convert`. + Compatible(None) } else { expr.dataType match { case _: DayTimeIntervalType => Compatible(None) @@ -70,6 +73,13 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { val dataType = expr.dataType val value = expr.value + val expanded = expandComplexLiteral(expr) + if (expanded.isDefined) { + return exprToProtoInternal(expanded.get, inputs, binding).orElse { + withFallbackReason(expr, s"Unsupported data type $dataType") + None + } + } val exprBuilder = LiteralOuterClass.Literal.newBuilder() if (value == null) { @@ -215,4 +225,175 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { } listLiteralBuilder } + + /** + * Rebuild a folded complex Literal as an equivalent tree of `CreateArray` / `CreateMap` over + * primitive-typed Literals, or `None` when [[canExpandComplexLiteral]] rejects the shape. + * `getSupportLevel` probes cheaply with [[canExpandComplexLiteral]], which shares this method's + * admission checks, so the two cannot diverge and only `convert` materializes the tree. The + * native `Literal` proto carries scalars and nested `ListLiteral`s but no map values, so a + * Literal whose type contains a `MapType` has to be expanded before serialization. Teaching the + * proto to transport maps directly would remove the need for this rewrite and for the declines + * below: https://github.com/apache/datafusion-comet/issues/1937 + * + * The rebuilt tree is the tree Spark itself had before `ConstantFolding` collapsed it, down to + * every container's declared nullability (see [[withNullability]]). That equivalence is the + * safety property this rewrite rests on: whatever the rebuilt expression does natively is what + * the same query already does with `ConstantFolding` disabled, so expansion cannot introduce a + * folding-only behaviour difference. A non-nullable-value map into `map_entries` is handled by + * the native planner widening its argument rather than declined here. + * + * Map key semantics are declined here at admission rather than left to consumers, because not + * every native map consumer gates on [[MapKeySupport]]; see [[mapKeyTypesExpandable]] for why. + * + * Declined shapes: + * - Null values and empty top-level containers: a synthesized `Create*` with no children + * cannot recover the original element type. Empty `ArrayType` literals still serialize via + * `makeListLiteral`, which keeps the type. + * - Any map key type native map kernels cannot reproduce Spark equality for (floating-point, + * collated string, complex), or whose interpreted ordering is undefined + * (`CalendarIntervalType`), at any nesting level. See [[mapKeyTypesExpandable]]. + * - Arrays whose elements are structs, because [[needsExpansion]] does not walk into a + * `StructType`. Native `CreateNamedStruct` builds a 1-row `StructArray` whenever all of its + * children are scalars (`values_to_arrays`), which collides with the surrounding batch's + * row count, and its proto message carries no type, so Spark's declared field nullability + * cannot survive the wire either way. A struct that is only a map value is safe: + * `CometCreateMap` hands the whole rebuilt `CreateMap` to the JVM codegen dispatcher, so + * Spark's own code builds the struct. + * - Folded maps with duplicate keys, see [[hasDuplicateMapKeys]]. + */ + private def expandComplexLiteral(expr: Literal): Option[Expression] = { + if (!canExpandComplexLiteral(expr)) return None + expr.dataType match { + case ArrayType(et, containsNull) => + val arr = expr.value.asInstanceOf[ArrayData] + val elements = (0 until arr.numElements()) + .map(i => withNullability(literalAt(arr, i, et), containsNull)) + Some(CreateArray(elements, useStringTypeWhenEmpty = false)) + case MapType(kt, vt, valueContainsNull) => + val mapData = expr.value.asInstanceOf[MapData] + val keys = mapData.keyArray() + val values = mapData.valueArray() + val children = (0 until keys.numElements()).flatMap(i => + Seq( + literalAt(keys, i, kt), + withNullability(literalAt(values, i, vt), valueContainsNull))) + Some(CreateMap(children, useStringTypeWhenEmpty = false)) + case _ => None + } + } + + /** + * Cheap admission test that mirrors [[expandComplexLiteral]] without materializing the rebuilt + * `Create*` tree, so `getSupportLevel` can probe a large folded literal without allocating the + * N `Literal`s `convert` would immediately rebuild. Declines a null value or empty top-level + * container (no children to recover the element type), a folded map with duplicate keys (see + * [[hasDuplicateMapKeys]]), any unsupported or non-orderable map key type at any nesting level + * (see [[mapKeyTypesExpandable]]), and an array of structs ([[needsExpansion]] stops at a + * `StructType`). [[expandComplexLiteral]] gates on this, so the two cannot diverge. + */ + private def canExpandComplexLiteral(expr: Literal): Boolean = { + if (expr.value == null || !mapKeyTypesExpandable(expr.dataType)) return false + expr.dataType match { + case ArrayType(et, _) if needsExpansion(et) => + expr.value.asInstanceOf[ArrayData].numElements() > 0 + case MapType(kt, _, _) => + val mapData = expr.value.asInstanceOf[MapData] + mapData.numElements() > 0 && !hasDuplicateMapKeys(mapData.keyArray(), kt) + case _ => false + } + } + + /** + * True when every map key type reachable inside `dataType` can be rebuilt and consumed + * natively. A folded map is rebuilt as a `CreateMap`; once native it may reach a map consumer + * that has no key-type gate of its own -- `map_contains_key` lowers to + * `array_contains(map_keys(...), key)`, and neither `map_keys` nor `array_contains` consults + * [[MapKeySupport]]. A map that is only the value of another map is handed to the JVM + * dispatcher whole by `CometCreateMap` and never revisits this serde, so a nested unsupported + * key type would otherwise slip through. Mirror the [[MapKeySupport]] gate here, walking into + * map values, array elements, and struct fields, so an unsupported key type (floating-point, + * collated string, complex) declines expansion instead of reaching a native kernel whose key + * equality disagrees with Spark. + * + * Also decline a key type whose interpreted ordering is undefined -- one that recursively + * contains a `CalendarIntervalType`, whose `PhysicalCalendarIntervalType` ordering throws "does + * not support ordered operations". `ArrayBasedMapBuilder` dedups such keys by hash equality and + * never asks for an ordering, but [[hasDuplicateMapKeys]] compares keys through that ordering, + * so gating here keeps it from being called on a type it cannot order. Such literals fell back + * before this rewrite too, so declining them matches the prior behaviour. + */ + private def mapKeyTypesExpandable(dataType: DataType): Boolean = dataType match { + case MapType(kt, vt, _) => + MapKeySupport.keySupport(kt).isInstanceOf[Compatible] && + !SupportLevel.containsType(kt, classOf[CalendarIntervalType]) && + mapKeyTypesExpandable(vt) + case ArrayType(et, _) => mapKeyTypesExpandable(et) + case StructType(fields) => fields.forall(f => mapKeyTypesExpandable(f.dataType)) + case _ => true + } + + /** + * True when a Literal of this type has to be expanded rather than serialized, because the + * native `Literal` proto carries no map values. Walks array nesting only: an array of structs + * is not expandable (see [[expandComplexLiteral]]), so the walk stops at a `StructType`. + */ + private def needsExpansion(dataType: DataType): Boolean = dataType match { + case _: MapType => true + case ArrayType(et, _) => needsExpansion(et) + case _ => false + } + + /** Element `i` of `arr` as a Literal of type `dt`, or a null Literal for a null slot. */ + private def literalAt(arr: ArrayData, i: Int, dt: DataType): Literal = + Literal(if (arr.isNullAt(i)) null else arr.get(i, dt), dt) + + /** + * True when the folded map's key array holds duplicates under Spark's own key equality, in + * which case rebuilding the value as a `CreateMap` would change semantics the folded `MapData` + * had already settled: `ArrayBasedMapBuilder` throws `DUPLICATED_MAP_KEY` under + * `MAP_KEY_DEDUP_POLICY=EXCEPTION` and silently drops the earlier entry under `LAST_WIN`, where + * the original literal (from `from_json`, or a cast of one) kept both entries. + * + * Only reached for key types [[mapKeyTypesExpandable]] admits, i.e. atomic non-floating-point, + * default-collation `StringType`, and `BinaryType` -- floating-point, collated, complex, and + * `CalendarIntervalType` keys are already declined before this runs. Every admitted key type is + * orderable, so `TypeUtils.getInterpretedOrdering` (which `ArrayBasedMapBuilder` itself uses + * for `BinaryType`) is safe and matches Spark's equality: `Array[Byte]` keys compare by content + * rather than by identity, and the atomic orderings agree with hash equality. A null key cannot + * occur in a map Spark built, so treat one as a duplicate and keep the projection on Spark + * instead of asking the ordering to compare it. + */ + private def hasDuplicateMapKeys(keys: ArrayData, keyType: DataType): Boolean = { + val n = keys.numElements() + if (n < 2) { + false + } else if ((0 until n).exists(keys.isNullAt)) { + true + } else { + val ordering = TypeUtils.getInterpretedOrdering(keyType) + val sorted = (0 until n).map(i => keys.get(i, keyType)).sorted(ordering) + sorted.sliding(2).exists(pair => ordering.compare(pair.head, pair.last) == 0) + } + } + + /** + * Wrap `expr` in `KnownNullable` when the container it came out of declares its elements or + * values nullable, so the rebuilt `Create*` reports exactly the literal's declared + * `ArrayType.containsNull` / `MapType.valueContainsNull`. Both directions matter: + * - when the flag is true, a null slot yields a nullable `Literal` while a populated slot + * does not, and DataFusion `make_array` asserts strict Arrow-type equality across siblings + * (apache/datafusion#22366), so all children have to agree. + * - when the flag is false, widening it would make the rebuilt map disagree with a non-folded + * `CreateMap` sibling whose type Spark's own coercion had already accepted, which + * `CometCreateArray` no longer normalizes away because DataFusion 54.1 cannot coerce a + * `MapType.valueContainsNull` mismatch. + * + * A container that declares itself non-nullable never holds a null slot: `CreateArray` and + * `CreateMap` only report `false` when every child is non-nullable, and Spark's other sources + * of a complex type (a `CAST` target, a `from_json` schema, `MapType.apply`) default the flag + * to `true`. `CometKnownNullable` forwards the child on the wire, so runtime is unaffected. + */ + private def withNullability(expr: Expression, nullable: Boolean): Expression = + if (nullable && !expr.nullable) KnownNullable(expr) else expr } 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 c889284129f..51fa428b543 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -23,9 +23,61 @@ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ -import org.apache.comet.serde.QueryPlanSerde.{createBinaryExpr, exprToProtoInternal, scalarFunctionExprToProto} +import org.apache.comet.DataTypeSupport.isComplexType +import org.apache.comet.serde.QueryPlanSerde.{createBinaryExpr, exprToProtoInternal, hasNonDefaultStringCollation, scalarFunctionExprToProto} import org.apache.comet.shims.CometTypeShim +/** + * Shared gate for the native map kernels that compare a lookup key against a map's stored keys + * (`map_extract`, reached from both `GetMapValue` and `ElementAt`). + */ +private[serde] object MapKeySupport { + + private val floatingPointReason: String = + "Spark normalizes floating-point map keys, so `-0.0` matches a `+0.0` key and all `NaN`s " + + "match each other; Comet's native map lookup compares the raw Arrow values." + + private val collationReason: String = + "Comet's native map lookup compares string keys as `UTF8_BINARY` and cannot honour a " + + "non-default collation." + + private val complexKeyReason: String = + "Comet's native `map_extract` casts the lookup key to the map's exact Arrow key type, which " + + "cannot reproduce Spark's equality for a complex key type (for example a `NULL` inside the " + + "lookup key aborts the cast against a non-nullable nested component)." + + /** + * The `SupportLevel` for a map-consuming expression whose stored-key type is `keyType`. Spark + * finds a key with `TypeUtils.getInterpretedOrdering` over the keys `ArrayBasedMapBuilder` + * stored, having first normalized them, while native `map_extract` compares the Arrow values as + * they are, so decline the key types where those disagree: + * - `ArrayBasedMapBuilder` rewrites a `-0.0` key to `+0.0` and canonicalises `NaN`, and + * `nanSafeCompareDoubles` treats `-0.0` and `+0.0` as equal, so Spark answers a `-0.0` + * lookup from a `+0.0` key where native finds nothing. + * - a non-default collation compares under rules native applies as `UTF8_BINARY`. + * - for a complex key type, `map_extract`'s `coerce_types` returns the map's exact key field + * type, so Comet's planner casts the lookup key to it. That cannot reproduce Spark's + * interpreted-ordering equality, and a `NULL` inside the lookup key aborts the cast with + * `Non-nullable field of ListArray "item" cannot contain nulls` rather than missing the + * lookup. Declined for every complex key type, since which of these trips is a runtime + * property of the lookup. + * + * `BinaryType` keys need no decline: Arrow compares them by content, as Spark's ordering does. + * The floating-point and collation checks walk every nesting level of the key type. + */ + def keySupport(keyType: DataType): SupportLevel = { + if (SupportLevel.containsType(keyType, classOf[FloatType], classOf[DoubleType])) { + Unsupported(Some(floatingPointReason)) + } else if (hasNonDefaultStringCollation(keyType)) { + Unsupported(Some(collationReason)) + } else if (isComplexType(keyType)) { + Unsupported(Some(complexKeyReason)) + } else { + Compatible() + } + } +} + object CometMapKeys extends CometExpressionSerde[MapKeys] { override def convert( @@ -64,6 +116,11 @@ object CometMapValues extends CometExpressionSerde[MapValues] { object CometMapExtract extends CometExpressionSerde[GetMapValue] { + override def getSupportLevel(expr: GetMapValue): SupportLevel = expr.child.dataType match { + case MapType(keyType, _, _) => MapKeySupport.keySupport(keyType) + case _ => Compatible() + } + override def convert( expr: GetMapValue, inputs: Seq[Attribute], 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 5d59df4577e..0acc1f25de3 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 @@ -15,6 +15,8 @@ -- specific language governing permissions and limitations -- under the License. +-- ConfigMatrix: parquet.enable.dictionary=false,true + statement CREATE TABLE test_create_array(a int, b int, c int) USING parquet @@ -26,3 +28,107 @@ SELECT array(a, b, c) FROM test_create_array query SELECT array(1, 2, 3, NULL) + +-- Boundary integer values including int min/max. +query +SELECT array(2147483647, -2147483648, 0, NULL) + +-- Table for column-based coverage of map/struct/nested-array children. +statement +CREATE TABLE test_create_array_complex(k int, v int, s string, arr array) USING parquet + +statement +INSERT INTO test_create_array_complex VALUES + (1, 10, 'a', array(1, 2, 3)), + (2, 20, NULL, array(4, NULL, 6)), + (NULL, NULL, 'c', array()), + (NULL, NULL, NULL, NULL) + +-- Array of nested arrays. +query +SELECT array(array(1, 2), array(3, NULL)) + +query +SELECT array(array(a, b), array(b, c)) FROM test_create_array + +-- Array of maps built from primitive int values. The SQL-file harness always excludes +-- ConstantFolding, so every `map(...)` below stays a `CreateMap` and reaches native +-- `make_array` through the constructor path. Folded-literal expansion in `CometLiteral` is +-- covered by `CometArrayExpressionSuite` instead, where folding is left enabled. +query +SELECT array(map(1, 10), map(2, 20)) + +-- Array of maps whose values are arrays: `MapType(IntegerType, ArrayType(IntegerType, false), true)`. +query +SELECT array(map(1, array(1, 2, 3)), map(2, array(4, 5, 6))) + +-- Array of maps with string keys. +query +SELECT array(map('x', 1, 'y', 2), map('z', 3)) + +-- Array of maps whose children disagree on `valueContainsNull`. Spark's CreateArray coercion +-- compares element types with `sameType`, which ignores nullability, so it inserts no unifying +-- cast; it reports the nullability-merged map type as the array's element type. `make_array` +-- needs identical Arrow types, so `CometCreateArray` casts each child to that merged type +-- (`cast_map_to_map` widens `valueContainsNull`) and runs natively. +query +SELECT array(map('x', CAST(NULL AS INT), 'y', 2), map('z', 3)) + +-- Both children agree that the value is nullable, so this needs no cast. +query +SELECT array(map('x', CAST(NULL AS INT)), map('z', CAST(NULL AS INT))) + +-- Array of maps whose values are arrays that disagree only on the array's `containsNull`: +-- `map(1, array(1))` has value `ArrayType(IntegerType, containsNull = false)` while +-- `map(2, array(a))` has `ArrayType(IntegerType, true)`. The difference is container nullability +-- nested inside the map value, which Comet's map cast widens, so `CometCreateArray` casts each +-- child to the merged element type and runs natively. +query +SELECT array(map(1, array(1)), map(2, array(a))) FROM test_create_array + +-- Single-element array of map. +query +SELECT array(map(1, 2)) + +-- Array of arrays of maps: `ArrayType(ArrayType(MapType(IntegerType, StringType)))`. +query +SELECT array(array(map(1, 'a')), array(map(2, 'b'))) + +-- Array of structs. +query +SELECT array(named_struct('a', 1, 'b', 'x'), named_struct('a', 2, 'b', 'y')) + +-- Array of structs differing only in a field's nullability: the first `ct` is a non-null literal, +-- the second is a nullable CASE. Spark compares element types with `sameType` (nullability ignored) +-- so it keeps distinct StructTypes; `CometCreateArray` casts each child to the merged struct type +-- (widening `ct` to nullable) before `make_array`, so this runs natively. +query +SELECT array( + named_struct('id', a, 'ct', 'x'), + named_struct('id', a, 'ct', CASE WHEN a = 0 THEN 'y' END)) FROM test_create_array + +-- Same nested-field-nullability divergence wrapped in a map value. +query +SELECT array( + map('k', named_struct('id', a, 'ct', 'x')), + map('k', named_struct('id', a, 'ct', CASE WHEN a = 0 THEN 'y' END))) FROM test_create_array + +-- Column-based array of maps. Filter out NULL keys because Spark forbids them. +query +SELECT array(map(k, v)) FROM test_create_array_complex WHERE k IS NOT NULL + +-- Column-based array of structs. +query +SELECT array(named_struct('k', k, 's', s)) FROM test_create_array_complex + +-- Column-based array of nested arrays. +query +SELECT array(arr, array(k, v)) FROM test_create_array_complex + +-- Array of maps whose value is a nested array column. +query +SELECT array(map(k, arr)) FROM test_create_array_complex WHERE k IS NOT NULL + +-- Empty array cast to a nested array type: the element type has to survive with no children. +query +SELECT CAST(array() AS ARRAY>) diff --git a/spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql b/spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql index 5a838e260d4..b9e505336e2 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql @@ -51,3 +51,48 @@ SELECT element_at(mi, CAST(1 AS BIGINT)), element_at(mi, CAST(2 AS SMALLINT)) FR -- literal map arguments query SELECT element_at(map('a', 1, 'b', 2), 'a'), element_at(map('a', 1, 'b', 2), 'missing'), element_at(map('a', 1, 'b', 2), NULL) + +-- Map key types whose Spark equality Comet's native `map_extract` cannot reproduce fall back to +-- Spark. These stay on the constructor path here because the SQL harness excludes +-- `ConstantFolding`; `CometMapExpressionSuite` covers the folded-literal form of each. + +-- Spark stores `-0.0` map keys as `+0.0` and compares with `nanSafeCompareDoubles`, so a `-0.0` +-- lookup finds the `+0.0` key. Native lookup compares the raw Arrow values. +query expect_fallback(Spark normalizes floating-point map keys) +SELECT element_at(map(CAST(0 AS DOUBLE), 7), CAST(-0.0 AS DOUBLE)) + +query expect_fallback(Spark normalizes floating-point map keys) +SELECT element_at(map(CAST(0 AS FLOAT), 7), CAST(-0.0 AS FLOAT)) + +-- The floating-point decline walks every nesting level of the key type, so an array-of-double key +-- falls back for the same reason. +query expect_fallback(Spark normalizes floating-point map keys) +SELECT element_at(map(array(CAST(0 AS DOUBLE)), 7), array(CAST(-0.0 AS DOUBLE))) + +-- A complex key type: `map_extract` casts the lookup key to the map's exact Arrow key type, so a +-- NULL inside the lookup key would abort the cast instead of missing the lookup. +query expect_fallback(casts the lookup key to the map's exact Arrow key type) +SELECT element_at(map(array(1), 7), array(CAST(NULL AS INT))) + +query expect_fallback(casts the lookup key to the map's exact Arrow key type) +SELECT element_at(map(named_struct('a', 1), 7), named_struct('a', 1)) + +-- `BinaryType` keys need no decline: Arrow compares them by content, as Spark's ordering does. +query +SELECT element_at(map(CAST('a' AS BINARY), 1, CAST('b' AS BINARY), 2), CAST('b' AS BINARY)) + +-- Nested INT-keyed map: the inner `element_at` returns NULL for ids not in the outer map (2, 3), +-- and the outer `element_at` looks it up with a per-row key `id % (id - 2)`. This harness runs with +-- ANSI disabled, so the remainder-by-zero at id = 2 evaluates to NULL rather than throwing, and +-- `element_at(NULL_map, NULL)` returns NULL -- matching Spark's 7, NULL, NULL natively. (Under ANSI, +-- native scalar functions evaluate the key eagerly and throw where Spark short-circuits after the +-- NULL inner map; that is a pre-existing eager-evaluation difference in native `ElementAt`.) +statement +CREATE TABLE test_element_at_nested(id int) USING parquet + +statement +INSERT INTO test_element_at_nested VALUES (1), (2), (3) + +query +SELECT id, element_at(element_at(map(1, map(0, 7)), id), id % (id - 2)) AS v +FROM test_element_at_nested diff --git a/spark/src/test/resources/sql-tests/expressions/map/element_at_map_collation.sql b/spark/src/test/resources/sql-tests/expressions/map/element_at_map_collation.sql new file mode 100644 index 00000000000..763e1b31f12 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/element_at_map_collation.sql @@ -0,0 +1,35 @@ +-- 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. + +-- MinSparkVersion: 4.0 + +-- Spark 4.0+ supports string collations. Spark compares a map key under its declared collation, +-- so a `UTF8_LCASE` map matches `A1` against a dynamic `a1` lookup and returns `7`. Comet's native +-- `map_extract` compares string keys as `UTF8_BINARY`, so `CometElementAt` declines the lookup and +-- Spark evaluates the projection. The `element_at` form is used rather than `m[key]` because +-- Spark's `SimplifyExtractValueOps` rewrites `map(...)[key]` over a literal map into a `CASE` +-- before it can reach the native map lookup. + +statement +CREATE TABLE test_element_at_collation(k string) USING parquet + +statement +INSERT INTO test_element_at_collation VALUES ('a1'), ('A1'), ('zz'), (NULL) + +query expect_fallback(cannot honour a non-default collation) +SELECT element_at(map(CAST('A1' AS STRING COLLATE UTF8_LCASE), 7), CAST(k AS STRING COLLATE UTF8_LCASE)) +FROM test_element_at_collation diff --git a/spark/src/test/resources/sql-tests/expressions/map/get_map_value.sql b/spark/src/test/resources/sql-tests/expressions/map/get_map_value.sql index 4b497177cee..2c6c23846eb 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/get_map_value.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/get_map_value.sql @@ -30,3 +30,19 @@ SELECT m['x'], m['missing'] FROM test_map -- literal arguments query spark_answer_only SELECT map('a', 1, 'b', 2)['a'], map('a', 1, 'b', 2)['missing'], map('a', 1, 'b', 2)[NULL] + +-- Map key types whose Spark equality Comet's native `map_extract` cannot reproduce fall back to +-- Spark. `x[key]` routes through `GetMapValue` / `CometMapExtract` (the `element_at` form in +-- `element_at_map.sql` exercises the same guard on `CometElementAt`). A `map` column +-- is used rather than a `map(...)` literal because Spark's `SimplifyExtractValueOps` rewrites +-- `map(...)[key]` over a literal map into a `CASE` before it can reach the native lookup. Spark +-- stores a `-0.0` key as `+0.0` and finds it with `nanSafeCompareDoubles`, so it returns `7` for a +-- `-0.0` lookup; native lookup compares the raw Arrow values. +statement +CREATE TABLE test_map_double(m map) USING parquet + +statement +INSERT INTO test_map_double VALUES (map(CAST(0 AS DOUBLE), 7)), (map(CAST(1 AS DOUBLE), 8)), (NULL) + +query expect_fallback(Spark normalizes floating-point map keys) +SELECT m[CAST(-0.0 AS DOUBLE)] FROM test_map_double 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 91720cc2ab1..2d172f436f8 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 @@ -21,5 +21,45 @@ CREATE TABLE test_map_entries(m map) USING parquet statement INSERT INTO test_map_entries VALUES (map('a', 1, 'b', 2)), (map()), (NULL) -query spark_answer_only +query SELECT map_entries(m) FROM test_map_entries + +-- `MapType(_, _, valueContainsNull = false)`, which every `map(...)` over non-null values +-- produces. DataFusion's `map_entries` declares the entry `value` field nullable but reuses the +-- input map's entries array, so the planner widens the argument before the call. +query +SELECT map_entries(map(1, 2, 3, 4)) + +query +SELECT map_entries(map('a', array(1, 2))) + +-- A map nested in a map value keeps `valueContainsNull = false` through the extract. +query +SELECT map_entries(element_at(map(1, map(1, 2)), 1)) + +-- The `map_entries` argument is widened so its entry `value` field is nullable, but ONLY that outer +-- field: the nested `map(1, 2)` value must keep `valueContainsNull = false`. Extracting it with +-- `[0].value` and pairing it with the `map(2, coalesce(id, 0))` sibling would otherwise hit +-- `make_array` with unequal map types (one widened to `valueContainsNull = true`) and panic. +statement +CREATE TABLE test_map_entries_nested(id int) USING parquet + +statement +INSERT INTO test_map_entries_nested VALUES (1), (2), (3) + +query +SELECT array( + map_entries(map(1, IF(id = 1, map(1, 2), NULL)))[0].value, + map(2, coalesce(id, 0))) AS a +FROM test_map_entries_nested + +-- `map_entries` declares its entry `value` field nullable, so the extracted entry struct has a +-- nullable `value` while the sibling `named_struct('key', 1, 'value', id IS NOT NULL)` has a +-- non-nullable one. The two struct children differ only in that field's nullability, so +-- `CometCreateArray` casts each to the merged struct type before `make_array` and this runs +-- natively rather than panicking on the struct-field mismatch. +query +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 diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 05a6e8e650d..40505562300 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1134,42 +1134,160 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } - test("array of structs with nullability-divergent children") { - // Spark's type coercion compares element types with `sameType`, which ignores nullability, - // so two struct children that differ ONLY in a nested field's nullability get no unifying - // cast -- CreateArray keeps children of different StructTypes. DataFusion's make_array asserts - // strict element-type equality (down to nested nullability) and panics on the mismatch. Comet - // must decline this CreateArray so Spark's evaluator handles it. - withParquetTable((0 until 5).map(i => (i, i.toLong)), "tbl") { - val df = spark - .table("tbl") - .select( - array( - // ct is NOT NULL (literal) - struct(col("_1").as("id"), lit("a").as("ct")), - // ct is NULLABLE (when without otherwise) -- same type, different nullability - struct(col("_1").as("id"), when(col("_1") === 0, lit("b")).as("ct"))).as("arr")) - checkSparkAnswerAndFallbackReason(df, "CreateArray children have mismatched data types") + // The tests below deliberately live in this suite, not a `sql-tests` fixture: constant folding is + // enabled by default here, so each `map(...)` collapses to a MapType Literal and the outer + // `array(...)` reaches `CometCreateArray` with folded-Literal children, which `CometLiteral` + // rebuilds as an equivalent `CreateMap` of primitive literals -- the folded-literal expansion path + // under review. `CometSqlFileTestSuite` force-disables `ConstantFolding`, so an equivalent SQL + // fixture would only exercise the constructor path (which `create_array.sql` already covers). + test("array of folded map literals with array values (multirow)") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator( + "SELECT array(map(1, array(1, 2, 3)), map(2, array(4, 5, 6))) FROM tbl") } } - test("array of maps with nullability-divergent struct values") { - // Same nested-nullability divergence as the struct case, but wrapped in a MapType value so we - // exercise normalizeContainerNullability's MapType branch: the two map children share a surface - // type and differ only in a nested struct field's nullability, so they survive container - // (`MapType.valueContainsNull`) normalization as distinct types and CreateArray must still - // decline -- DataFusion's make_array would otherwise panic on the struct-field mismatch. - withParquetTable((0 until 5).map(i => (i, i.toLong)), "tbl") { - val df = spark - .table("tbl") - .select( - array( - // map value struct has ct NOT NULL (literal) - map(lit("k"), struct(col("_1").as("id"), lit("a").as("ct"))), - // map value struct has ct NULLABLE -- same type, different nested nullability - map(lit("k"), struct(col("_1").as("id"), when(col("_1") === 0, lit("b")).as("ct")))) - .as("arr")) - checkSparkAnswerAndFallbackReason(df, "CreateArray children have mismatched data types") + // A folded `map(1, array(1))` (value `ArrayType(IntegerType, containsNull = false)`) sits beside a + // dynamic `map(2, array(_1))` (value `ArrayType(IntegerType, true)`). Both maps still declare + // `valueContainsNull = false`, so only the nested array's `containsNull` differs. `CometCreateArray` + // casts each child to the nullability-merged element type (`cast_map_to_map` widens the nested + // array), so `make_array` sees identical Arrow types and this runs natively. + test( + "folded map with non-null nested array beside a dynamic map sibling runs natively (multirow)") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator( + "SELECT array(map(1, array(1)), map(2, array(_1))) AS a FROM tbl") + } + } + + test("array of folded map literals (multirow)") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator("SELECT array(map(1, 10), map(2, 20)) FROM tbl") + } + } + + // The folded `map(1, 2)` and the dynamic `map(2, coalesce(...))` both declare + // `MapType(IntegerType, IntegerType, valueContainsNull = false)`, so `CometCreateArray` admits + // the pair. Rebuilding the literal has to report that exact type: widening the rebuilt map's + // value to nullable would leave the sibling behind and `make_array` would panic, because + // DataFusion 54.1 cannot coerce a `MapType.valueContainsNull` mismatch away. + test("folded map literal keeps non-nullable values next to a dynamic map sibling (multirow)") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator( + "SELECT _1 AS id, array(map(1, 2), map(2, coalesce(_1, 0))) AS arr FROM tbl") + } + } + + // The whole `array(...)` folds to one `ArrayType(MapType(IntegerType, IntegerType, true))` + // literal, so every rebuilt element must report the unified nullable value type. + test("folded array of maps with a NULL value (multirow)") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator( + "SELECT array(map(1, CAST(NULL AS INT)), map(2, 3)) AS arr FROM tbl") + } + } + + // A NULL element sits next to a populated one inside a single folded + // `ArrayType(MapType(IntegerType, IntegerType, false), containsNull = true)` literal. The null + // slot serializes as a typed null literal carrying the declared map type, so the rebuilt sibling + // has to report that same type for `make_array` to accept the pair. + test("folded array mixing a map literal and a NULL element (multirow)") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator("SELECT array(map(1, 2), NULL) AS arr FROM tbl") + checkSparkAnswerAndOperator("SELECT array(NULL, map(1, 2)) AS arr FROM tbl") + } + } + + // A map value that is itself a map or an array of maps is built by Spark's own generated code + // inside `CometCreateMap`, so it keeps its declared nullability all the way to the consumer. + test("folded map literals with nested map values (multirow)") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator("SELECT _1 AS id, map(1, map(1, 2)) AS m FROM tbl") + checkSparkAnswerAndOperator("SELECT _1 AS id, map(1, array(map(2, 3))) AS m FROM tbl") + checkSparkAnswerAndOperator( + "SELECT _1 AS id, map(1, named_struct('a', 1, 'b', 'x')) AS m FROM tbl") + } + } + + // An empty container has no children to recover the element type from, and a NULL literal + // serializes as a typed null without expansion, so only the empty cases fall back. + test("folded empty and NULL complex literals") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, CAST(map() AS MAP) AS m FROM tbl", + "Unsupported data type MapType") + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, CAST(array() AS ARRAY>) AS a FROM tbl", + "Unsupported data type ArrayType") + checkSparkAnswerAndOperator("SELECT _1 AS id, CAST(NULL AS MAP) AS m FROM tbl") + checkSparkAnswerAndOperator( + "SELECT _1 AS id, CAST(NULL AS ARRAY>) AS a FROM tbl") + } + } + + // A folded struct literal would reach native `CreateNamedStruct` with all-scalar children, + // which builds a 1-row `StructArray` regardless of batch size. `CometLiteral` excludes + // StructType from expansion so the projection falls back instead of truncating the result. + test("folded struct literal in multirow projection falls back") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, named_struct('a', 1) AS s FROM tbl", + "Unsupported data type StructType") + } + } + + // Folds to one ArrayType(StructType) Literal whose structs disagree on field nullability. + // `KnownNullable` is dropped on the wire, so expansion could not make the sibling struct + // arrays agree and `make_array` would panic. Excluding StructType keeps this on Spark. + test("folded array of structs with divergent field nullability falls back (multirow)") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndFallbackReason( + "SELECT array(named_struct('a', CAST(NULL AS INT)), named_struct('a', 1)) AS arr FROM tbl", + "Unsupported data type ArrayType") + } + } + + // `from_json` can fold to a MapData with duplicate keys. Rebuilding it as a `CreateMap` would + // run `ArrayBasedMapBuilder`, which throws under `MAP_KEY_DEDUP_POLICY=EXCEPTION` and silently + // drops the earlier entry under `LAST_WIN`, so `CometLiteral` declines expansion under either + // policy and Spark evaluates the projection. + test("folded map literal with duplicate keys falls back (multirow)") { + assume(isSpark35Plus) + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + Seq("EXCEPTION", "LAST_WIN").foreach { policy => + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> policy) { + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, from_json('{\"a\":1,\"a\":2}', 'MAP') AS m FROM tbl", + "Unsupported data type MapType") + } + } + } + } + + // Spark's map cast preserves both entries, so the folded value still holds duplicate keys after + // the key type becomes BinaryType. `Array[Byte]` has no value-based `equals`, so the duplicate + // check has to compare keys the way `ArrayBasedMapBuilder` does, through + // `TypeUtils.getInterpretedOrdering`. + test("folded map literal with duplicate binary keys falls back (multirow)") { + assume(isSpark35Plus) + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + Seq("EXCEPTION", "LAST_WIN").foreach { policy => + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> policy) { + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, CAST(from_json('{\"a\":1,\"a\":2}', 'MAP') " + + "AS MAP) AS m FROM tbl", + "Unsupported data type MapType") + } + } + } + } + + // Distinct binary keys are fine: Arrow compares them by content, as Spark's ordering does. + test("folded map literal with distinct binary keys (multirow)") { + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator( + "SELECT _1 AS id, element_at(map(CAST('1' AS BINARY), 10, CAST('2' AS BINARY), 20), " + + "CAST(CAST(_1 AS STRING) AS BINARY)) AS v FROM tbl") } } diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 31c36a07c72..7311ef38d14 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -272,4 +272,182 @@ class CometMapExpressionSuite extends CometTestBase { } } + // ============================================================================================== + // Folded complex-literal tests. These live in this suite, not a `sql-tests` fixture, on purpose: + // they need `ConstantFolding` ON so `map(...)` / `array(...)` collapses to a `Literal` that + // `CometLiteral` then rebuilds -- the folded-literal expansion path under review. The SQL harness + // (`CometSqlFileTestSuite`) force-disables `ConstantFolding`, so an equivalent fixture would only + // exercise the constructor path. The `*.sql` files cover that constructor path where the same + // guard fires regardless of folding. + // ============================================================================================== + + // A map that reaches `map_entries` through an expression rather than a scan keeps + // `valueContainsNull = false`, which every `map(...)` over non-null values produces. DataFusion's + // `map_entries` declares the entry `value` field nullable but reuses the input map's entries + // array, so the planner widens the argument before the call. Here the inner `map(1, map(1, 2))` + // folds to a literal that `CometLiteral` rebuilds, and the outer `element_at` yields a + // non-nullable-value map straight into native `map_entries`. This exercises the widening on the + // default folding-on path; `map_entries.sql` covers the constructor path, where the harness + // excludes `ConstantFolding`. + test("map_entries on a folded non-nullable-value map (multirow)") { + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator( + "SELECT _1 AS id, map_entries(element_at(map(1, map(1, 2)), _1)) AS e FROM tbl") + } + } + + // Finding E: a map nested inside a map value is handed to the JVM dispatcher whole, so its double + // keys never revisit `CometLiteral`. The guard has to live on the lookup instead. The outer key + // type is INT, so inspecting only the outermost map would admit the query; the fallback comes + // from the inner `element_at`, whose child is `MapType(DoubleType, IntegerType)`. Constant folding + // is on here, which is the only configuration where the inner map becomes such a literal, so this + // regression cannot be expressed in a SQL fixture (the harness disables folding). The direct + // single-level double-key lookups live in `element_at_map.sql` / `get_map_value.sql`. + test("nested map lookup with floating-point keys falls back") { + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + val negZero = "CAST(concat('-', CAST(_1 - 1 AS STRING), '.0') AS DOUBLE)" + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, element_at(element_at(map(1, map(CAST(0 AS DOUBLE), 7)), _1), " + + s"$negZero) AS v FROM tbl", + "Spark normalizes floating-point map keys") + } + } + + // Finding E for a collated inner key. Same folding-on-only bypass as the double-key case above; + // the direct single-level collated lookup lives in `element_at_map_collation.sql`. + test("nested map lookup with collated string keys falls back") { + assume(isSpark40Plus) + withParquetTable(Seq(("a1", 0)), "tbl") { + checkSparkAnswerAndFallbackReason( + "SELECT element_at(element_at(map(1, map(CAST('A1' AS STRING COLLATE UTF8_LCASE), 7)), 1), " + + "CAST(_1 AS STRING COLLATE UTF8_LCASE)) AS v FROM tbl", + "cannot honour a non-default collation") + } + } + + // A folded map with `CalendarIntervalType` keys reaches `CometLiteral`. `ArrayBasedMapBuilder` + // dedups such keys by hash equality, but the folded-literal duplicate-key check needs an + // interpreted ordering and `PhysicalCalendarIntervalType` has none ("does not support ordered + // operations"). Expansion must decline these keys and keep the projection on Spark rather than + // crash planning by asking for that ordering; such literals fell back before this rewrite too. + test("folded map literal with calendar-interval keys falls back (multirow)") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, map(make_interval(1), 1, make_interval(2), 2) AS m FROM tbl", + "Unsupported data type MapType") + } + } + + // `map_entries` widens its argument so the entry `value` field is nullable, but it must widen ONLY + // that outer field. The outer `map(1, IF(...))` has `valueContainsNull = true`, and its value is a + // folded `map(1, 2)` with `valueContainsNull = false`. Widening the nested map too would extract a + // `valueContainsNull = true` map that disagrees with the dynamic `map(2, coalesce(...))` sibling, + // and native `make_array` would panic; the shallow widen keeps the nested type intact. + test("map_entries widening keeps nested map value nullability (multirow)") { + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator( + "SELECT _1 AS id, array(" + + "map_entries(map(1, IF(_1 = 1, map(1, 2), NULL)))[0].value, " + + "map(2, coalesce(_1, 0))) AS a FROM tbl") + } + } + + // `map_contains_key` lowers to `array_contains(map_keys(...), key)`; neither `map_keys` nor + // `array_contains` gates the key type, so a folded map with floating-point keys would reach them + // and compare `-0.0` against a normalized `+0.0` key bytewise, unlike Spark. The double-keyed map + // is nested inside an INT-keyed outer map, so the lookup guard on the outer `element_at` does not + // see it; expansion has to inspect nested map key types and decline. Spark returns `7, NULL, NULL`. + test("map_contains_key over nested floating-point map keys falls back (multirow)") { + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + val negZero = "CAST(concat('-', CAST(_1 - 1 AS STRING), '.0') AS DOUBLE)" + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, map_contains_key(" + + s"element_at(map(1, map(CAST(0 AS DOUBLE), 7)), _1), $negZero) AS present FROM tbl", + "Unsupported data type MapType") + } + } + + // The collated counterpart of the double-key `map_contains_key` case above: a nested + // `UTF8_LCASE`-keyed map would reach `array_contains` with bytewise comparison. Expansion declines + // the folded literal so the case-insensitive lookup stays on Spark. The outer lookup key is the + // dynamic `_1` so the inner map survives as a literal (a constant key would let Spark fold + // `map_keys` into a collated-string array literal instead, which never reaches this guard). + test("map_contains_key over nested collated map keys falls back (multirow)") { + assume(isSpark40Plus) + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, map_contains_key(" + + "element_at(map(1, map(CAST('A1' AS STRING COLLATE UTF8_LCASE), 7)), _1), " + + "CAST('a1' AS STRING COLLATE UTF8_LCASE)) AS present FROM tbl", + "Unsupported data type MapType") + } + } + + // Direct single-level folded map with floating-point keys: `map(CAST(0 AS DOUBLE), 7)` folds to a + // literal and `element_at` with a dynamic `-0.0` lookup must match Spark's `+0.0`-normalized key. + // Native `map_extract` compares raw Arrow values, so `MapKeySupport` declines it at `element_at`. + // Spark returns 7, NULL, NULL. (`element_at_map.sql` covers the folding-off constructor form.) + test("folded map literal with floating-point keys in element_at falls back (multirow)") { + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + val lookup = "CAST(concat('-', CAST(_1 - 1 AS STRING), '.0') AS DOUBLE)" + checkSparkAnswerAndFallbackReason( + s"SELECT _1 AS id, element_at(map(CAST(0 AS DOUBLE), 7), $lookup) AS v FROM tbl", + "Spark normalizes floating-point map keys") + } + } + + // Direct single-level folded map with collated string keys. Native lookup is bytewise, so a + // case-insensitive `a1` lookup against a stored `A1` cannot match; `MapKeySupport` declines it. + test("folded map literal with collated string keys in element_at falls back (multirow)") { + assume(isSpark40Plus) + withParquetTable(Seq(("a1", 0)), "tbl") { + checkSparkAnswerAndFallbackReason( + "SELECT element_at(map(CAST('A1' AS STRING COLLATE UTF8_LCASE), 7), " + + "CAST(_1 AS STRING COLLATE UTF8_LCASE)) AS v FROM tbl", + "cannot honour a non-default collation") + } + } + + // Folded map with a complex (array) key. Spark permits a dynamic lookup array containing a NULL + // element; native `map_extract` casts the lookup to the key's exact Arrow type and cannot + // reproduce Spark's equality, so `MapKeySupport` declines every complex key type. Spark returns + // 7, NULL, NULL. (`element_at_map.sql` covers the constructor form with a non-null lookup.) + test("folded map literal with complex array key in element_at falls back (multirow)") { + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, element_at(map(array(1), 7), " + + "array(IF(_1 = 2, CAST(NULL AS INT), _1))) AS v FROM tbl", + "casts the lookup key to the map's exact Arrow key type") + } + } + + // A folded nested INT-keyed map is admitted natively (all key-type guards pass), so the outer + // `element_at` runs as native `map_extract`. With ANSI disabled, a remainder-by-zero lookup key + // evaluates to NULL instead of throwing, so `map_extract(NULL_map, NULL)` returns NULL and the + // result matches Spark's 7, NULL, NULL. Under ANSI, native scalar functions evaluate the key + // eagerly and throw where Spark short-circuits after the NULL inner map -- a pre-existing eager + // evaluation difference in native `ElementAt`, not specific to folded literals. + test("folded nested map lookup with per-row key evaluation (multirow, non-ansi)") { + withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") { + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator( + "SELECT _1 AS id, element_at(element_at(map(1, map(0, 7)), _1), _1 % (_1 - 2)) AS v " + + "FROM tbl") + } + } + } + + // A large folded map (`map_from_arrays(sequence(...), sequence(...))` collapses to a many-entry + // MapType literal) rebuilds as a big `CreateMap` whose generated code Spark's codegen splits into + // nested helper classes. Those helpers read `CometBatchKernel.references`, which must be public + // (not protected) or the split code raises `IllegalAccessError` at runtime. 100k entries is large + // enough to force the split. + test("large folded map literal in element_at runs natively (multirow)") { + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator( + "SELECT _1 AS id, element_at(" + + "map_from_arrays(sequence(1, 100000), sequence(1, 100000)), _1) AS v FROM tbl") + } + } + }