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 12:11:06 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6178-c56f6a0209dba0600103a37fac7ca9b006639857 created (now 6a6d71c52f)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6178-c56f6a0209dba0600103a37fac7ca9b006639857
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 6a6d71c52f fix: refuse codegen dispatch for a TRY cast that can put a null key in a map (#6178)
This branch includes the following new commits:
new 6a6d71c52f fix: refuse codegen dispatch for a TRY cast that can put a null key in a map (#6178)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: fix: refuse codegen dispatch for a TRY cast that can put a null key in a map (#6178)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6178-c56f6a0209dba0600103a37fac7ca9b006639857
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
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