You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/23 13:32:00 UTC
(datafusion-comet) branch main updated: fix: refuse codegen dispatch for a TRY cast that can put a null key in a map (#6178)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 6a6d71c52f fix: refuse codegen dispatch for a TRY cast that can put a null key in a map (#6178)
6a6d71c52f is described below
commit 6a6d71c52f3452ac92ebe1c30c92c8515508f59e
Author: Han-Yin Chang <ni...@gmail.com>
AuthorDate: Fri Oct 23 10:18:04 2026 +0000
fix: refuse codegen dispatch for a TRY cast that can put a null key in a map (#6178)
* fix: refuse codegen dispatch for a TRY cast that can put a null key in a map
The dispatcher's map writer writes each key with no null check, so a null key comes out as the key type's default value. A TRY cast whose key cast fails produces one: the failing key becomes null, and Spark's Cast.castMap builds the map without rejecting it. map_keys(transform_values(try_cast(m AS map<int,int>), ...)) returned [1, 0] where Spark returns [1, NULL].
Arrow's map format cannot hold a null key, and Spark's own readers disagree about one, so refuse any tree holding such a TRY cast in canHandle and let Spark evaluate it. A key cast is treated as able to fail unless it is on a short list of casts that cannot throw. That list is narrower than Spark's Cast.canUpCast, which admits date and timestamp conversions that overflow for extreme values.
Closes #6172.
* review: refuse only trees whose output can carry a map, and catch a folded null-key literal
---------
Co-authored-by: Andy Grove <ag...@apache.org>
---
.../comet/codegen/CometBatchKernelCodegen.scala | 43 ++++++++++-
.../org/apache/comet/expressions/CometCast.scala | 62 +++++++++++++++-
.../misc/codegen_dispatch_null_map_key.sql | 65 +++++++++++++++++
.../scala/org/apache/comet/CometCodegenSuite.scala | 83 +++++++++++++++++++++-
4 files changed, 247 insertions(+), 6 deletions(-)
diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala
index f906334101..fb72412d72 100644
--- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala
+++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala
@@ -25,15 +25,19 @@ import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
import org.apache.arrow.vector.types.pojo.Field
import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.DeserializerBuildHelper.createDeserializerForTypesSupportValueOf
+import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.SerializerBuildHelper.{createSerializerForBoolean, createSerializerForByte, createSerializerForDouble, createSerializerForFloat, createSerializerForInteger, createSerializerForLong, createSerializerForShort}
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
-import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, Literal, ScalaUDF, Unevaluable}
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, Cast, Expression, Literal, ScalaUDF, Unevaluable}
import org.apache.spark.sql.catalyst.expressions.codegen._
+import org.apache.spark.sql.catalyst.util.{ArrayData, MapData}
import org.apache.spark.sql.execution.ExecSubqueryExpression
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
import org.apache.comet.CometArrowAllocator
+import org.apache.comet.expressions.CometCast
+import org.apache.comet.serde.SupportLevel
import org.apache.comet.shims.{CometExprTraitShim, CometTypeShim}
/**
@@ -139,19 +143,52 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
case _ => 1
}
+ /** Whether a literal value holds a map with a null key at any depth. */
+ private def hasNullMapKey(value: Any, dataType: DataType): Boolean = (value, dataType) match {
+ case (null, _) => false
+ case (map: MapData, MapType(keyType, valueType, _)) =>
+ val keys = map.keyArray()
+ (0 until keys.numElements()).exists(keys.isNullAt) ||
+ hasNullMapKey(keys, ArrayType(keyType)) ||
+ hasNullMapKey(map.valueArray(), ArrayType(valueType))
+ case (array: ArrayData, ArrayType(elementType, _))
+ if SupportLevel.containsType(elementType, classOf[MapType]) =>
+ (0 until array.numElements()).exists(i =>
+ !array.isNullAt(i) && hasNullMapKey(array.get(i, elementType), elementType))
+ case (row: InternalRow, StructType(fields)) =>
+ fields.indices.exists(i =>
+ !row.isNullAt(i) && hasNullMapKey(row.get(i, fields(i).dataType), fields(i).dataType))
+ case _ => false
+ }
+
/**
* Plan-time predicate. `None` greenlights the serde to emit the codegen proto; `Some(reason)`
* forces a Spark fallback (typically `withFallbackReason(...) + None`) so the operator falls
* back cleanly rather than crashing the Janino compile at execute time.
*
* Checks every `BoundReference`'s data type and the root `expr.dataType` against
- * [[isSupportedDataType]], rejects aggregates / generators / `Unevaluable`, and gates total
- * nested-field count on `spark.sql.codegen.maxFields`.
+ * [[isSupportedDataType]], refuses a tree whose output can carry a map with a null key, rejects
+ * aggregates / generators / `Unevaluable`, and gates total nested-field count on
+ * `spark.sql.codegen.maxFields`.
*/
def canHandle(boundExpr: Expression): Option[String] = {
if (!isSupportedDataType(boundExpr.dataType)) {
return Some(s"codegen dispatch: unsupported output type ${boundExpr.dataType}")
}
+ // Arrow cannot hold a null map key, so the output writer would lose one. Only refuse when the
+ // output can carry a map, since inside the kernel it is Spark's own generated code.
+ // https://github.com/apache/datafusion-comet/issues/6172
+ if (SupportLevel.containsType(boundExpr.dataType, classOf[MapType]) &&
+ boundExpr.find {
+ case c: Cast => CometCast.canProduceNullMapKey(c)
+ // A TRY cast over a foldable map is already folded into a literal holding the null key.
+ case l: Literal => hasNullMapKey(l.value, l.dataType)
+ case _ => false
+ }.isDefined) {
+ return Some(
+ "codegen dispatch: the output can carry a null map key, which Arrow cannot hold " +
+ "(https://github.com/apache/datafusion-comet/issues/6172)")
+ }
// Mirror WSCG's `spark.sql.codegen.maxFields` gate. Wide schemas blow the generated class's
// typed input field count, the typed-getter switch, and the constant pool. Refuse here so the
// operator falls back to Spark cleanly rather than tripping a Janino compile failure
diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala
index 0abd6dd52d..2abfff01b3 100644
--- a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala
+++ b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala
@@ -23,7 +23,7 @@ import scala.util.Try
import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Expression, Literal}
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, DecimalType, MapType, NullType, StructType, TimestampNTZType, TimestampType}
+import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, DecimalType, MapType, NullType, NumericType, StructType, TimestampNTZType, TimestampType}
import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, withFallbackReason}
import org.apache.comet.DataTypeSupport.isComplexType
@@ -563,6 +563,66 @@ object CometCast
case _ => Unsupported(Some(s"Cast from DateType to $toType is not supported"))
}
+ /**
+ * Whether evaluating `cast` can produce a map with a null key. Under TRY a key cast that throws
+ * becomes a null key, and Spark's `Cast.castMap` builds the result without rejecting it, so any
+ * TRY cast that changes a map's key type at any depth can, unless the key cast cannot fail.
+ * Arrow's map format requires a non-nullable key field, so no Comet path can hold that value.
+ */
+ private[comet] def canProduceNullMapKey(cast: Cast): Boolean =
+ evalMode(cast) == CometEvalMode.TRY && mayNullMapKey(cast.child.dataType, cast.dataType)
+
+ private def mayNullMapKey(fromType: DataType, toType: DataType): Boolean =
+ (fromType, toType) match {
+ case (MapType(fromKey, fromValue, _), MapType(toKey, toValue, _)) =>
+ !keyCastCannotFail(fromKey, toKey) || mayNullMapKey(fromValue, toValue)
+ case (ArrayType(fromElement, _), ArrayType(toElement, _)) =>
+ mayNullMapKey(fromElement, toElement)
+ case (StructType(fromFields), StructType(toFields)) =>
+ fromFields.zip(toFields).exists { case (from, to) =>
+ mayNullMapKey(from.dataType, to.dataType)
+ }
+ case _ => false
+ }
+
+ private val numericWidening =
+ Seq(
+ DataTypes.ByteType,
+ DataTypes.ShortType,
+ DataTypes.IntegerType,
+ DataTypes.LongType,
+ DataTypes.FloatType,
+ DataTypes.DoubleType)
+
+ private def integralDigits(dataType: DataType): Option[Int] = dataType match {
+ case DataTypes.ByteType => Some(3)
+ case DataTypes.ShortType => Some(5)
+ case DataTypes.IntegerType => Some(10)
+ case DataTypes.LongType => Some(20)
+ case _ => None
+ }
+
+ /**
+ * Key casts that cannot throw, so cannot become a null key under TRY. Deliberately narrower
+ * than Spark's `Cast.canUpCast`, which also admits date and timestamp conversions that overflow
+ * `Math.multiplyExact` for extreme values. Anything not listed is treated as able to fail.
+ */
+ private def keyCastCannotFail(fromType: DataType, toType: DataType): Boolean =
+ (fromType, toType) match {
+ case _ if fromType == toType => true
+ case (from, to) if numericWidening.contains(from) && numericWidening.contains(to) =>
+ numericWidening.indexOf(from) < numericWidening.indexOf(to)
+ case (from, to: DecimalType) if integralDigits(from).isDefined =>
+ to.precision - to.scale >= integralDigits(from).get
+ case (from: DecimalType, to: DecimalType) =>
+ to.scale >= from.scale && to.precision - to.scale >= from.precision - from.scale
+ case (
+ _: NumericType | DataTypes.BooleanType | DataTypes.DateType | DataTypes.TimestampType,
+ DataTypes.StringType) =>
+ true
+ case _ => false
+ }
+
private def unsupported(fromType: DataType, toType: DataType): Unsupported = {
Unsupported(Some(s"Cast from $fromType to $toType is not supported"))
}
diff --git a/spark/src/test/resources/sql-tests/expressions/misc/codegen_dispatch_null_map_key.sql b/spark/src/test/resources/sql-tests/expressions/misc/codegen_dispatch_null_map_key.sql
new file mode 100644
index 0000000000..b6242f5c5a
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/misc/codegen_dispatch_null_map_key.sql
@@ -0,0 +1,65 @@
+-- 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.
+-- A TRY cast whose key cast can fail turns an out-of-range key into a null key, and Spark keeps
+-- the row. Arrow cannot hold a null key, so a dispatched tree whose output can carry a map must
+-- fall back to Spark. See #6172.
+
+statement
+CREATE TABLE test_null_map_key(id int, m map<bigint, int>) USING parquet
+
+statement
+INSERT INTO test_null_map_key VALUES
+ (1, map(1, 10, 9999999999, 20)),
+ (2, map(2, 30)),
+ (3, map()),
+ (4, NULL)
+
+-- Before the fix these returned [1, 0] for the first row, where Spark returns [1, NULL]
+query expect_fallback(null map key)
+SELECT id, map_keys(transform_values(try_cast(m AS map<int, int>), (k, v) -> v + 1))
+FROM test_null_map_key ORDER BY id
+
+query expect_fallback(null map key)
+SELECT id, map_keys(map_filter(try_cast(m AS map<int, int>), (k, v) -> true))
+FROM test_null_map_key ORDER BY id
+
+-- The output carries no map, so the kernel's answer is Spark's own generated code
+query expect_dispatch(cast)
+SELECT id, cast(try_cast(m AS map<int, int>) AS string) FROM test_null_map_key ORDER BY id
+
+query expect_dispatch(exists)
+SELECT id, exists(map_keys(try_cast(m AS map<int, int>)), k -> k IS NULL)
+FROM test_null_map_key ORDER BY id
+
+-- A widening key cast cannot fail, so the tree stays in the dispatcher
+query expect_dispatch(transform_values)
+SELECT id, map_keys(transform_values(try_cast(m AS map<decimal(20, 0), int>), (k, v) -> v))
+FROM test_null_map_key ORDER BY id
+
+-- Spark counts date to timestamp as an upcast, but it overflows for extreme dates and so can
+-- still produce a null key under TRY
+statement
+CREATE TABLE test_null_map_key_date(id int, m map<date, int>) USING parquet
+
+statement
+INSERT INTO test_null_map_key_date VALUES
+ (1, map(date_from_unix_date(1), 10, date_from_unix_date(2000000000), 20)),
+ (2, NULL)
+
+query expect_fallback(null map key)
+SELECT id, map_keys(transform_values(try_cast(m AS map<timestamp, int>), (k, v) -> v))
+FROM test_null_map_key_date ORDER BY id
diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
index 6bdad65ebf..63343c23c7 100644
--- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
@@ -28,10 +28,10 @@ import org.apache.spark.{SparkConf, SparkEnv, TaskContext}
import org.apache.spark.sql.{CometTestBase, Row}
import org.apache.spark.sql.api.java.UDF1
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, AttributeSeq, BindReferences, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, If, IsNull, LessThan, Literal, MapConcat, MonotonicallyIncreasingID, Or, Rand, ScalaUDF}
+import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, AttributeSeq, BindReferences, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, EvalMode, Expression, GenericInternalRow, Hypot, If, IsNull, LessThan, Literal, MapConcat, MapKeys, MonotonicallyIncreasingID, Or, Rand, ScalaUDF}
import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial}
import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke}
-import org.apache.spark.sql.catalyst.util.GenericArrayData
+import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData, MapData}
import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec, CometProjectExec}
import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryCatalog}
import org.apache.spark.sql.connector.catalog.functions.{BoundFunction, ScalarFunction, UnboundFunction}
@@ -1209,6 +1209,85 @@ class CometCodegenSuite
}
}
+ test("canHandle refuses a TRY cast whose map key cast can fail") {
+ // https://github.com/apache/datafusion-comet/issues/6172
+ def refused(from: DataType, to: DataType, mode: EvalMode.Value = EvalMode.TRY): Boolean =
+ CometBatchKernelCodegen
+ .canHandle(Cast(BoundReference(0, from, nullable = true), to, None, mode))
+ .exists(_.contains("null map key"))
+ def map(k: DataType): MapType = MapType(k, IntegerType)
+
+ // Narrowing can overflow, and date to timestamp can overflow for extreme dates even though
+ // Spark counts it as an upcast.
+ assert(refused(map(LongType), map(IntegerType)))
+ assert(refused(map(DateType), map(TimestampType)))
+ assert(refused(map(StringType), map(IntegerType)))
+ // Widening numeric, integral to a wide enough decimal, and numeric to string cannot fail.
+ assert(!refused(map(IntegerType), map(LongType)))
+ assert(!refused(map(LongType), map(DecimalType(20, 0))))
+ assert(!refused(map(IntegerType), map(StringType)))
+ // The list takes a long as 20 digits, as Spark's own `DecimalType.forType` does, so
+ // decimal(19, 0) is refused even though every long fits. That is only conservative.
+ assert(refused(map(LongType), map(DecimalType(19, 0))))
+ // Only TRY turns a failing key cast into a null key.
+ assert(!refused(map(LongType), map(IntegerType), EvalMode.LEGACY))
+ assert(!refused(map(LongType), map(IntegerType), EvalMode.ANSI))
+ // Reached at any depth: inside an array, a struct field and a map value.
+ assert(refused(ArrayType(map(LongType)), ArrayType(map(IntegerType))))
+ assert(
+ refused(
+ StructType(Seq(StructField("a", IntegerType), StructField("m", map(LongType)))),
+ StructType(Seq(StructField("a", IntegerType), StructField("m", map(IntegerType))))))
+ assert(refused(MapType(IntegerType, map(LongType)), MapType(IntegerType, map(IntegerType))))
+ // Only refused when the output can carry a map. Inside the kernel it is Spark's own code.
+ val narrowing =
+ Cast(
+ BoundReference(0, map(LongType), nullable = true),
+ map(IntegerType),
+ None,
+ EvalMode.TRY)
+ assert(CometBatchKernelCodegen.canHandle(MapKeys(narrowing)).isEmpty)
+
+ // A TRY cast over a foldable map is folded into a literal that already holds the null key.
+ def nullKeyMap: MapData =
+ new ArrayBasedMapData(
+ new GenericArrayData(Array[Any](1, null)),
+ new GenericArrayData(Array[Any](10, 20)))
+ def literalRefused(lit: Literal): Boolean =
+ CometBatchKernelCodegen.canHandle(lit).exists(_.contains("null map key"))
+ assert(literalRefused(Literal(nullKeyMap, map(IntegerType))))
+ assert(
+ literalRefused(
+ Literal(new GenericArrayData(Array[Any](nullKeyMap)), ArrayType(map(IntegerType)))))
+ assert(
+ literalRefused(
+ Literal(
+ InternalRow(1, nullKeyMap),
+ StructType(Seq(StructField("a", IntegerType), StructField("m", map(IntegerType)))))))
+ assert(
+ !literalRefused(
+ Literal(
+ new ArrayBasedMapData(
+ new GenericArrayData(Array[Any](1, 2)),
+ new GenericArrayData(Array[Any](10, 20))),
+ map(IntegerType))))
+ }
+
+ test("a TRY cast over a foldable map with a failing key falls back") {
+ // https://github.com/apache/datafusion-comet/issues/6172
+ // `ConstantFolding` turns the cast into a literal that already holds the null key, so no
+ // `Cast` is left for the guard. The SQL file tests run with constant folding off.
+ withTable("t") {
+ sql("CREATE TABLE t(id int) USING parquet")
+ sql("INSERT INTO t VALUES (1), (2)")
+ checkSparkAnswerAndFallbackReason(
+ """SELECT id, map_keys(transform_values(
+ | try_cast(map(1L, 10, 9999999999L, 20) AS map<int, int>), (k, v) -> v + id))
+ |FROM t ORDER BY id""".stripMargin,
+ "the output can carry a null map key")
+ }
+ }
+
test("a primitive ScalaUDF's null guard runs in the UDF's kernel (#6704)") {
// Spark's `HandleNullInputsForUDF` wraps each call below as
// `if (isnull(a) or ...) null else f(knownnotnull(a), ...)`. Run natively, the `if` is a
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org