From ce1cfe973ee916f6e8de6373b7b4a48f99bcd776 Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 24 Aug 2026 10:22:00 -0700 Subject: [PATCH 1/8] feat: support Map for `CreateArray` literal --- .../org/apache/comet/serde/literals.scala | 96 ++++++++++++++++++- .../expressions/array/create_array.sql | 87 +++++++++++++++++ 2 files changed, 179 insertions(+), 4 deletions(-) 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..6cb7a05122d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/literals.scala +++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala @@ -22,9 +22,10 @@ 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.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, CreateArray, CreateMap, CreateNamedStruct, Expression, KnownNullable, Literal} +import org.apache.spark.sql.catalyst.util.{ArrayData, MapData} +import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, 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 @@ -32,7 +33,7 @@ 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.QueryPlanSerde.{exprToProtoInternal, isTimeType, serializeDataType, supportedDataType} import org.apache.comet.serde.Types.ListLiteral object CometLiteral extends CometExpressionSerde[Literal] with Logging { @@ -55,6 +56,9 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { .elementType .isInstanceOf[ArrayType])))) { Compatible(None) + } else if (canExpandComplexLiteral(expr)) { + // Rebuilt as a Create[Array|Map|NamedStruct] tree in `convert`. + Compatible(None) } else { expr.dataType match { case _: DayTimeIntervalType => Compatible(None) @@ -70,6 +74,12 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { val dataType = expr.dataType val value = expr.value + if (canExpandComplexLiteral(expr)) { + return exprToProtoInternal(expandComplexLiteral(value, dataType), inputs, binding).orElse { + withFallbackReason(expr, s"Unsupported data type $dataType") + None + } + } val exprBuilder = LiteralOuterClass.Literal.newBuilder() if (value == null) { @@ -215,4 +225,82 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { } listLiteralBuilder } + + /** + * True when a non-null Literal of this type is not encodable in the native `Literal` proto and + * `convert` should try to rebuild it from primitive-typed Literals. The native proto today + * carries scalars and nested `ListLiteral`s (arrays of arrays / arrays of scalars). It does not + * carry Map or Struct values, so any Literal whose type contains a `MapType` or a `StructType` + * needs to be expanded before serialization. + */ + private def needsExpansion(dataType: DataType): Boolean = dataType match { + case _: MapType | _: StructType => true + case ArrayType(et, _) => needsExpansion(et) + case _ => false + } + + /** + * True when the Literal is a non-null complex value that we can rebuild from primitive Literals + * via [[expandComplexLiteral]]. Empty top-level containers are excluded because a synthesized + * `Create[Array|Map]` with no children cannot recover the original element type. + */ + private def canExpandComplexLiteral(expr: Literal): Boolean = { + val value = expr.value + if (value == null || !needsExpansion(expr.dataType)) return false + expr.dataType match { + case _: ArrayType => value.asInstanceOf[ArrayData].numElements() > 0 + case _: MapType => value.asInstanceOf[MapData].numElements() > 0 + case _: StructType => true + case _ => false + } + } + + /** + * Rebuild a folded complex Literal as an equivalent tree of `CreateArray` / `CreateMap` / + * `CreateNamedStruct` over primitive-typed Literals. Callers must gate on + * [[canExpandComplexLiteral]] so `dataType` is one of the three complex cases and top-level + * containers are non-empty. + */ + private def expandComplexLiteral(value: Any, dataType: DataType): Expression = + dataType match { + case ArrayType(et, _) => + val arr = value.asInstanceOf[ArrayData] + val elems = (0 until arr.numElements()).map { i => + if (arr.isNullAt(i)) Literal(null, et) + else asNullable(Literal(arr.get(i, et), et)) + } + CreateArray(elems, useStringTypeWhenEmpty = false) + case MapType(kt, vt, _) => + val mapData = value.asInstanceOf[MapData] + val keys = mapData.keyArray() + val vals = mapData.valueArray() + val children = (0 until keys.numElements()).flatMap { i => + val k = if (keys.isNullAt(i)) Literal(null, kt) else Literal(keys.get(i, kt), kt) + val v = + if (vals.isNullAt(i)) Literal(null, vt) + else asNullable(Literal(vals.get(i, vt), vt)) + Seq(k, v) + } + CreateMap(children, useStringTypeWhenEmpty = false) + case StructType(fields) => + val row = value.asInstanceOf[InternalRow] + val children = fields.zipWithIndex.flatMap { case (f, i) => + val v = + if (row.isNullAt(i)) Literal(null, f.dataType) + else asNullable(Literal(row.get(i, f.dataType), f.dataType)) + Seq(Literal(f.name), v) + } + CreateNamedStruct(children.toSeq) + case other => + throw new IllegalStateException(s"expandComplexLiteral called on $other") + } + + /** + * Wrap a non-null expression in `KnownNullable` so a surrounding `Create*` sees every sibling + * as nullable. DataFusion `make_array` asserts strict Arrow-type equality across siblings and + * would panic if a folded literal produced a non-nullable child next to a nullable one. + * `CometKnownNullable` forwards the child on the wire, so runtime is unaffected. + */ + private def asNullable(expr: Expression): Expression = + if (expr.nullable) expr else KnownNullable(expr) } 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..34fcda4a4ab 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,88 @@ 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. Under Spark constant folding this collapses +-- each `map(...)` to a MapType Literal; CometLiteral expands it back to a CreateMap of int +-- literals so the outer `array(...)` runs natively via `make_array`. +query +SELECT array(map(1, 10), map(2, 20)) + +-- Array of maps whose values are arrays. This is the exact fallback shape reported in +-- https://github.com/apache/datafusion-comet — `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 a NULL value inside (map value expansion must handle NULLs). +query +SELECT array(map('x', CAST(NULL AS INT), 'y', 2), map('z', 3)) + +-- Array containing a folded NULL map literal alongside a non-null map. +query +SELECT array(CAST(NULL AS MAP), map(1, 2)) + +-- Deeply nested folded complex literal: array of struct of map of array of struct. +query +SELECT array(named_struct('m', map(1, array(named_struct('id', 1, 's', 'x'))))) + +-- Array of maps with string keys. +query +SELECT array(map('x', 1, 'y', 2), map('z', 3)) + +-- Single-element array of map. +query +SELECT array(map(1, 2)) + +-- Array of arrays of maps (recursive expansion of ArrayType(ArrayType(MapType)) literal). +query +SELECT array(array(map(1, 'a')), array(map(2, 'b'))) + +-- Array of structs. Under constant folding each named_struct(...) folds to a StructType +-- Literal; expansion rebuilds it as CreateNamedStruct of primitive literals. +query +SELECT array(named_struct('a', 1, 'b', 'x'), named_struct('a', 2, 'b', 'y')) + +-- Column-based array of maps. With constant folding disabled by the SQL-file harness this +-- flows through CometCreateMap's codegen dispatch, not through Literal expansion, but it +-- must remain natively supported. 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 complex ArrayType literal takes the pre-expansion makeListLiteral path (empty +-- ListLiteral) so its element type survives even without any children. +query +SELECT CAST(array() AS ARRAY>) From f121109ef94abf027c1f29c8ddc9750dd6fb9e55 Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 24 Aug 2026 11:48:20 -0700 Subject: [PATCH 2/8] feat: support Map for `CreateArray` literal --- .../org/apache/comet/serde/literals.scala | 74 ++++++++++++------- .../expressions/array/create_array.sql | 20 +---- .../comet/CometArrayExpressionSuite.scala | 57 ++++++++++++++ 3 files changed, 110 insertions(+), 41 deletions(-) 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 6cb7a05122d..83490b6689f 100644 --- a/spark/src/main/scala/org/apache/comet/serde/literals.scala +++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala @@ -22,10 +22,9 @@ package org.apache.comet.serde.literals import java.lang import org.apache.spark.internal.Logging -import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, CreateArray, CreateMap, CreateNamedStruct, Expression, KnownNullable, Literal} +import org.apache.spark.sql.catalyst.expressions.{Attribute, CreateArray, CreateMap, Expression, KnownNullable, Literal} import org.apache.spark.sql.catalyst.util.{ArrayData, MapData} -import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, StructType, TimestampNTZType, TimestampType} +import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.UTF8String import com.google.protobuf.ByteString @@ -230,11 +229,21 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { * True when a non-null Literal of this type is not encodable in the native `Literal` proto and * `convert` should try to rebuild it from primitive-typed Literals. The native proto today * carries scalars and nested `ListLiteral`s (arrays of arrays / arrays of scalars). It does not - * carry Map or Struct values, so any Literal whose type contains a `MapType` or a `StructType` - * needs to be expanded before serialization. + * carry Map values, so any Literal whose type contains a `MapType` needs to be expanded before + * serialization. + * + * `StructType` (at any nesting depth) is deliberately excluded. Native `CometCreateNamedStruct` + * (`spark/src/main/scala/org/apache/comet/serde/structs.scala`) uses `values_to_arrays` in + * `native/spark-expr/src/struct_funcs/create_named_struct.rs`, which returns a 1-row + * `StructArray` whenever all children are scalar values. That collides with the row count of + * the surrounding batch and fails `make_array`'s length check. Field nullability is likewise + * inferred from the concrete child expressions, and `CometKnownNullable` + * (`spark/src/main/scala/org/apache/comet/serde/contraintExpressions.scala:99`) drops the tag + * on the wire, so wrapping a non-null child in `KnownNullable` does not carry across. Fall back + * to Spark for those shapes. */ private def needsExpansion(dataType: DataType): Boolean = dataType match { - case _: MapType | _: StructType => true + case _: MapType => true case ArrayType(et, _) => needsExpansion(et) case _ => false } @@ -242,24 +251,48 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { /** * True when the Literal is a non-null complex value that we can rebuild from primitive Literals * via [[expandComplexLiteral]]. Empty top-level containers are excluded because a synthesized - * `Create[Array|Map]` with no children cannot recover the original element type. + * `Create[Array|Map]` with no children cannot recover the original element type. A folded + * `MapData` with duplicate keys is also excluded: Spark's `CreateMap.eval` + * (`sql/catalyst/.../complexTypeCreator.scala:250`) feeds every entry through + * `ArrayBasedMapBuilder`, which throws under the default `MAP_KEY_DEDUP_POLICY=EXCEPTION` + * (`sql/catalyst/.../util/ArrayBasedMapBuilder.scala`). Rebuilding a folded literal that came + * from `from_json` or a similar source would then throw where the original literal had executed + * cleanly. */ private def canExpandComplexLiteral(expr: Literal): Boolean = { - val value = expr.value - if (value == null || !needsExpansion(expr.dataType)) return false + if (expr.value == null) return false expr.dataType match { - case _: ArrayType => value.asInstanceOf[ArrayData].numElements() > 0 - case _: MapType => value.asInstanceOf[MapData].numElements() > 0 - case _: StructType => true + case at: ArrayType if needsExpansion(at) => + expr.value.asInstanceOf[ArrayData].numElements() > 0 + case MapType(kt, _, _) => + val mapData = expr.value.asInstanceOf[MapData] + mapData.numElements() > 0 && !hasDuplicateMapKeys(mapData.keyArray(), kt) case _ => false } } /** - * Rebuild a folded complex Literal as an equivalent tree of `CreateArray` / `CreateMap` / - * `CreateNamedStruct` over primitive-typed Literals. Callers must gate on - * [[canExpandComplexLiteral]] so `dataType` is one of the three complex cases and top-level - * containers are non-empty. + * True when the folded map's key array contains any duplicate values. Uses Catalyst internal + * equality (`UTF8String`, `Decimal`, `ArrayData`, `InternalRow`, and boxed primitives all + * implement `equals` / `hashCode`), matching `ArrayBasedMapBuilder`'s own dedup semantics. + */ + private def hasDuplicateMapKeys(keys: ArrayData, keyType: DataType): Boolean = { + val n = keys.numElements() + if (n < 2) return false + val seen = new java.util.HashSet[Any](n) + var i = 0 + while (i < n) { + val k = if (keys.isNullAt(i)) null else keys.get(i, keyType) + if (!seen.add(k)) return true + i += 1 + } + false + } + + /** + * Rebuild a folded complex Literal as an equivalent tree of `CreateArray` / `CreateMap` over + * primitive-typed Literals. Callers must gate on [[canExpandComplexLiteral]] so `dataType` is + * one of the two complex cases and top-level containers are non-empty and have unique keys. */ private def expandComplexLiteral(value: Any, dataType: DataType): Expression = dataType match { @@ -282,15 +315,6 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { Seq(k, v) } CreateMap(children, useStringTypeWhenEmpty = false) - case StructType(fields) => - val row = value.asInstanceOf[InternalRow] - val children = fields.zipWithIndex.flatMap { case (f, i) => - val v = - if (row.isNullAt(i)) Literal(null, f.dataType) - else asNullable(Literal(row.get(i, f.dataType), f.dataType)) - Seq(Literal(f.name), v) - } - CreateNamedStruct(children.toSeq) case other => throw new IllegalStateException(s"expandComplexLiteral called on $other") } 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 34fcda4a4ab..ffcdc2d2620 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 @@ -57,23 +57,10 @@ SELECT array(array(a, b), array(b, c)) FROM test_create_array query SELECT array(map(1, 10), map(2, 20)) --- Array of maps whose values are arrays. This is the exact fallback shape reported in --- https://github.com/apache/datafusion-comet — `MapType(IntegerType, ArrayType(IntegerType, false), true)`. +-- 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 a NULL value inside (map value expansion must handle NULLs). -query -SELECT array(map('x', CAST(NULL AS INT), 'y', 2), map('z', 3)) - --- Array containing a folded NULL map literal alongside a non-null map. -query -SELECT array(CAST(NULL AS MAP), map(1, 2)) - --- Deeply nested folded complex literal: array of struct of map of array of struct. -query -SELECT array(named_struct('m', map(1, array(named_struct('id', 1, 's', 'x'))))) - -- Array of maps with string keys. query SELECT array(map('x', 1, 'y', 2), map('z', 3)) @@ -86,8 +73,9 @@ SELECT array(map(1, 2)) query SELECT array(array(map(1, 'a')), array(map(2, 'b'))) --- Array of structs. Under constant folding each named_struct(...) folds to a StructType --- Literal; expansion rebuilds it as CreateNamedStruct of primitive literals. +-- Array of structs. The SQL-file harness disables ConstantFolding, so `named_struct(...)` +-- stays as `CreateNamedStruct` (not a folded Literal) and this exercises the constructor +-- path, not `CometLiteral` expansion. query SELECT array(named_struct('a', 1, 'b', 'x'), named_struct('a', 2, 'b', 'y')) diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 05a6e8e650d..249076ee9fd 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1173,6 +1173,63 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } + // 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. + // `CometLiteral` rebuilds each folded map as an equivalent `CreateMap` of primitive literals. + test("array of folded map literals with array values (multirow)") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + val df = + spark.table("tbl").selectExpr("array(map(1, array(1, 2, 3)), map(2, array(4, 5, 6)))") + checkSparkAnswerAndOperator(df) + } + } + + test("array of folded map literals (multirow)") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator(spark.table("tbl").selectExpr("array(map(1, 10), map(2, 20))")) + } + } + + // P2 regression: standalone folded struct literal reaches native `CreateNamedStruct` with all + // scalar children. `values_to_arrays` returns a 1-row `StructArray` regardless of batch size, + // which fails length checks. `CometLiteral` excludes StructType from expansion so the projection + // falls back to Spark rather than producing a truncated result. + test("folded struct literal in multirow projection falls back") { + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + // Spark's `named_struct(...)` folds to a single struct Literal. Comet must not run + // `CreateNamedStruct` with all-scalar children against a 3-row batch. + val df = spark.table("tbl").selectExpr("_1 AS id", "named_struct('a', 1) AS s") + checkSparkAnswerAndFallbackReason(df, "Unsupported data type StructType") + } + } + + // P2 regression: `array(named_struct(...), named_struct(...))` with divergent field + // nullabilities folds to a single ArrayType(StructType) Literal. `KnownNullable` is dropped on + // the wire, so recursive expansion would produce sibling struct arrays with divergent per-field + // nullabilities 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") { + val df = spark + .table("tbl") + .selectExpr("array(named_struct('a', CAST(NULL AS INT)), named_struct('a', 1)) AS arr") + checkSparkAnswerAndFallbackReason(df, "Unsupported data type ArrayType") + } + } + + // P2 regression: `from_json` produces a MapData whose keys may be duplicated. Rebuilding via + // `CreateMap` invokes `ArrayBasedMapBuilder`, which throws under the default + // `MAP_KEY_DEDUP_POLICY=EXCEPTION` policy. `CometLiteral` inspects the folded map's key array + // and declines expansion when duplicates are present, letting Spark evaluate the projection. + test("folded map literal with duplicate keys falls back (multirow)") { + assume(isSpark35Plus) + withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { + val df = spark + .table("tbl") + .selectExpr("_1 AS id", "from_json('{\"a\":1,\"a\":2}', 'MAP') AS m") + checkSparkAnswerAndFallbackReason(df, "Unsupported data type MapType") + } + } + // Local table scan carries non-null array child fields (an in-memory Seq encodes // containsNull=false) into native kernels that promise nullable elements. ConvertToLocalRelation // must be disabled or the optimizer folds the expression at plan time and nothing runs natively. From 31f28b2e0ea5341716f2824be3b7a26bc8c82acb Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 24 Aug 2026 13:41:14 -0700 Subject: [PATCH 3/8] feat: support Map for `CreateArray` literal --- .../org/apache/comet/serde/literals.scala | 166 +++++++++--------- .../expressions/array/create_array.sql | 20 +-- .../comet/CometArrayExpressionSuite.scala | 70 +++++--- 3 files changed, 132 insertions(+), 124 deletions(-) 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 83490b6689f..c3e60b1dd6d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/literals.scala +++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala @@ -32,7 +32,7 @@ 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.{exprToProtoInternal, isTimeType, serializeDataType, supportedDataType} +import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, hasNonDefaultStringCollation, isTimeType, serializeDataType, supportedDataType} import org.apache.comet.serde.Types.ListLiteral object CometLiteral extends CometExpressionSerde[Literal] with Logging { @@ -55,8 +55,8 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { .elementType .isInstanceOf[ArrayType])))) { Compatible(None) - } else if (canExpandComplexLiteral(expr)) { - // Rebuilt as a Create[Array|Map|NamedStruct] tree in `convert`. + } else if (expandComplexLiteral(expr).isDefined) { + // Rebuilt as a `CreateArray` / `CreateMap` tree in `convert`. Compatible(None) } else { expr.dataType match { @@ -73,8 +73,9 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { val dataType = expr.dataType val value = expr.value - if (canExpandComplexLiteral(expr)) { - return exprToProtoInternal(expandComplexLiteral(value, dataType), inputs, binding).orElse { + val expanded = expandComplexLiteral(expr) + if (expanded.isDefined) { + return exprToProtoInternal(expanded.get, inputs, binding).orElse { withFallbackReason(expr, s"Unsupported data type $dataType") None } @@ -226,21 +227,58 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { } /** - * True when a non-null Literal of this type is not encodable in the native `Literal` proto and - * `convert` should try to rebuild it from primitive-typed Literals. The native proto today - * carries scalars and nested `ListLiteral`s (arrays of arrays / arrays of scalars). It does not - * carry Map values, so any Literal whose type contains a `MapType` needs to be expanded before - * serialization. + * Rebuild a folded complex Literal as an equivalent tree of `CreateArray` / `CreateMap` over + * primitive-typed Literals, or `None` when the shape cannot be rebuilt. 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 * - * `StructType` (at any nesting depth) is deliberately excluded. Native `CometCreateNamedStruct` - * (`spark/src/main/scala/org/apache/comet/serde/structs.scala`) uses `values_to_arrays` in - * `native/spark-expr/src/struct_funcs/create_named_struct.rs`, which returns a 1-row - * `StructArray` whenever all children are scalar values. That collides with the row count of - * the surrounding batch and fails `make_array`'s length check. Field nullability is likewise - * inferred from the concrete child expressions, and `CometKnownNullable` - * (`spark/src/main/scala/org/apache/comet/serde/contraintExpressions.scala:99`) drops the tag - * on the wire, so wrapping a non-null child in `KnownNullable` does not carry across. Fall back - * to Spark for those shapes. + * 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. + * - `StructType` at any depth. 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. Its proto message also carries no type, so Spark's + * declared field nullability cannot survive the wire either way. + * - Map key types whose Spark equality semantics native lookup cannot honor, see + * [[hasUnsafeMapKeyType]]. + * - Folded maps with duplicate keys. The rebuilt `CreateMap` evaluates through + * `ArrayBasedMapBuilder`, which throws under `MAP_KEY_DEDUP_POLICY=EXCEPTION`, where the + * original literal had already folded cleanly. + */ + private def expandComplexLiteral(expr: Literal): Option[Expression] = { + if (expr.value == null) return None + expr.dataType match { + case ArrayType(et, _) if needsExpansion(et) => + val arr = expr.value.asInstanceOf[ArrayData] + if (arr.numElements() == 0) { + None + } else { + val elements = (0 until arr.numElements()).map(i => asNullable(literalAt(arr, i, et))) + Some(CreateArray(elements, useStringTypeWhenEmpty = false)) + } + case MapType(kt, vt, _) => + val mapData = expr.value.asInstanceOf[MapData] + val keys = mapData.keyArray() + if (mapData.numElements() == 0 || hasUnsafeMapKeyType(kt) || + hasDuplicateMapKeys(keys, kt)) { + None + } else { + val values = mapData.valueArray() + val children = (0 until keys.numElements()).flatMap(i => + Seq(literalAt(keys, i, kt), asNullable(literalAt(values, i, vt)))) + Some(CreateMap(children, useStringTypeWhenEmpty = false)) + } + case _ => None + } + } + + /** + * 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: a `StructType` is not + * expandable at all (see [[expandComplexLiteral]]), so the walk stops there. */ private def needsExpansion(dataType: DataType): Boolean = dataType match { case _: MapType => true @@ -248,81 +286,41 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { case _ => false } + /** Element `i` of `arr`, or `null` for a null slot. */ + private def valueAt(arr: ArrayData, i: Int, dt: DataType): Any = + if (arr.isNullAt(i)) null else arr.get(i, dt) + + /** Element `i` of `arr` as a Literal of type `dt`. */ + private def literalAt(arr: ArrayData, i: Int, dt: DataType): Literal = + Literal(valueAt(arr, i, dt), dt) + /** - * True when the Literal is a non-null complex value that we can rebuild from primitive Literals - * via [[expandComplexLiteral]]. Empty top-level containers are excluded because a synthesized - * `Create[Array|Map]` with no children cannot recover the original element type. A folded - * `MapData` with duplicate keys is also excluded: Spark's `CreateMap.eval` - * (`sql/catalyst/.../complexTypeCreator.scala:250`) feeds every entry through - * `ArrayBasedMapBuilder`, which throws under the default `MAP_KEY_DEDUP_POLICY=EXCEPTION` - * (`sql/catalyst/.../util/ArrayBasedMapBuilder.scala`). Rebuilding a folded literal that came - * from `from_json` or a similar source would then throw where the original literal had executed - * cleanly. + * True when the map key type has Spark equality semantics that native map lookup (Arrow + * bytewise comparison) cannot honor, in which case declining expansion keeps the projection on + * Spark. `NormalizeFloatingNumbers` makes Spark treat `+0.0` and `-0.0` as the same key and + * canonicalises NaN, and a non-default collation compares under rules native applies as + * `UTF8_BINARY`. Both are checked at every nesting level of the key type. */ - private def canExpandComplexLiteral(expr: Literal): Boolean = { - if (expr.value == null) return false - expr.dataType match { - case at: ArrayType if needsExpansion(at) => - expr.value.asInstanceOf[ArrayData].numElements() > 0 - case MapType(kt, _, _) => - val mapData = expr.value.asInstanceOf[MapData] - mapData.numElements() > 0 && !hasDuplicateMapKeys(mapData.keyArray(), kt) - case _ => false - } - } + private def hasUnsafeMapKeyType(kt: DataType): Boolean = + hasNonDefaultStringCollation(kt) || + SupportLevel.containsType(kt, classOf[FloatType], classOf[DoubleType]) /** - * True when the folded map's key array contains any duplicate values. Uses Catalyst internal - * equality (`UTF8String`, `Decimal`, `ArrayData`, `InternalRow`, and boxed primitives all - * implement `equals` / `hashCode`), matching `ArrayBasedMapBuilder`'s own dedup semantics. + * True when the folded map's key array holds duplicates under Catalyst internal equality + * (`UTF8String`, `Decimal`, `ArrayData`, `InternalRow` and boxed primitives all implement + * `equals` / `hashCode`), which is what `ArrayBasedMapBuilder` uses for the key types reachable + * here. */ private def hasDuplicateMapKeys(keys: ArrayData, keyType: DataType): Boolean = { val n = keys.numElements() - if (n < 2) return false - val seen = new java.util.HashSet[Any](n) - var i = 0 - while (i < n) { - val k = if (keys.isNullAt(i)) null else keys.get(i, keyType) - if (!seen.add(k)) return true - i += 1 - } - false + n > 1 && (0 until n).map(i => valueAt(keys, i, keyType)).distinct.size < n } /** - * Rebuild a folded complex Literal as an equivalent tree of `CreateArray` / `CreateMap` over - * primitive-typed Literals. Callers must gate on [[canExpandComplexLiteral]] so `dataType` is - * one of the two complex cases and top-level containers are non-empty and have unique keys. - */ - private def expandComplexLiteral(value: Any, dataType: DataType): Expression = - dataType match { - case ArrayType(et, _) => - val arr = value.asInstanceOf[ArrayData] - val elems = (0 until arr.numElements()).map { i => - if (arr.isNullAt(i)) Literal(null, et) - else asNullable(Literal(arr.get(i, et), et)) - } - CreateArray(elems, useStringTypeWhenEmpty = false) - case MapType(kt, vt, _) => - val mapData = value.asInstanceOf[MapData] - val keys = mapData.keyArray() - val vals = mapData.valueArray() - val children = (0 until keys.numElements()).flatMap { i => - val k = if (keys.isNullAt(i)) Literal(null, kt) else Literal(keys.get(i, kt), kt) - val v = - if (vals.isNullAt(i)) Literal(null, vt) - else asNullable(Literal(vals.get(i, vt), vt)) - Seq(k, v) - } - CreateMap(children, useStringTypeWhenEmpty = false) - case other => - throw new IllegalStateException(s"expandComplexLiteral called on $other") - } - - /** - * Wrap a non-null expression in `KnownNullable` so a surrounding `Create*` sees every sibling - * as nullable. DataFusion `make_array` asserts strict Arrow-type equality across siblings and - * would panic if a folded literal produced a non-nullable child next to a nullable one. + * Wrap a non-null expression in `KnownNullable` so the surrounding `Create*` derives a nullable + * element / value type. DataFusion `make_array` asserts strict Arrow-type equality across + * siblings and would panic on a non-nullable child next to a nullable one. Same root cause as + * the mismatched-type decline in `CometCreateArray` (apache/datafusion#22366). * `CometKnownNullable` forwards the child on the wire, so runtime is unaffected. */ private def asNullable(expr: Expression): Expression = 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 ffcdc2d2620..42fa2309dfb 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 @@ -51,9 +51,10 @@ 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. Under Spark constant folding this collapses --- each `map(...)` to a MapType Literal; CometLiteral expands it back to a CreateMap of int --- literals so the outer `array(...)` runs natively via `make_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)) @@ -69,19 +70,15 @@ SELECT array(map('x', 1, 'y', 2), map('z', 3)) query SELECT array(map(1, 2)) --- Array of arrays of maps (recursive expansion of ArrayType(ArrayType(MapType)) literal). +-- Array of arrays of maps: `ArrayType(ArrayType(MapType(IntegerType, StringType)))`. query SELECT array(array(map(1, 'a')), array(map(2, 'b'))) --- Array of structs. The SQL-file harness disables ConstantFolding, so `named_struct(...)` --- stays as `CreateNamedStruct` (not a folded Literal) and this exercises the constructor --- path, not `CometLiteral` expansion. +-- Array of structs. query SELECT array(named_struct('a', 1, 'b', 'x'), named_struct('a', 2, 'b', 'y')) --- Column-based array of maps. With constant folding disabled by the SQL-file harness this --- flows through CometCreateMap's codegen dispatch, not through Literal expansion, but it --- must remain natively supported. Filter out NULL keys because Spark forbids them. +-- 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 @@ -97,7 +94,6 @@ SELECT array(arr, array(k, v)) FROM test_create_array_complex query SELECT array(map(k, arr)) FROM test_create_array_complex WHERE k IS NOT NULL --- Empty complex ArrayType literal takes the pre-expansion makeListLiteral path (empty --- ListLiteral) so its element type survives even without any children. +-- 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/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 249076ee9fd..9eb2003a50a 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1178,55 +1178,69 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp // `CometLiteral` rebuilds each folded map as an equivalent `CreateMap` of primitive literals. test("array of folded map literals with array values (multirow)") { withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { - val df = - spark.table("tbl").selectExpr("array(map(1, array(1, 2, 3)), map(2, array(4, 5, 6)))") - checkSparkAnswerAndOperator(df) + checkSparkAnswerAndOperator( + "SELECT array(map(1, array(1, 2, 3)), map(2, array(4, 5, 6))) FROM tbl") } } test("array of folded map literals (multirow)") { withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { - checkSparkAnswerAndOperator(spark.table("tbl").selectExpr("array(map(1, 10), map(2, 20))")) + checkSparkAnswerAndOperator("SELECT array(map(1, 10), map(2, 20)) FROM tbl") } } - // P2 regression: standalone folded struct literal reaches native `CreateNamedStruct` with all - // scalar children. `values_to_arrays` returns a 1-row `StructArray` regardless of batch size, - // which fails length checks. `CometLiteral` excludes StructType from expansion so the projection - // falls back to Spark rather than producing a truncated result. + // 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") { - // Spark's `named_struct(...)` folds to a single struct Literal. Comet must not run - // `CreateNamedStruct` with all-scalar children against a 3-row batch. - val df = spark.table("tbl").selectExpr("_1 AS id", "named_struct('a', 1) AS s") - checkSparkAnswerAndFallbackReason(df, "Unsupported data type StructType") + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, named_struct('a', 1) AS s FROM tbl", + "Unsupported data type StructType") } } - // P2 regression: `array(named_struct(...), named_struct(...))` with divergent field - // nullabilities folds to a single ArrayType(StructType) Literal. `KnownNullable` is dropped on - // the wire, so recursive expansion would produce sibling struct arrays with divergent per-field - // nullabilities and `make_array` would panic. Excluding StructType keeps this on Spark. + // 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") { - val df = spark - .table("tbl") - .selectExpr("array(named_struct('a', CAST(NULL AS INT)), named_struct('a', 1)) AS arr") - checkSparkAnswerAndFallbackReason(df, "Unsupported data type ArrayType") + checkSparkAnswerAndFallbackReason( + "SELECT array(named_struct('a', CAST(NULL AS INT)), named_struct('a', 1)) AS arr FROM tbl", + "Unsupported data type ArrayType") } } - // P2 regression: `from_json` produces a MapData whose keys may be duplicated. Rebuilding via - // `CreateMap` invokes `ArrayBasedMapBuilder`, which throws under the default - // `MAP_KEY_DEDUP_POLICY=EXCEPTION` policy. `CometLiteral` inspects the folded map's key array - // and declines expansion when duplicates are present, letting Spark evaluate the projection. + // `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`, so + // `CometLiteral` declines expansion 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") { - val df = spark - .table("tbl") - .selectExpr("_1 AS id", "from_json('{\"a\":1,\"a\":2}', 'MAP') AS m") - checkSparkAnswerAndFallbackReason(df, "Unsupported data type MapType") + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, from_json('{\"a\":1,\"a\":2}', 'MAP') AS m FROM tbl", + "Unsupported data type MapType") + } + } + + // Spark treats `+0.0` and `-0.0` as the same map key; native lookup compares bytes. + test("folded map literal with double keys falls back to Spark for +/-0.0 equality") { + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, element_at(map(CAST(0 AS DOUBLE), 7), " + + "CAST(concat('-', CAST(_1 - 1 AS STRING), '.0') AS DOUBLE)) AS v FROM tbl", + "Unsupported data type MapType") + } + } + + // Spark compares the key under its declared collation; native compares under UTF8_BINARY. + test("folded map literal with collated string keys falls back") { + 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", + "Unsupported data type MapType") } } From 97544ebe6e60ef558ebb5dc0f6b34aa7ec1ea242 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 25 Aug 2026 10:25:06 -0700 Subject: [PATCH 4/8] feat: support Map for `CreateArray` literal --- native/core/src/execution/planner.rs | 41 +++++-- .../scala/org/apache/comet/serde/arrays.scala | 36 +++--- .../org/apache/comet/serde/literals.scala | 109 ++++++++++------- .../scala/org/apache/comet/serde/maps.scala | 63 +++++++++- .../expressions/array/create_array.sql | 13 +++ .../expressions/map/element_at_map.sql | 24 ++++ .../sql-tests/expressions/map/map_entries.sql | 13 +++ .../comet/CometArrayExpressionSuite.scala | 110 ++++++++++++++---- .../comet/CometMapExpressionSuite.scala | 78 +++++++++++++ 9 files changed, 400 insertions(+), 87 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index d109627e825..25667a96027 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -2897,13 +2897,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 +3384,16 @@ impl PhysicalPlanner { .collect::, _>>()?; let fun_name = &expr.func; + // `map_entries` reuses the input map's entries array but declares the entry `value` field + // nullable, so widen the argument first or Arrow rejects the child type. See + // `coerce_child_fields_nullable`. + let args = if fun_name == "map_entries" { + args.into_iter() + .map(|arg| Self::coerce_child_fields_nullable(arg, &input_schema)) + .collect::, ExecutionError>>()? + } else { + args + }; let input_expr_types = args .iter() .map(|x| x.data_type(input_schema.as_ref())) @@ -3516,15 +3526,24 @@ 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` to the all-nullable variant of its own type, or returns it unchanged when it + /// already is. Used where a kernel's declared output type marks every nested field nullable + /// but its implementation reuses the child's arrays, so a non-nullable inner field makes the + /// declared and produced types disagree: + /// + /// - `collect_list` / `collect_set` 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"). + /// - `map_entries` declares its list element as `Struct(key non-null, value nullable)` but + /// reuses the input map's own entries array as the list values, so `ListArray::new` + /// rejects a map whose entry `value` field is non-nullable. Every Spark `map(...)` + /// constructor over non-null values produces exactly that type. + /// + /// `make_all_fields_nullable` keeps an Arrow map's key field non-nullable, so widening a map + /// stays a legal map type. + fn coerce_child_fields_nullable( child: Arc, schema: &SchemaRef, ) -> Result, ExecutionError> { 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..a730f5c47a2 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -474,11 +474,12 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] { // `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. + // nullability get no unifying cast. DataFusion coerces an `ArrayType.containsNull` mismatch + // away before that assert, but NOT a struct field's or a `MapType.valueContainsNull` + // mismatch, so `array(struct(a not null), struct(a nullable))` and + // `array(map('x', CAST(NULL AS INT)), map('z', 3))` both panic inside `make_array_inner`. + // Decline only those cases (i.e. children that still differ after normalizing the container + // nullability DataFusion does coerce) 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. @@ -502,20 +503,25 @@ 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. + * Rewrites a type so that `ArrayType.containsNull` is forced to `true` everywhere, while struct + * field nullability and `MapType.valueContainsNull` are left intact. Two CreateArray children + * whose types differ ONLY in `containsNull` are tolerated by DataFusion's `make_array`, because + * `type_union_resolution_coercion` routes `List` arguments through `list_coercion`, which + * merges the element field's nullability and lets Comet's planner insert the unifying cast. Its + * fallback chain has no `map_coercion` arm in DataFusion 54.1, so two `Map` arguments that + * differ in value nullability fail to unify, no cast is inserted, and `MutableArrayData` + * panics. Keep `valueContainsNull` significant here so those children are declined instead. The + * same is true of a struct field's nullability, which `coerce_struct_by_*` would merge but + * which `MutableArrayData` still rejects for the array's own element type. */ private def normalizeContainerNullability(dt: DataType): DataType = dt match { case ArrayType(elementType, _) => ArrayType(normalizeContainerNullability(elementType), containsNull = true) - case MapType(keyType, valueType, _) => + case MapType(keyType, valueType, valueContainsNull) => MapType( normalizeContainerNullability(keyType), normalizeContainerNullability(valueType), - valueContainsNull = true) + valueContainsNull) case StructType(fields) => StructType(fields.map(f => f.copy(dataType = normalizeContainerNullability(f.dataType)))) case other => other @@ -593,7 +599,11 @@ 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.unsupportedKeyTypeReason(keyType) match { + case Some(reason) => Unsupported(Some(reason)) + case None => Compatible() + } 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 c3e60b1dd6d..e11b79d9fe5 100644 --- a/spark/src/main/scala/org/apache/comet/serde/literals.scala +++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala @@ -23,7 +23,7 @@ import java.lang import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.expressions.{Attribute, CreateArray, CreateMap, Expression, KnownNullable, Literal} -import org.apache.spark.sql.catalyst.util.{ArrayData, MapData} +import org.apache.spark.sql.catalyst.util.{ArrayData, MapData, TypeUtils} import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.UTF8String @@ -32,7 +32,7 @@ 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.{exprToProtoInternal, hasNonDefaultStringCollation, isTimeType, serializeDataType, supportedDataType} +import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, serializeDataType, supportedDataType} import org.apache.comet.serde.Types.ListLiteral object CometLiteral extends CometExpressionSerde[Literal] with Logging { @@ -234,41 +234,49 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { * 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. Shapes the native map kernels cannot handle are therefore + * declined by their own serdes rather than here, see [[MapKeySupport]] and [[CometMapEntries]]. + * * 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. - * - `StructType` at any depth. 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. Its proto message also carries no type, so Spark's - * declared field nullability cannot survive the wire either way. - * - Map key types whose Spark equality semantics native lookup cannot honor, see - * [[hasUnsafeMapKeyType]]. - * - Folded maps with duplicate keys. The rebuilt `CreateMap` evaluates through - * `ArrayBasedMapBuilder`, which throws under `MAP_KEY_DEDUP_POLICY=EXCEPTION`, where the - * original literal had already folded cleanly. + * - 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 (expr.value == null) return None expr.dataType match { - case ArrayType(et, _) if needsExpansion(et) => + case ArrayType(et, containsNull) if needsExpansion(et) => val arr = expr.value.asInstanceOf[ArrayData] if (arr.numElements() == 0) { None } else { - val elements = (0 until arr.numElements()).map(i => asNullable(literalAt(arr, i, et))) + val elements = (0 until arr.numElements()) + .map(i => withNullability(literalAt(arr, i, et), containsNull)) Some(CreateArray(elements, useStringTypeWhenEmpty = false)) } - case MapType(kt, vt, _) => + case MapType(kt, vt, valueContainsNull) => val mapData = expr.value.asInstanceOf[MapData] val keys = mapData.keyArray() - if (mapData.numElements() == 0 || hasUnsafeMapKeyType(kt) || - hasDuplicateMapKeys(keys, kt)) { + if (mapData.numElements() == 0 || hasDuplicateMapKeys(keys, kt)) { None } else { val values = mapData.valueArray() val children = (0 until keys.numElements()).flatMap(i => - Seq(literalAt(keys, i, kt), asNullable(literalAt(values, i, vt)))) + Seq( + literalAt(keys, i, kt), + withNullability(literalAt(values, i, vt), valueContainsNull))) Some(CreateMap(children, useStringTypeWhenEmpty = false)) } case _ => None @@ -277,8 +285,8 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { /** * 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: a `StructType` is not - * expandable at all (see [[expandComplexLiteral]]), so the walk stops there. + * 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 @@ -295,34 +303,51 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { Literal(valueAt(arr, i, dt), dt) /** - * True when the map key type has Spark equality semantics that native map lookup (Arrow - * bytewise comparison) cannot honor, in which case declining expansion keeps the projection on - * Spark. `NormalizeFloatingNumbers` makes Spark treat `+0.0` and `-0.0` as the same key and - * canonicalises NaN, and a non-default collation compares under rules native applies as - * `UTF8_BINARY`. Both are checked at every nesting level of the key type. - */ - private def hasUnsafeMapKeyType(kt: DataType): Boolean = - hasNonDefaultStringCollation(kt) || - SupportLevel.containsType(kt, classOf[FloatType], classOf[DoubleType]) - - /** - * True when the folded map's key array holds duplicates under Catalyst internal equality - * (`UTF8String`, `Decimal`, `ArrayData`, `InternalRow` and boxed primitives all implement - * `equals` / `hashCode`), which is what `ArrayBasedMapBuilder` uses for the key types reachable - * here. + * 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. + * + * `ArrayBasedMapBuilder` hashes most atomic keys but compares `BinaryType`, collated + * `StringType` and complex keys through `TypeUtils.getInterpretedOrdering`, and it normalizes + * `-0.0` and `NaN` before either. The interpreted ordering is the one comparison that agrees + * with all of those: `Array[Byte]` keys compare by content rather than by identity, a collated + * `StringType` compares under its collation, and `nanSafeCompareDoubles` already treats `-0.0` + * and `+0.0`, and any two `NaN`s, as equal. 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() - n > 1 && (0 until n).map(i => valueAt(keys, i, keyType)).distinct.size < n + 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 a non-null expression in `KnownNullable` so the surrounding `Create*` derives a nullable - * element / value type. DataFusion `make_array` asserts strict Arrow-type equality across - * siblings and would panic on a non-nullable child next to a nullable one. Same root cause as - * the mismatched-type decline in `CometCreateArray` (apache/datafusion#22366). - * `CometKnownNullable` forwards the child on the wire, so runtime is unaffected. + * 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 asNullable(expr: Expression): Expression = - if (expr.nullable) expr else KnownNullable(expr) + 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..1195f3aaeb4 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 { + + 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." + + val collationReason: String = + "Comet's native map lookup compares string keys as `UTF8_BINARY` and cannot honour a " + + "non-default collation." + + val complexKeyReason: String = + "Comet's native `map_extract` casts the lookup key to the map's exact Arrow key type, so a " + + "complex key type whose components are declared non-nullable rejects a lookup key that " + + "holds a `NULL` inside one of them." + + /** + * The reason Comet's native map lookup cannot reproduce Spark's key equality for `keyType`, or + * `None` when it can. 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: + * - `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. When the key type declares a nested + * component non-nullable, a `NULL` inside the lookup key aborts that cast with + * `Non-nullable field of ListArray "item" cannot contain nulls` rather than missing the + * lookup, which is what Spark does. + * + * `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, so they + * keep holding if native `map_extract` ever gains complex-key support. + */ + def unsupportedKeyTypeReason(keyType: DataType): Option[String] = { + if (SupportLevel.containsType(keyType, classOf[FloatType], classOf[DoubleType])) { + Some(floatingPointReason) + } else if (hasNonDefaultStringCollation(keyType)) { + Some(collationReason) + } else if (isComplexType(keyType)) { + Some(complexKeyReason) + } else { + None + } + } +} + object CometMapKeys extends CometExpressionSerde[MapKeys] { override def convert( @@ -64,6 +116,15 @@ object CometMapValues extends CometExpressionSerde[MapValues] { object CometMapExtract extends CometExpressionSerde[GetMapValue] { + override def getSupportLevel(expr: GetMapValue): SupportLevel = expr.child.dataType match { + case MapType(keyType, _, _) => + MapKeySupport.unsupportedKeyTypeReason(keyType) match { + case Some(reason) => Unsupported(Some(reason)) + case None => Compatible() + } + 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 42fa2309dfb..5dfda61fabd 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 @@ -66,6 +66,19 @@ SELECT array(map(1, array(1, 2, 3)), map(2, array(4, 5, 6))) 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. DataFusion 54.1 cannot unify two Map arguments either: the fallback chain in +-- `type_union_resolution_coercion` has no `map_coercion` arm, so `make_array` receives the two +-- differently typed map arrays and `MutableArrayData` panics. `CometCreateArray` therefore keeps +-- `MapType.valueContainsNull` significant when it normalizes container nullability and declines. +query expect_fallback(CreateArray children have mismatched data types) +SELECT array(map('x', CAST(NULL AS INT), 'y', 2), map('z', 3)) + +-- Both children agree that the value is nullable, so this stays native. +query +SELECT array(map('x', CAST(NULL AS INT)), map('z', CAST(NULL AS INT))) + -- Single-element array of map. query SELECT array(map(1, 2)) 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..5bb322b7f32 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,27 @@ 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)) + +-- 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)) 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..1b1cd26eccf 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 @@ -23,3 +23,16 @@ INSERT INTO test_map_entries VALUES (map('a', 1, 'b', 2)), (map()), (NULL) query spark_answer_only 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)) diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 9eb2003a50a..cd3e40b8f6d 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1189,6 +1189,65 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } + // 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. @@ -1212,35 +1271,46 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } // `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`, so - // `CometLiteral` declines expansion and Spark evaluates the projection. + // 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") { - checkSparkAnswerAndFallbackReason( - "SELECT _1 AS id, from_json('{\"a\":1,\"a\":2}', 'MAP') AS m FROM tbl", - "Unsupported data type MapType") + 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 treats `+0.0` and `-0.0` as the same map key; native lookup compares bytes. - test("folded map literal with double keys falls back to Spark for +/-0.0 equality") { - withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { - checkSparkAnswerAndFallbackReason( - "SELECT _1 AS id, element_at(map(CAST(0 AS DOUBLE), 7), " + - "CAST(concat('-', CAST(_1 - 1 AS STRING), '.0') AS DOUBLE)) AS v 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") + } + } } } - // Spark compares the key under its declared collation; native compares under UTF8_BINARY. - test("folded map literal with collated string keys falls back") { - 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", - "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..ad3769c1a41 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -272,4 +272,82 @@ class CometMapExpressionSuite extends CometTestBase { } } + // 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. Both queries below depend on a + // column, so `ConstantFolding` cannot collapse the whole `map_entries(...)` call into an + // `ArrayType(StructType)` literal that `CometLiteral` would decline. `map_entries.sql` covers the + // all-constant shapes, where the harness excludes `ConstantFolding`. + test("map_entries on a non-nullable-value map expression (multirow)") { + withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { + checkSparkAnswerAndOperator("SELECT _1 AS id, map_entries(map(_1, 2)) AS e FROM tbl") + checkSparkAnswerAndOperator( + "SELECT _1 AS id, map_entries(element_at(map(1, map(1, 2)), _1)) AS e FROM tbl") + } + } + + // Spark stores a `-0.0` map key as `+0.0` and finds it with `nanSafeCompareDoubles`, which + // treats `-0.0` and `+0.0` as equal. Native `map_extract` compares the raw Arrow values, so + // `CometMapExtract` / `CometElementAt` decline the lookup. The map here folds to a literal, so + // the decline has to come from the lookup rather than from `CometLiteral`, which cannot see how + // a map nested in a map value will be consumed. + test("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( + s"SELECT _1 AS id, element_at(map(CAST(0 AS DOUBLE), 7), $negZero) AS v FROM tbl", + "Spark normalizes floating-point map keys") + checkSparkAnswerAndFallbackReason( + s"SELECT _1 AS id, map(CAST(0 AS DOUBLE), 7)[$negZero] AS v FROM tbl", + "Spark normalizes floating-point map keys") + // The outer key type is INT, so only inspecting it would admit the lookup; the inner map is + // handed to the JVM dispatcher whole and never revisits `CometLiteral`. + 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") + } + } + + // Spark compares the key under its declared collation; native compares under UTF8_BINARY. + test("map lookup with collated string keys falls back") { + 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") + 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") + } + } + + // `map_extract`'s `coerce_types` returns the map's exact Arrow key field type, so Comet casts the + // lookup key to it. A NULL inside the lookup key then aborts the cast with `Non-nullable field of + // ListArray "item" cannot contain nulls` instead of missing the lookup, which is what Spark does. + test("map lookup with complex keys falls back") { + 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") + checkSparkAnswerAndFallbackReason( + "SELECT _1 AS id, element_at(map(named_struct('a', 1), 7), " + + "named_struct('a', _1)) AS v FROM tbl", + "casts the lookup key to the map's exact Arrow key type") + } + } + + // `BinaryType` keys need no decline: Arrow compares them by content, as Spark's ordering does. + test("map lookup with binary keys stays native") { + 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") + } + } + } From fecf175ab8f7ad434fa53a255851d4df90349728 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 25 Aug 2026 12:52:27 -0700 Subject: [PATCH 5/8] feat: support Map for `CreateArray` literal --- .../scala/org/apache/comet/serde/arrays.scala | 11 +-- .../org/apache/comet/serde/literals.scala | 55 +++++++++------ .../scala/org/apache/comet/serde/maps.scala | 46 ++++++------- .../map/element_at_map_collation.sql | 35 ++++++++++ .../expressions/map/get_map_value.sql | 16 +++++ .../sql-tests/expressions/map/map_entries.sql | 2 +- .../comet/CometMapExpressionSuite.scala | 68 +++++-------------- 7 files changed, 130 insertions(+), 103 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/map/element_at_map_collation.sql 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 a730f5c47a2..502ee85e8cf 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -513,6 +513,11 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] { * panics. Keep `valueContainsNull` significant here so those children are declined instead. The * same is true of a struct field's nullability, which `coerce_struct_by_*` would merge but * which `MutableArrayData` still rejects for the array's own element type. + * + * TODO: DataFusion 55.0.0 adds a `map_coercion` arm to `type_union_resolution_coercion` + * (apache/datafusion#23521), which coerces a `MapType.valueContainsNull` mismatch away. When + * Comet upgrades onto it, stop keeping `valueContainsNull` significant here or this decline + * stays correct but grows over-conservative, falling maps back that DataFusion can now unify. */ private def normalizeContainerNullability(dt: DataType): DataType = dt match { case ArrayType(elementType, _) => @@ -599,11 +604,7 @@ object CometElementAt extends CometExpressionSerde[ElementAt] { override def getSupportLevel(expr: ElementAt): SupportLevel = { expr.left.dataType match { case _: ArrayType => Compatible() - case MapType(keyType, _, _) => - MapKeySupport.unsupportedKeyTypeReason(keyType) match { - case Some(reason) => Unsupported(Some(reason)) - case None => 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 e11b79d9fe5..19bd20c8086 100644 --- a/spark/src/main/scala/org/apache/comet/serde/literals.scala +++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala @@ -55,7 +55,7 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { .elementType .isInstanceOf[ArrayType])))) { Compatible(None) - } else if (expandComplexLiteral(expr).isDefined) { + } else if (canExpandComplexLiteral(expr)) { // Rebuilt as a `CreateArray` / `CreateMap` tree in `convert`. Compatible(None) } else { @@ -239,7 +239,9 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { * 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. Shapes the native map kernels cannot handle are therefore - * declined by their own serdes rather than here, see [[MapKeySupport]] and [[CometMapEntries]]. + * declined by their consumers rather than here: a map lookup with an unsupported key type by + * [[MapKeySupport]], and a non-nullable-value map into `map_entries` by the native planner + * widening its argument. * * Declined shapes: * - Null values and empty top-level containers: a synthesized `Create*` with no children @@ -255,34 +257,45 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { * - Folded maps with duplicate keys, see [[hasDuplicateMapKeys]]. */ private def expandComplexLiteral(expr: Literal): Option[Expression] = { - if (expr.value == null) return None + if (!canExpandComplexLiteral(expr)) return None expr.dataType match { - case ArrayType(et, containsNull) if needsExpansion(et) => + case ArrayType(et, containsNull) => val arr = expr.value.asInstanceOf[ArrayData] - if (arr.numElements() == 0) { - None - } else { - val elements = (0 until arr.numElements()) - .map(i => withNullability(literalAt(arr, i, et), containsNull)) - Some(CreateArray(elements, useStringTypeWhenEmpty = false)) - } + 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() - if (mapData.numElements() == 0 || hasDuplicateMapKeys(keys, kt)) { - None - } else { - 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)) - } + 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 } } + /** + * True when [[expandComplexLiteral]] can rebuild this folded literal, tested without allocating + * the rebuilt tree so `getSupportLevel` can probe cheaply (`convert` runs the build only once, + * on the value it keeps). Mirrors the shapes [[expandComplexLiteral]] declines: a non-null + * value, a non-empty top-level container, an array whose elements need expansion, and a map + * with no duplicate keys. + */ + private def canExpandComplexLiteral(expr: Literal): Boolean = { + if (expr.value == null) 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 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 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 1195f3aaeb4..51fa428b543 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -33,47 +33,47 @@ import org.apache.comet.shims.CometTypeShim */ private[serde] object MapKeySupport { - val floatingPointReason: String = + 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." - val collationReason: String = + private val collationReason: String = "Comet's native map lookup compares string keys as `UTF8_BINARY` and cannot honour a " + "non-default collation." - val complexKeyReason: String = - "Comet's native `map_extract` casts the lookup key to the map's exact Arrow key type, so a " + - "complex key type whose components are declared non-nullable rejects a lookup key that " + - "holds a `NULL` inside one of them." + 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 reason Comet's native map lookup cannot reproduce Spark's key equality for `keyType`, or - * `None` when it can. 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: + * 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. When the key type declares a nested - * component non-nullable, a `NULL` inside the lookup key aborts that cast with + * 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, which is what Spark does. + * 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, so they - * keep holding if native `map_extract` ever gains complex-key support. + * The floating-point and collation checks walk every nesting level of the key type. */ - def unsupportedKeyTypeReason(keyType: DataType): Option[String] = { + def keySupport(keyType: DataType): SupportLevel = { if (SupportLevel.containsType(keyType, classOf[FloatType], classOf[DoubleType])) { - Some(floatingPointReason) + Unsupported(Some(floatingPointReason)) } else if (hasNonDefaultStringCollation(keyType)) { - Some(collationReason) + Unsupported(Some(collationReason)) } else if (isComplexType(keyType)) { - Some(complexKeyReason) + Unsupported(Some(complexKeyReason)) } else { - None + Compatible() } } } @@ -117,11 +117,7 @@ object CometMapValues extends CometExpressionSerde[MapValues] { object CometMapExtract extends CometExpressionSerde[GetMapValue] { override def getSupportLevel(expr: GetMapValue): SupportLevel = expr.child.dataType match { - case MapType(keyType, _, _) => - MapKeySupport.unsupportedKeyTypeReason(keyType) match { - case Some(reason) => Unsupported(Some(reason)) - case None => Compatible() - } + case MapType(keyType, _, _) => MapKeySupport.keySupport(keyType) case _ => Compatible() } 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 1b1cd26eccf..46e4f7c111a 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,7 +21,7 @@ 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 diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index ad3769c1a41..d91dd671b82 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -275,34 +275,28 @@ class CometMapExpressionSuite extends CometTestBase { // 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. Both queries below depend on a - // column, so `ConstantFolding` cannot collapse the whole `map_entries(...)` call into an - // `ArrayType(StructType)` literal that `CometLiteral` would decline. `map_entries.sql` covers the - // all-constant shapes, where the harness excludes `ConstantFolding`. - test("map_entries on a non-nullable-value map expression (multirow)") { + // 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(map(_1, 2)) AS e FROM tbl") checkSparkAnswerAndOperator( "SELECT _1 AS id, map_entries(element_at(map(1, map(1, 2)), _1)) AS e FROM tbl") } } - // Spark stores a `-0.0` map key as `+0.0` and finds it with `nanSafeCompareDoubles`, which - // treats `-0.0` and `+0.0` as equal. Native `map_extract` compares the raw Arrow values, so - // `CometMapExtract` / `CometElementAt` decline the lookup. The map here folds to a literal, so - // the decline has to come from the lookup rather than from `CometLiteral`, which cannot see how - // a map nested in a map value will be consumed. - test("map lookup with floating-point keys falls back") { + // 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( - s"SELECT _1 AS id, element_at(map(CAST(0 AS DOUBLE), 7), $negZero) AS v FROM tbl", - "Spark normalizes floating-point map keys") - checkSparkAnswerAndFallbackReason( - s"SELECT _1 AS id, map(CAST(0 AS DOUBLE), 7)[$negZero] AS v FROM tbl", - "Spark normalizes floating-point map keys") - // The outer key type is INT, so only inspecting it would admit the lookup; the inner map is - // handed to the JVM dispatcher whole and never revisits `CometLiteral`. checkSparkAnswerAndFallbackReason( "SELECT _1 AS id, element_at(element_at(map(1, map(CAST(0 AS DOUBLE), 7)), _1), " + s"$negZero) AS v FROM tbl", @@ -310,14 +304,11 @@ class CometMapExpressionSuite extends CometTestBase { } } - // Spark compares the key under its declared collation; native compares under UTF8_BINARY. - test("map lookup with collated string keys falls back") { + // 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(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") 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", @@ -325,29 +316,4 @@ class CometMapExpressionSuite extends CometTestBase { } } - // `map_extract`'s `coerce_types` returns the map's exact Arrow key field type, so Comet casts the - // lookup key to it. A NULL inside the lookup key then aborts the cast with `Non-nullable field of - // ListArray "item" cannot contain nulls` instead of missing the lookup, which is what Spark does. - test("map lookup with complex keys falls back") { - 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") - checkSparkAnswerAndFallbackReason( - "SELECT _1 AS id, element_at(map(named_struct('a', 1), 7), " + - "named_struct('a', _1)) AS v FROM tbl", - "casts the lookup key to the map's exact Arrow key type") - } - } - - // `BinaryType` keys need no decline: Arrow compares them by content, as Spark's ordering does. - test("map lookup with binary keys stays native") { - 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") - } - } - } From 541f4e59c1d5b4953248662dcd041eb94948984c Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 25 Aug 2026 18:59:18 -0700 Subject: [PATCH 6/8] feat: support Map for `CreateArray` literal --- native/core/src/execution/planner.rs | 93 +++++++++---- .../scala/org/apache/comet/serde/arrays.scala | 55 ++++---- .../org/apache/comet/serde/literals.scala | 128 +++++++++++------- .../expressions/array/create_array.sql | 19 ++- .../sql-tests/expressions/map/map_entries.sql | 16 +++ .../comet/CometArrayExpressionSuite.scala | 13 ++ .../comet/CometMapExpressionSuite.scala | 58 ++++++++ 7 files changed, 273 insertions(+), 109 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 25667a96027..74f0d395f6b 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 @@ -3385,11 +3420,14 @@ impl PhysicalPlanner { let fun_name = &expr.func; // `map_entries` reuses the input map's entries array but declares the entry `value` field - // nullable, so widen the argument first or Arrow rejects the child type. See - // `coerce_child_fields_nullable`. + // nullable, so widen ONLY that outer field first or Arrow rejects the child type. Widening + // nested types would corrupt a nested map's `valueContainsNull` and break a downstream + // consumer. See `widen_map_entry_value_nullable`. let args = if fun_name == "map_entries" { args.into_iter() - .map(|arg| Self::coerce_child_fields_nullable(arg, &input_schema)) + .map(|arg| { + Self::coerce_child_to(arg, &input_schema, widen_map_entry_value_nullable) + }) .collect::, ExecutionError>>()? } else { args @@ -3526,36 +3564,41 @@ impl PhysicalPlanner { Ok(scalar_expr) } - /// Casts `child` to the all-nullable variant of its own type, or returns it unchanged when it - /// already is. Used where a kernel's declared output type marks every nested field nullable - /// but its implementation reuses the child's arrays, so a non-nullable inner field makes the - /// declared and produced types disagree: - /// - /// - `collect_list` / `collect_set` 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"). - /// - `map_entries` declares its list element as `Struct(key non-null, value nullable)` but - /// reuses the input map's own entries array as the list values, so `ListArray::new` - /// rejects a map whose entry `value` field is non-nullable. Every Spark `map(...)` - /// constructor over non-null values produces exactly that type. - /// - /// `make_all_fields_nullable` keeps an Arrow map's key field non-nullable, so widening a map - /// stays a legal map type. - fn coerce_child_fields_nullable( + /// 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/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 502ee85e8cf..50adf158765 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -473,16 +473,19 @@ 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 coerces an `ArrayType.containsNull` mismatch - // away before that assert, but NOT a struct field's or a `MapType.valueContainsNull` - // mismatch, so `array(struct(a not null), struct(a nullable))` and - // `array(map('x', CAST(NULL AS INT)), map('z', 3))` both panic inside `make_array_inner`. - // Decline only those cases (i.e. children that still differ after normalizing the container - // nullability DataFusion does coerce) so Spark's evaluator handles them. + // nullability, so children that share a surface type but differ only in nested nullability get + // no unifying cast, and Spark reports a merged element type (nullability-OR across children) + // as the array's own type. Reconcile that here by casting every child to the merged element + // type, so `make_array` sees identical Arrow types. Comet's cast widens container nullability + // (`ArrayType.containsNull`, `MapType.valueContainsNull`, and either nested) via the recursive + // `cast_map_to_map` / array casts, which is what lets `array(map('x', CAST(NULL AS INT)), + // map('z', 3))` and `array(map(1, array(1)), map(2, array(col)))` run natively rather than + // falling back. // - // 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. + // A struct FIELD's own nullability is the one difference kept significant (see + // [[normalizeContainerNullability]]): `array(struct(a not null), struct(a nullable))` is + // declined so Spark evaluates it, matching the coverage the `make_array` struct-field panic + // drove. val normalizedTypes = children.map(c => normalizeContainerNullability(c.dataType)) if (normalizedTypes.distinct.size > 1) { withFallbackReason( @@ -492,7 +495,13 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] { return None } - val childExprs = children.map(exprToProtoInternal(_, inputs, binding)) + // Cast each child that is not already the merged element type to it. The element data types + // are identical, so the cast only widens nullability metadata and never changes values. + val elementType = 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) + } if (childExprs.forall(_.isDefined)) { scalarFunctionExprToProto("make_array", childExprs: _*) @@ -503,30 +512,22 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] { } /** - * Rewrites a type so that `ArrayType.containsNull` is forced to `true` everywhere, while struct - * field nullability and `MapType.valueContainsNull` are left intact. Two CreateArray children - * whose types differ ONLY in `containsNull` are tolerated by DataFusion's `make_array`, because - * `type_union_resolution_coercion` routes `List` arguments through `list_coercion`, which - * merges the element field's nullability and lets Comet's planner insert the unifying cast. Its - * fallback chain has no `map_coercion` arm in DataFusion 54.1, so two `Map` arguments that - * differ in value nullability fail to unify, no cast is inserted, and `MutableArrayData` - * panics. Keep `valueContainsNull` significant here so those children are declined instead. The - * same is true of a struct field's nullability, which `coerce_struct_by_*` would merge but - * which `MutableArrayData` still rejects for the array's own element type. - * - * TODO: DataFusion 55.0.0 adds a `map_coercion` arm to `type_union_resolution_coercion` - * (apache/datafusion#23521), which coerces a `MapType.valueContainsNull` mismatch away. When - * Comet upgrades onto it, stop keeping `valueContainsNull` significant here or this decline - * stays correct but grows over-conservative, falling maps back that DataFusion can now unify. + * Rewrites a type so that container nullability Comet can cast away (`ArrayType.containsNull` + * and `MapType.valueContainsNull`, at every nesting level) is forced to `true`, while each + * struct FIELD's own nullability is left intact. Two CreateArray children whose types differ + * only in container nullability normalize equal, so [[convert]] unifies them by casting each to + * the merged element type before `make_array` (DataFusion 54.1 only coerces `List` nullability + * on its own, not `Map`; the explicit cast covers both). Children that differ in a struct + * field's own nullability normalize distinct and are declined instead, so Spark evaluates them. */ private def normalizeContainerNullability(dt: DataType): DataType = dt match { case ArrayType(elementType, _) => ArrayType(normalizeContainerNullability(elementType), containsNull = true) - case MapType(keyType, valueType, valueContainsNull) => + case MapType(keyType, valueType, _) => MapType( normalizeContainerNullability(keyType), normalizeContainerNullability(valueType), - valueContainsNull) + valueContainsNull = true) case StructType(fields) => StructType(fields.map(f => f.copy(dataType = normalizeContainerNullability(f.dataType)))) case other => other 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 19bd20c8086..889c6ff5e31 100644 --- a/spark/src/main/scala/org/apache/comet/serde/literals.scala +++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala @@ -24,14 +24,14 @@ import java.lang import org.apache.spark.internal.Logging 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, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, TimestampNTZType, TimestampType} +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.{CometExpressionSerde, Compatible, ExprOuterClass, LiteralOuterClass, MapKeySupport, SupportLevel, Unsupported} import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, serializeDataType, supportedDataType} import org.apache.comet.serde.Types.ListLiteral @@ -55,7 +55,7 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { .elementType .isInstanceOf[ArrayType])))) { Compatible(None) - } else if (canExpandComplexLiteral(expr)) { + } else if (expandComplexLiteral(expr).isDefined) { // Rebuilt as a `CreateArray` / `CreateMap` tree in `convert`. Compatible(None) } else { @@ -228,25 +228,37 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { /** * Rebuild a folded complex Literal as an equivalent tree of `CreateArray` / `CreateMap` over - * primitive-typed Literals, or `None` when the shape cannot be rebuilt. 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: + * primitive-typed Literals, or `None` when the shape cannot be rebuilt. `getSupportLevel` + * probes with this same function and reports `Compatible` when it returns a `Some`, so + * admission and expansion cannot diverge. 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. Shapes the native map kernels cannot handle are therefore - * declined by their consumers rather than here: a map lookup with an unsupported key type by - * [[MapKeySupport]], and a non-nullable-value map into `map_entries` by the native planner - * widening its argument. + * 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, however, are declined here at admission (see [[mapKeyTypesExpandable]]) + * rather than left to consumers. A map lookup gates itself on [[MapKeySupport]], but not every + * native map consumer does: `map_contains_key` lowers to `array_contains(map_keys(...), key)`, + * and neither kernel checks the key type. A map that is only the value of another map is also + * opaque once rebuilt, because `CometCreateMap` hands the whole `CreateMap` to the JVM codegen + * dispatcher, so a nested unsupported key type never revisits this serde. Declining any folded + * literal that (recursively) contains an unsupported or non-orderable map key type keeps all of + * those paths on Spark. * * 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 @@ -257,43 +269,61 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { * - Folded maps with duplicate keys, see [[hasDuplicateMapKeys]]. */ private def expandComplexLiteral(expr: Literal): Option[Expression] = { - if (!canExpandComplexLiteral(expr)) return None + if (expr.value == null || !mapKeyTypesExpandable(expr.dataType)) return None expr.dataType match { - case ArrayType(et, containsNull) => + case ArrayType(et, containsNull) if needsExpansion(et) => 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)) + if (arr.numElements() == 0) { + None + } else { + 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)) + if (mapData.numElements() == 0 || hasDuplicateMapKeys(keys, kt)) { + None + } else { + 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 } } /** - * True when [[expandComplexLiteral]] can rebuild this folded literal, tested without allocating - * the rebuilt tree so `getSupportLevel` can probe cheaply (`convert` runs the build only once, - * on the value it keeps). Mirrors the shapes [[expandComplexLiteral]] declines: a non-null - * value, a non-empty top-level container, an array whose elements need expansion, and a map - * with no duplicate keys. + * 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 canExpandComplexLiteral(expr: Literal): Boolean = { - if (expr.value == null) 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 - } + 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 } /** @@ -307,13 +337,9 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { case _ => false } - /** Element `i` of `arr`, or `null` for a null slot. */ - private def valueAt(arr: ArrayData, i: Int, dt: DataType): Any = - if (arr.isNullAt(i)) null else arr.get(i, dt) - - /** Element `i` of `arr` as a Literal of type `dt`. */ + /** 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(valueAt(arr, i, dt), dt) + 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 @@ -322,14 +348,14 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { * `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. * - * `ArrayBasedMapBuilder` hashes most atomic keys but compares `BinaryType`, collated - * `StringType` and complex keys through `TypeUtils.getInterpretedOrdering`, and it normalizes - * `-0.0` and `NaN` before either. The interpreted ordering is the one comparison that agrees - * with all of those: `Array[Byte]` keys compare by content rather than by identity, a collated - * `StringType` compares under its collation, and `nanSafeCompareDoubles` already treats `-0.0` - * and `+0.0`, and any two `NaN`s, as equal. 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. + * 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() 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 5dfda61fabd..320d4985a25 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 @@ -68,17 +68,24 @@ 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. DataFusion 54.1 cannot unify two Map arguments either: the fallback chain in --- `type_union_resolution_coercion` has no `map_coercion` arm, so `make_array` receives the two --- differently typed map arrays and `MutableArrayData` panics. `CometCreateArray` therefore keeps --- `MapType.valueContainsNull` significant when it normalizes container nullability and declines. -query expect_fallback(CreateArray children have mismatched data types) +-- 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 stays native. +-- 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)) 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 46e4f7c111a..a6c145e9f3c 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 @@ -36,3 +36,19 @@ 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 diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index cd3e40b8f6d..c72d94866c8 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1183,6 +1183,19 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } + // 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") diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index d91dd671b82..c26c352915f 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -316,4 +316,62 @@ class CometMapExpressionSuite extends CometTestBase { } } + // 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") + } + } + } From ba011bbf1538f3cf1e6fc22d3ea189e81ae0a13f Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 26 Aug 2026 10:01:02 -0700 Subject: [PATCH 7/8] feat: support Map for `CreateArray` literal --- .../expressions/map/element_at_map.sql | 5 ++ .../comet/CometMapExpressionSuite.scala | 54 +++++++++++++++++++ 2 files changed, 59 insertions(+) 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 5bb322b7f32..ded57f490e2 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 @@ -64,6 +64,11 @@ 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) diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index c26c352915f..74e88231c54 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -374,4 +374,58 @@ class CometMapExpressionSuite extends CometTestBase { } } + // 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") + } + } + } + } From b64837f9057813d0791ab53b896388f9f4b433d8 Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 26 Aug 2026 16:29:10 -0700 Subject: [PATCH 8/8] feat: support Map for `CreateArray` literal --- native/core/src/execution/planner.rs | 6 +- .../comet/codegen/CometBatchKernel.java | 8 +- .../scala/org/apache/comet/serde/arrays.scala | 63 +++++---------- .../org/apache/comet/serde/literals.scala | 77 ++++++++++--------- .../expressions/array/create_array.sql | 15 ++++ .../expressions/map/element_at_map.sql | 16 ++++ .../sql-tests/expressions/map/map_entries.sql | 11 +++ .../comet/CometArrayExpressionSuite.scala | 48 ++---------- .../comet/CometMapExpressionSuite.scala | 22 ++++++ 9 files changed, 142 insertions(+), 124 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 74f0d395f6b..3536073daa0 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -3419,10 +3419,8 @@ impl PhysicalPlanner { .collect::, _>>()?; let fun_name = &expr.func; - // `map_entries` reuses the input map's entries array but declares the entry `value` field - // nullable, so widen ONLY that outer field first or Arrow rejects the child type. Widening - // nested types would corrupt a nested map's `valueContainsNull` and break a downstream - // consumer. See `widen_map_entry_value_nullable`. + // `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| { 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 50adf158765..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,32 +472,15 @@ 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 nullability get - // no unifying cast, and Spark reports a merged element type (nullability-OR across children) - // as the array's own type. Reconcile that here by casting every child to the merged element - // type, so `make_array` sees identical Arrow types. Comet's cast widens container nullability - // (`ArrayType.containsNull`, `MapType.valueContainsNull`, and either nested) via the recursive - // `cast_map_to_map` / array casts, which is what lets `array(map('x', CAST(NULL AS INT)), - // map('z', 3))` and `array(map(1, array(1)), map(2, array(col)))` run natively rather than - // falling back. - // - // A struct FIELD's own nullability is the one difference kept significant (see - // [[normalizeContainerNullability]]): `array(struct(a not null), struct(a nullable))` is - // declined so Spark evaluates it, matching the coverage the `make_array` struct-field panic - // drove. - 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 - } - - // Cast each child that is not already the merged element type to it. The element data types - // are identical, so the cast only widens nullability metadata and never changes values. - val elementType = expr.dataType.asInstanceOf[ArrayType].elementType + // 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) @@ -512,24 +495,20 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] { } /** - * Rewrites a type so that container nullability Comet can cast away (`ArrayType.containsNull` - * and `MapType.valueContainsNull`, at every nesting level) is forced to `true`, while each - * struct FIELD's own nullability is left intact. Two CreateArray children whose types differ - * only in container nullability normalize equal, so [[convert]] unifies them by casting each to - * the merged element type before `make_array` (DataFusion 54.1 only coerces `List` nullability - * on its own, not `Map`; the explicit cast covers both). Children that differ in a struct - * field's own nullability normalize distinct and are declined instead, so Spark evaluates them. + * 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 } } 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 889c6ff5e31..d50a0af5d60 100644 --- a/spark/src/main/scala/org/apache/comet/serde/literals.scala +++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala @@ -55,7 +55,7 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { .elementType .isInstanceOf[ArrayType])))) { Compatible(None) - } else if (expandComplexLiteral(expr).isDefined) { + } else if (canExpandComplexLiteral(expr)) { // Rebuilt as a `CreateArray` / `CreateMap` tree in `convert`. Compatible(None) } else { @@ -228,13 +228,13 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { /** * Rebuild a folded complex Literal as an equivalent tree of `CreateArray` / `CreateMap` over - * primitive-typed Literals, or `None` when the shape cannot be rebuilt. `getSupportLevel` - * probes with this same function and reports `Compatible` when it returns a `Some`, so - * admission and expansion cannot diverge. 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 + * 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 @@ -243,14 +243,8 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { * 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, however, are declined here at admission (see [[mapKeyTypesExpandable]]) - * rather than left to consumers. A map lookup gates itself on [[MapKeySupport]], but not every - * native map consumer does: `map_contains_key` lowers to `array_contains(map_keys(...), key)`, - * and neither kernel checks the key type. A map that is only the value of another map is also - * opaque once rebuilt, because `CometCreateMap` hands the whole `CreateMap` to the JVM codegen - * dispatcher, so a nested unsupported key type never revisits this serde. Declining any folded - * literal that (recursively) contains an unsupported or non-orderable map key type keeps all of - * those paths on Spark. + * 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 @@ -269,34 +263,47 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging { * - Folded maps with duplicate keys, see [[hasDuplicateMapKeys]]. */ private def expandComplexLiteral(expr: Literal): Option[Expression] = { - if (expr.value == null || !mapKeyTypesExpandable(expr.dataType)) return None + if (!canExpandComplexLiteral(expr)) return None expr.dataType match { - case ArrayType(et, containsNull) if needsExpansion(et) => + case ArrayType(et, containsNull) => val arr = expr.value.asInstanceOf[ArrayData] - if (arr.numElements() == 0) { - None - } else { - val elements = (0 until arr.numElements()) - .map(i => withNullability(literalAt(arr, i, et), containsNull)) - Some(CreateArray(elements, useStringTypeWhenEmpty = false)) - } + 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() - if (mapData.numElements() == 0 || hasDuplicateMapKeys(keys, kt)) { - None - } else { - 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)) - } + 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 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 320d4985a25..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 @@ -98,6 +98,21 @@ SELECT array(array(map(1, 'a')), array(map(2, 'b'))) 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 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 ded57f490e2..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 @@ -80,3 +80,19 @@ 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/map_entries.sql b/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql index a6c145e9f3c..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 @@ -52,3 +52,14 @@ 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 c72d94866c8..40505562300 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1134,48 +1134,12 @@ 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") - } - } - - 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") - } - } - - // 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. - // `CometLiteral` rebuilds each folded map as an equivalent `CreateMap` of primitive literals. + // 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( diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 74e88231c54..7311ef38d14 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -272,6 +272,15 @@ 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 @@ -428,4 +437,17 @@ class CometMapExpressionSuite extends CometTestBase { } } + // 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") + } + } + }