comphead commented on code in PR #6178:
URL: https://github.com/apache/datafusion-comet/pull/6178#discussion_r4354855255
##########
spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala:
##########
@@ -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)))
Review Comment:
Four assertions in this test (lines 1222, 1223, 1227 and 1249) repeat cases
that `codegen_dispatch_null_map_key.sql` already runs end to end: `long` to
`int`, `date` to `timestamp`, the `long` to `decimal(20, 0)` widening and
`MapKeys` over a narrowing cast. Could they go, keeping the unit test for what
the fixture cannot reach (string to int, int to string, `decimal(19, 0)`, the
LEGACY and ANSI modes and the nested positions)?
##########
spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala:
##########
@@ -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)) ||
Review Comment:
The `ArrayData` arm only recurses when the element type contains a
`MapType`, so for the key array this call needs a map key type that contains a
map. Spark's map constructors reject that through
`TypeUtils.checkForMapKeyType`, and none of the literal assertions has a nested
map in key position. I could not find a way to build such a literal, but I only
looked at the constructors. Could this line be dropped, leaving the `isNullAt`
scan on the keys and the value recursion?
##########
spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala:
##########
@@ -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") {
Review Comment:
The comment says the SQL file tests run with constant folding off, but a
fixture can opt in with `-- ConstantFolding: enabled` (documented in
`sql-file-tests.md`, used by
`expressions/array/folded_array_map_literals.sql`). Could this query move to a
second fixture next to `codegen_dispatch_null_map_key.sql` using that directive
and `query expect_fallback(null map key)`? The directive is file-wide, so it
needs its own file. That would keep the #6172 coverage in one place. I have not
run such a fixture.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
