This is an automated email from the ASF dual-hosted git repository.
andygrove 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 f13cc977ca fix: codegen dispatcher null short-circuit swallowed ANSI
errors, plus two latent TIME-type gaps (#5219)
f13cc977ca is described below
commit f13cc977cac30df422f99fa9cf04132ff3229390
Author: Andy Grove <[email protected]>
AuthorDate: Mon Aug 3 13:58:24 2026 -0600
fix: codegen dispatcher null short-circuit swallowed ANSI errors, plus two
latent TIME-type gaps (#5219)
* fix: codegen dispatcher null short-circuit swallowed errors Spark raises
The kernel applied an "any input null -> null output" short-circuit on the
union of input ordinals whenever every node in the bound tree was
NullIntolerant. Spark's null handling is per-node and left-to-right
(BinaryExpression.nullSafeCodeGen emits the left child's code
unconditionally, then tests the left child's null, then the right child's),
so a union short-circuit skips a subtree Spark would have evaluated and
loses any error it raises.
Under ANSI, `add_months(cast(s as date), i)` on ('notadate', NULL) raises
CAST_INVALID_INPUT in Spark; Comet returned NULL. ANSI is on by default in
Spark 4, and ~70 built-in expressions route through CometCodegenDispatch.
Restrict the short-circuit to trees that read exactly one input ordinal.
With one input, Spark also evaluates nothing when that input is null, so
the optimization stays exact and the common unary shapes (upper, length,
date_format) keep the fast path.
Also close two TIME-type gaps in the same dispatcher, where canHandle
accepted a type the runtime rejected:
- CometScalaUDFCodegen.specFor omitted TimeNanoVector and threw
UnsupportedOperationException at execute time, after the plan had
committed to the kernel, with no fallback available.
- CometSpecializedGettersDispatch.get had no time-type branch, so
get(ordinal, TimeType) threw. That is the path SafeProjection and every
CodegenFallback.eval(row) use.
Closes #5218
* ci: register CometSpecializedGettersDispatchSuite in PR workflows
* fix: keep the null short-circuit for leaf-only multi-input trees
Addresses review feedback on #5219. The previous fix dropped the
union-of-ordinals short-circuit for every multi-argument NullIntolerant
tree, not only the ones that could swallow an error.
The bug needs a subtree between the root and a leaf that can raise
independently of nullness, like the Cast in add_months(CAST(s AS DATE), i).
A BoundReference or Literal read cannot raise on its own, so when every
direct child of the root is a leaf the tree is one level deep and the union
short-circuit is still exact: any error comes from the root's own logic,
which runs only once every input is known non-null, in Spark's shape and in
ours alike.
canShortCircuitNulls now accepts either a single input ordinal or a
leaf-only-children root, restoring the fast path for the common no-cast
multi-argument case (pmod, conv, make_timestamp) across the ~70 expressions
that route through this dispatcher. defaultBody emits the disjunction over
every ordinal the tree reads rather than only the first.
Verified against Spark's generated code that the leaf-only rule holds even
where Spark reorders children: DivModLike.doGenCode and Pmod.doGenCode
evaluate the divisor first but still test the dividend's null before their
divide/remainder-by-zero throw, so pmod(NULL, 0) returns NULL rather than
raising.
* test: make the pmod-by-zero error class version-aware
Spark 4.1 introduced REMAINDER_BY_ZERO; 3.4, 3.5 and 4.0 raise
DIVIDE_BY_ZERO for pmod. The new #5218 test hard-coded the 4.1 class,
so the [expressions] job failed on those three profiles.
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
.../comet/codegen/CometBatchKernelCodegen.scala | 191 +++++++++++++--------
.../codegen/CometSpecializedGettersDispatch.scala | 12 +-
.../comet/udf/codegen/CometScalaUDFCodegen.scala | 10 +-
.../org/apache/comet/CometCodegenSourceSuite.scala | 131 +++++++++++++-
.../scala/org/apache/comet/CometCodegenSuite.scala | 142 ++++++++++++++-
.../CometSpecializedGettersDispatchSuite.scala | 74 ++++++++
8 files changed, 483 insertions(+), 79 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml
b/.github/workflows/pr_build_linux.yml
index 2d01ecc104..004f10b037 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -400,6 +400,7 @@ jobs:
org.apache.comet.CometCodegenHOFSuite
org.apache.comet.CometFuzzMathSuite
org.apache.comet.CometCodegenFuzzSuite
+ org.apache.comet.codegen.CometSpecializedGettersDispatchSuite
org.apache.comet.CometStringDecodeSuite
org.apache.comet.CometWidthBucketSuite
org.apache.comet.CometUuidExpressionSuite
diff --git a/.github/workflows/pr_build_macos.yml
b/.github/workflows/pr_build_macos.yml
index 5f60c07451..d290714411 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -216,6 +216,7 @@ jobs:
org.apache.comet.CometCodegenHOFSuite
org.apache.comet.CometFuzzMathSuite
org.apache.comet.CometCodegenFuzzSuite
+ org.apache.comet.codegen.CometSpecializedGettersDispatchSuite
org.apache.comet.CometStringDecodeSuite
org.apache.comet.CometWidthBucketSuite
org.apache.comet.CometUuidExpressionSuite
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 cce8d79476..c29300f9ed 100644
---
a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala
+++
b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala
@@ -337,10 +337,10 @@ object CometBatchKernelCodegen extends Logging with
CometExprTraitShim with Come
}
/**
- * Per-row body. For `NullIntolerant` expressions where the entire tree
propagates nulls,
- * prepends a short-circuit on the union of input ordinals so the whole
`ev.code` cost is
- * skipped on null rows. Otherwise the standard shape: run `ev.code`, then
`setNull` or write
- * based on `ev.isNull`.
+ * Per-row body. For `NullIntolerant` expressions whose input nulls fully
determine a null
+ * output (see [[canShortCircuitNulls]]), prepends a short-circuit on those
ordinals so the
+ * whole `ev.code` cost is skipped on null rows. Otherwise the standard
shape: run `ev.code`,
+ * then `setNull` or write based on `ev.isNull`.
*
* `subExprsCode` is the CSE helper-invocation block. It must run before
`ev.code`. Inside the
* short-circuit it lives in the else branch so null rows skip CSE too.
@@ -351,82 +351,131 @@ object CometBatchKernelCodegen extends Logging with
CometExprTraitShim with Come
ev: ExprCode,
writeSnippet: String,
subExprsCode: String): String = {
- boundExpr match {
- case _ if isNullIntolerant(boundExpr) && allNullIntolerant(boundExpr) =>
- // Every node from root to leaf is `NullIntolerant` or a leaf, so "any
BoundReference null
- // -> whole expression null". A non-null-propagating node like
`coalesce` or `if` would
- // make this incorrect (`coalesce(null, x)` is `x`);
`allNullIntolerant` rejects those.
- val inputOrdinals =
- boundExpr.collect { case b: BoundReference => b.ordinal }.distinct
- // Primitive Arrow vectors are wrapped in `CometPlainVector` at
input-cast time, which
- // exposes `isNullAt(int)` rather than the raw Arrow `isNull(int)`.
Pick the right method
- // per ordinal so the short-circuit compiles for timestamp / int /
float columns too,
- // not just VarChar / Decimal vectors that stay as raw Arrow types.
- def nullCheckCall(ord: Int): String = {
- val method =
CometBatchKernelCodegenInput.nullCheckMethod(inputSchema(ord))
- s"this.col$ord.$method(i)"
- }
- val nullCheck =
- if (inputOrdinals.isEmpty) "false"
- else inputOrdinals.map(nullCheckCall).mkString(" || ")
- // `NullIntolerant` only constrains "any input null -> output null";
it does NOT promise
- // that non-null inputs always produce non-null output.
`MakeTimestamp(failOnError=false)`
- // is `NullIntolerant=true` but its `doGenCode` catches
`DateTimeException` for invalid
- // year/month/day/hour/min/sec components and sets `ev.isNull = true`.
Honor `ev.isNull`
- // post-eval whenever the expression is nullable; skip the guard only
when the root is
- // statically non-nullable (`ev.isNull` is then a literal `false`).
- if (boundExpr.nullable) {
- s"""
- |if ($nullCheck) {
- | output.setNull(i);
- |} else {
- | $subExprsCode
- | ${ev.code}
- | if (${ev.isNull}) {
- | output.setNull(i);
- | } else {
- | $writeSnippet
- | }
- |}
+ val inputOrdinals = boundExpr.collect { case b: BoundReference =>
b.ordinal }.distinct
+ if (canShortCircuitNulls(boundExpr, inputOrdinals)) {
+ // Primitive Arrow vectors are wrapped in `CometPlainVector` at
input-cast time, which
+ // exposes `isNullAt(int)` rather than the raw Arrow `isNull(int)`. Pick
the right method
+ // for the ordinal so the short-circuit compiles for timestamp / int /
float columns too,
+ // not just VarChar / Decimal vectors that stay as raw Arrow types.
+ //
+ // Multi-ordinal trees test the disjunction of their ordinals. That is
only reachable for
+ // leaf-only-children roots, where it is exact; see
[[canShortCircuitNulls]].
+ val nullCheck = inputOrdinals
+ .map(ord =>
+
s"this.col$ord.${CometBatchKernelCodegenInput.nullCheckMethod(inputSchema(ord))}(i)")
+ .mkString(" || ")
+ // `NullIntolerant` only constrains "any input null -> output null"; it
does NOT promise
+ // that non-null inputs always produce non-null output.
`MakeTimestamp(failOnError=false)`
+ // is `NullIntolerant=true` but its `doGenCode` catches
`DateTimeException` for invalid
+ // year/month/day/hour/min/sec components and sets `ev.isNull = true`.
Honor `ev.isNull`
+ // post-eval whenever the expression is nullable; skip the guard only
when the root is
+ // statically non-nullable (`ev.isNull` is then a literal `false`).
+ if (boundExpr.nullable) {
+ s"""
+ |if ($nullCheck) {
+ | output.setNull(i);
+ |} else {
+ | $subExprsCode
+ | ${ev.code}
+ | if (${ev.isNull}) {
+ | output.setNull(i);
+ | } else {
+ | $writeSnippet
+ | }
+ |}
""".stripMargin
- } else {
- s"""
- |if ($nullCheck) {
- | output.setNull(i);
- |} else {
- | $subExprsCode
- | ${ev.code}
- | $writeSnippet
- |}
+ } else {
+ s"""
+ |if ($nullCheck) {
+ | output.setNull(i);
+ |} else {
+ | $subExprsCode
+ | ${ev.code}
+ | $writeSnippet
+ |}
""".stripMargin
- }
- case _ =>
- // NonNullableOutputShortCircuit: when `nullable = false`, drop the
`if (ev.isNull)`
- // guard at source level rather than relying on JIT folding.
- if (!boundExpr.nullable) {
- s"""
- |$subExprsCode
- |${ev.code}
- |$writeSnippet
+ }
+ } else {
+ // NonNullableOutputShortCircuit: when `nullable = false`, drop the `if
(ev.isNull)`
+ // guard at source level rather than relying on JIT folding.
+ if (!boundExpr.nullable) {
+ s"""
+ |$subExprsCode
+ |${ev.code}
+ |$writeSnippet
""".stripMargin
- } else {
- s"""
- |$subExprsCode
- |${ev.code}
- |if (${ev.isNull}) {
- | output.setNull(i);
- |} else {
- | $writeSnippet
- |}
+ } else {
+ s"""
+ |$subExprsCode
+ |${ev.code}
+ |if (${ev.isNull}) {
+ | output.setNull(i);
+ |} else {
+ | $writeSnippet
+ |}
""".stripMargin
- }
+ }
}
}
+ /**
+ * Gates the [[defaultBody]] null short-circuit ("any input null implies
null output", tested as
+ * the disjunction of the tree's input ordinals before `ev.code` runs).
+ *
+ * The short-circuit is only equivalent to Spark when no subtree that Spark
would have evaluated
+ * ahead of a null check gets skipped. Spark's null handling is per-node and
left-to-right:
+ * `BinaryExpression.nullSafeCodeGen` emits the left child's code
unconditionally, then tests
+ * the left child's null, then the right child's. So for `add_months(cast(s
as date), i)` on
+ * `('notadate', NULL)` Spark evaluates the cast and, under ANSI, raises
`CAST_INVALID_INPUT`; a
+ * short-circuit on the union of ordinals would skip the cast and return
null, swallowing the
+ * error (#5218).
+ *
+ * Conditions, all necessary:
+ *
+ * - The tree reads at least one input ordinal. A literal-only tree has
nothing to test.
+ * (Catalyst's `ConstantFolding` evaluates those at analysis time, so
this is defensive.)
+ * - The root is `NullIntolerant`, so a null input really does mean a null
result.
+ * - Every node in the tree is null-propagating ([[allNullIntolerant]]); a
`Coalesce` / `If` /
+ * `CaseWhen` anywhere would break the chain.
+ * - Either the tree reads exactly one ordinal, or every direct child of
the root is a leaf
+ * ([[rootChildrenAreLeaves]]). Both shapes make the short-circuit exact:
+ * - One ordinal: there is nothing left for Spark to evaluate ahead of
that ordinal's own
+ * null check.
+ * - Leaf-only children: the tree is one level deep, so the only code
Spark runs ahead of
+ * its null checks is `BoundReference` / `Literal` reads, which cannot
raise. Any error
+ * comes from the root's own logic, which runs only once every input
is known non-null --
+ * in both Spark's shape and ours. This keeps the fast path for the
common no-cast
+ * multi-argument case (`pmod(a, b)`, `conv(a, b, c)`,
`make_timestamp(y, m, d, h, mi,
+ * s)`) across the ~70 expressions that route through this dispatcher.
+ *
+ * The leaf-only rule relies on a `NullIntolerant` root not raising before
it has checked every
+ * input for null. Spark's own generated code honors that even where it
reorders children: the
+ * divide/remainder family (`DivModLike.doGenCode`, and `Pmod.doGenCode`,
which carries its own
+ * copy of the same shape) evaluates the divisor first, but still tests the
dividend's null
+ * *before* its divide/remainder-by-zero throw, so `pmod(NULL, 0)` returns
null rather than
+ * raising. A root that raised ahead of its null checks would contradict
`NullIntolerant`.
+ */
+ private def canShortCircuitNulls(expr: Expression, inputOrdinals: Seq[Int]):
Boolean =
+ inputOrdinals.nonEmpty && isNullIntolerant(expr) &&
allNullIntolerant(expr) &&
+ (inputOrdinals.size == 1 || rootChildrenAreLeaves(expr))
+
+ /**
+ * True iff every direct child of the root is a leaf (`BoundReference` or
`Literal`), i.e. the
+ * tree is one level deep with no subtree between the root and its inputs.
One of the conditions
+ * on [[canShortCircuitNulls]]: it is what makes a multi-ordinal
short-circuit exact, because a
+ * leaf read cannot raise an error that Spark's left-to-right evaluation
would have surfaced.
+ */
+ private def rootChildrenAreLeaves(expr: Expression): Boolean =
+ expr.children.forall {
+ case _: BoundReference | _: Literal => true
+ case _ => false
+ }
+
/**
* True iff every node in the tree propagates nulls (`NullIntolerant`,
`BoundReference`, or
- * `Literal`). Gates the [[defaultBody]] short-circuit, which is only
correct when no node
- * (`Coalesce`, `If`, `CaseWhen`, `Concat`, ...) breaks the propagation
chain.
+ * `Literal`). One of the conditions on [[canShortCircuitNulls]]: the
short-circuit is only
+ * correct when no node (`Coalesce`, `If`, `CaseWhen`, `Concat`, ...) breaks
the propagation
+ * chain.
*/
private def allNullIntolerant(expr: Expression): Boolean =
!expr.exists {
diff --git
a/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala
b/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala
index d258e9e6ef..957790d180 100644
---
a/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala
+++
b/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala
@@ -22,17 +22,23 @@ package org.apache.comet.codegen
import org.apache.spark.sql.catalyst.expressions.SpecializedGetters
import org.apache.spark.sql.types._
+import org.apache.comet.shims.CometTypeShim
+
/**
* Shared `SpecializedGetters.get(ordinal, dataType)` dispatch used by
[[CometInternalRow]] and
* [[CometArrayData]]. Spark codegen paths (notably `SafeProjection` for
ScalaUDF struct args) and
* interpreted-eval fallbacks (`ArrayDistinct.nullSafeEval` etc.) call the
generic `get` instead
* of the typed getter, so both kernel-side bases need a non-throwing
implementation.
*
+ * The type surface here must stay in step with
`CometBatchKernelCodegen.isSupportedDataType`,
+ * which is what gates the expression at plan time. A type accepted there but
missing here throws
+ * at execute time, after the plan has already committed to the kernel.
+ *
* For complex types, the typed getter allocates a fresh `InputStruct_*` /
`InputArray_*` /
* `InputMap_*` per call (`ColumnarRow`-style), so retain-by-reference
consumers like
* `OpenHashSet` get distinct identities.
*/
-private[codegen] object CometSpecializedGettersDispatch {
+private[codegen] object CometSpecializedGettersDispatch extends CometTypeShim {
def get(g: SpecializedGetters, ordinal: Int, dataType: DataType): AnyRef = {
if (g.isNullAt(ordinal)) return null
@@ -45,6 +51,10 @@ private[codegen] object CometSpecializedGettersDispatch {
case LongType | TimestampType | TimestampNTZType | _:
DayTimeIntervalType =>
java.lang.Long.valueOf(g.getLong(ordinal))
case CalendarIntervalType => g.getInterval(ordinal)
+ // Spark 4.1's `TimeType` is nanos-since-midnight in a long, read
through `getLong` just
+ // like the other long-backed temporal types. Matches
`elementGetterCall` and
+ // `emitSpecializedGetterExpr`, which both route time types to `getLong`.
+ case dt if isTimeType(dt) => java.lang.Long.valueOf(g.getLong(ordinal))
case FloatType => java.lang.Float.valueOf(g.getFloat(ordinal))
case DoubleType => java.lang.Double.valueOf(g.getDouble(ordinal))
case _: StringType => g.getUTF8String(ordinal)
diff --git
a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala
b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala
index 56746f5156..ef541fb3e2 100644
---
a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala
+++
b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala
@@ -185,6 +185,11 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
/**
* Build the compile-time spec for one input Arrow vector. Recurses on
complex types.
*
+ * The vector classes matched here must cover every Spark type
+ * `CometBatchKernelCodegen.isSupportedDataType` admits, because that
predicate is what
+ * `canHandle` gates the plan on. A type accepted at plan time but unmatched
here throws at
+ * execute time, when the operator can no longer fall back to Spark.
+ *
* Top-level `nullable=true` is hardcoded: the cache key does not specialize
on per-batch null
* density. Schema-declared nullability still reaches the kernel via
`BoundReference.nullable`
* embedded in `bytesKey`, so `BoundReference.doGenCode` elides its own
`isNullAt` probe on
@@ -219,8 +224,9 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
StructColumnSpec(nullable = true, fieldSpecs)
case _: BitVector | _: TinyIntVector | _: SmallIntVector | _: IntVector |
_: BigIntVector |
_: Float4Vector | _: Float8Vector | _: DecimalVector | _:
VarCharVector |
- _: VarBinaryVector | _: DateDayVector | _: DurationVector | _:
TimeStampMicroVector |
- _: TimeStampMicroTZVector | _: IntervalYearVector | _:
IntervalMonthDayNanoVector =>
+ _: VarBinaryVector | _: DateDayVector | _: DurationVector | _:
TimeNanoVector |
+ _: TimeStampMicroVector | _: TimeStampMicroTZVector | _:
IntervalYearVector |
+ _: IntervalMonthDayNanoVector =>
ScalarColumnSpec(v.getClass.asInstanceOf[Class[_ <: ValueVector]],
nullable = true)
case other =>
throw new UnsupportedOperationException(
diff --git
a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
index 87c52f153c..adc4440915 100644
--- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
@@ -24,7 +24,7 @@ import org.scalatest.funsuite.AnyFunSuite
import org.apache.spark.SparkConf
import org.apache.spark.serializer.JavaSerializer
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Add, AddMonths,
BoundReference, Coalesce, Concat, CreateArray, CreateMap, DateFormatClass,
ElementAt, Expression, GetStructField, LeafExpression, Length, Literal,
MakeTimestamp, MicrosToTimestamp, MillisToTimestamp, MonthsBetween,
Nondeterministic, Rand, Size, ToUnixTimestamp, Unevaluable, UnixMicros,
UnixMillis, UnixSeconds, Upper}
+import org.apache.spark.sql.catalyst.expressions.{Add, AddMonths,
BoundReference, Cast, Coalesce, Concat, CreateArray, CreateMap,
DateFormatClass, ElementAt, Expression, GetStructField, LeafExpression, Length,
Literal, MakeTimestamp, MicrosToTimestamp, MillisToTimestamp, MonthsBetween,
Nondeterministic, Rand, Size, ToUnixTimestamp, Unevaluable, UnixMicros,
UnixMillis, UnixSeconds, Upper}
import org.apache.spark.sql.catalyst.expressions.codegen.{CodeFormatter,
CodegenContext, CodegenFallback, ExprCode}
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String
@@ -41,8 +41,10 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen
* assert on the emitted Java directly, without invoking Janino. The goal is
to catch regressions
* in the optimizations we claim the dispatcher applies:
*
- * - `NullIntolerant` short-circuit wraps `ev.code` in `if (any-input-null)
{ setNull } else {
- * ev.code; write }`.
+ * - `NullIntolerant` short-circuit wraps `ev.code` in `if (input-null) {
setNull } else {
+ * ev.code; write }`. It is emitted for a tree reading a single ordinal,
and for a
+ * multi-ordinal tree whose root has only leaf children; it is *not*
emitted once a subtree
+ * sits between the root and a leaf (see the multi-input tests for why).
* - Non-nullable column declaration emits `return false;` from
`isNullAt(ord)`, and a
* `BoundReference.nullable=false` (Catalyst sets this from
schema-declared nullability) makes
* Spark's `doGenCode` skip emitting its own `row.isNullAt(ord)` probe
entirely.
@@ -166,6 +168,53 @@ class CometCodegenSourceSuite extends AnyFunSuite {
s"got:\n$src")
}
+ test("NullIntolerant short-circuit skipped when a subtree sits between root
and leaf (#5218)") {
+ // Even when every node is NullIntolerant, a short-circuit on the *union*
of input ordinals is
+ // not equivalent to Spark once a real subtree sits between the root and a
leaf. Spark's null
+ // handling is per-node and left-to-right:
`BinaryExpression.nullSafeCodeGen` emits the left
+ // child's code unconditionally, then tests the left child's null, then
the right child's. So
+ // for `Add(cast(c0 AS INT), c1)` on a row where c0 = 'abc' and c1 IS
NULL, Spark evaluates the
+ // cast first and (under ANSI) raises CAST_INVALID_INPUT. A union
short-circuit would skip the
+ // cast and return null, swallowing the error. See CometCodegenSuite's
add_months test for the
+ // end-to-end witness.
+ //
+ // The `Cast` is what disqualifies the short-circuit here: it is a
non-leaf child of the root
+ // that can raise on its own. The sibling test below covers the
leaf-only-children shape, which
+ // keeps the short-circuit.
+ val strCol = ArrowColumnSpec(varCharVectorClass, nullable = true)
+ val intCol = ArrowColumnSpec(
+ CometBatchKernelCodegen.vectorClassBySimpleName("IntVector"),
+ nullable = true)
+ val expr = Add(
+ Cast(BoundReference(0, StringType, nullable = true), IntegerType,
ansiEnabled = true),
+ BoundReference(1, IntegerType, nullable = true))
+ val src = gen(expr, strCol, intCol)
+ assert(
+ !src.contains("this.col0.isNull(i) || this.col1.isNullAt(i)"),
+ s"expected no union-of-inputs short-circuit when a Cast sits under the
root; got:\n$src")
+ assert(
+ !src.contains("if (this.col0.isNull(i))") && !src.contains("if
(this.col1.isNullAt(i))"),
+ s"expected no pre-eval input-null short-circuit at all for this shape;
got:\n$src")
+ }
+
+ test("NullIntolerant short-circuit kept for a multi-input leaf-only tree
(#5218)") {
+ // Counterpart to the test above. When every direct child of the root is a
`BoundReference` or
+ // `Literal`, the tree is one level deep: the only code Spark runs ahead
of its own null checks
+ // is leaf reads, which cannot raise. Any error comes from the root's own
logic, which runs only
+ // once every input is known non-null -- in Spark's shape and in ours
alike. So the union
+ // short-circuit is exact and worth keeping for the common no-cast
multi-argument case.
+ val intCol = ArrowColumnSpec(
+ CometBatchKernelCodegen.vectorClassBySimpleName("IntVector"),
+ nullable = true)
+ val expr = Add(
+ BoundReference(0, IntegerType, nullable = true),
+ BoundReference(1, IntegerType, nullable = true))
+ val src = gen(expr, intCol, intCol)
+ assert(
+ src.contains("if (this.col0.isNullAt(i) || this.col1.isNullAt(i))"),
+ s"expected union-of-inputs short-circuit for a leaf-only two-input tree;
got:\n$src")
+ }
+
test("canHandle accepts CodegenFallback expressions (delegates to
eval(row))") {
// CodegenFallback.doGenCode emits ((Expression) references[N]).eval(row)
which is the same
// mechanism that backs HigherOrderFunction support: the eval reads
through the kernel's typed
@@ -404,6 +453,9 @@ class CometCodegenSourceSuite extends AnyFunSuite {
// non-null output and dropped the post-eval guard; that wrote stale
`ev.value` bytes for
// every invalid row. The short-circuit on input nulls must coexist with a
post-eval
// `if (ev.isNull) setNull` check whenever the expression itself is
nullable.
+ //
+ // All six children are `BoundReference`s, so this is a leaf-only tree and
keeps the
+ // union-of-ordinals short-circuit (#5218).
val expr = MakeTimestamp(
BoundReference(0, IntegerType, nullable = true),
BoundReference(1, IntegerType, nullable = true),
@@ -428,12 +480,83 @@ class CometCodegenSourceSuite extends AnyFunSuite {
val formatted = CodeFormatter.format(result.code)
// Two distinct setNull sites must exist: the input-null short-circuit
before `ev.code` runs,
// and the post-eval guard that propagates `ev.isNull = true` set by
MakeTimestamp's catch
- // block on invalid components. Pre-fix there was only one (the input
short-circuit).
+ // block on invalid components. Pre-#4554 there was only one (the input
short-circuit).
val setNullOccurrences = "output\\.setNull\\(i\\);".r.findAllIn(src).length
assert(
setNullOccurrences >= 2,
"expected at least two setNull sites (input short-circuit + post-eval
ev.isNull guard); " +
s"found $setNullOccurrences. Source:\n$formatted")
+ // The short-circuit must test every ordinal the tree reads, not just the
first.
+ assert(
+ src.contains(
+ "if (this.col0.isNullAt(i) || this.col1.isNullAt(i) || " +
+ "this.col2.isNullAt(i) || this.col3.isNullAt(i) ||
this.col4.isNullAt(i) || " +
+ "this.col5.isNull(i))"),
+ s"expected the short-circuit to test all six ordinals;
source:\n$formatted")
+ }
+
+ test("multi-input tree with a non-leaf child keeps only the post-eval guard
(#5218)") {
+ // The counterpart to the test above: the same
nullable-NullIntolerant-root shape, but with a
+ // `Cast` between the root and one leaf. That disqualifies the
short-circuit, so exactly one
+ // setNull site remains -- the post-eval `ev.isNull` guard, which is the
actual #4554
+ // regression and must survive independently of the short-circuit.
+ val expr = MakeTimestamp(
+ Cast(BoundReference(0, StringType, nullable = true), IntegerType,
ansiEnabled = true),
+ BoundReference(1, IntegerType, nullable = true),
+ BoundReference(2, IntegerType, nullable = true),
+ Literal(0),
+ Literal(0),
+ Literal(Decimal(0, 8, 6), DecimalType(8, 6)),
+ timezone = None,
+ timeZoneId = Some("UTC"),
+ failOnError = false)
+ assert(expr.nullable, "MakeTimestamp(failOnError=false) must be nullable
for this test")
+ val strCol = ArrowColumnSpec(varCharVectorClass, nullable = true)
+ val intCol = ArrowColumnSpec(
+ CometBatchKernelCodegen.vectorClassBySimpleName("IntVector"),
+ nullable = true)
+ val result =
+ CometBatchKernelCodegen.generateSource(expr, IndexedSeq(strCol, intCol,
intCol))
+ val src = result.body
+ val formatted = CodeFormatter.format(result.code)
+ val setNullOccurrences = "output\\.setNull\\(i\\);".r.findAllIn(src).length
+ assert(
+ setNullOccurrences == 1,
+ "expected exactly one setNull site (the post-eval ev.isNull guard, with
no short-circuit); " +
+ s"found $setNullOccurrences. Source:\n$formatted")
+ assert(
+ !src.contains("if (this.col0.isNull(i))"),
+ s"expected no input-null short-circuit when a Cast sits under the root;
source:\n$formatted")
+ }
+
+ test("single-input nullable NullIntolerant root emits both short-circuit and
post-eval guard") {
+ // Counterpart to the test above, for the shape that does keep the
short-circuit: one input
+ // ordinal. Two distinct setNull sites must exist -- the input-null
short-circuit before
+ // `ev.code` runs, and the post-eval `ev.isNull` guard.
+ val expr = MakeTimestamp(
+ BoundReference(0, IntegerType, nullable = true),
+ Literal(1),
+ Literal(1),
+ Literal(0),
+ Literal(0),
+ Literal(Decimal(0, 8, 6), DecimalType(8, 6)),
+ timezone = None,
+ timeZoneId = Some("UTC"),
+ failOnError = false)
+ assert(expr.nullable, "MakeTimestamp(failOnError=false) must be nullable
for this test")
+ val intCol = ArrowColumnSpec(
+ CometBatchKernelCodegen.vectorClassBySimpleName("IntVector"),
+ nullable = true)
+ val result = CometBatchKernelCodegen.generateSource(expr,
IndexedSeq(intCol))
+ val src = result.body
+ assert(
+ src.contains("if (this.col0.isNullAt(i))"),
+ s"expected input-null short-circuit for the single-input tree;
got:\n$src")
+ val setNullOccurrences = "output\\.setNull\\(i\\);".r.findAllIn(src).length
+ assert(
+ setNullOccurrences >= 2,
+ "expected two setNull sites (input short-circuit + post-eval ev.isNull
guard); " +
+ s"found $setNullOccurrences.
Source:\n${CodeFormatter.format(result.code)}")
}
test("ArrayType(StringType) output emits ListVector startNewValue/endValue
recursion") {
diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
index afb61d09cc..5806cb3501 100644
--- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
@@ -22,7 +22,7 @@ package org.apache.comet
import scala.util.Random
import org.apache.arrow.vector._
-import org.apache.spark.{SparkConf, TaskContext}
+import org.apache.spark.{SparkConf, SparkEnv, TaskContext}
import org.apache.spark.sql.CometTestBase
import org.apache.spark.sql.api.java.UDF1
import org.apache.spark.sql.catalyst.expressions.{BoundReference, CreateArray,
CreateMap, CreateNamedStruct, Expression, Literal, MapConcat}
@@ -31,6 +31,7 @@ import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String
+import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus
import org.apache.comet.codegen.CometBatchKernelCodegen
import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec
import org.apache.comet.udf.codegen.CometScalaUDFCodegen
@@ -1443,6 +1444,145 @@ class CometCodegenSuite
}
}
+ test("multi-input NullIntolerant tree does not swallow an ANSI error
(#5218)") {
+ // `add_months` is NullIntolerant and a plain BinaryExpression, so Spark's
`nullSafeCodeGen`
+ // emits the LEFT child's code unconditionally before testing the right
child's null. On the
+ // ('notadate', NULL) row Spark therefore evaluates the cast and raises
CAST_INVALID_INPUT.
+ // The dispatcher used to short-circuit on the union of input ordinals,
see the null on `i`,
+ // and return NULL -- silently losing the error. Note `pmod` / `div` are
not witnesses here:
+ // `DivModLike` deliberately evaluates its right child first, so Spark
also returns NULL.
+ withTable("t") {
+ sql("CREATE TABLE t (s STRING, i INT) USING parquet")
+ sql("INSERT INTO t VALUES ('notadate', NULL), ('2024-01-31', 1)")
+ withSQLConf("spark.sql.ansi.enabled" -> "true") {
+ CometScalaUDFCodegen.resetStats()
+ val (sparkErr, cometErr) =
+ checkSparkAnswerMaybeThrows(sql("SELECT add_months(CAST(s AS DATE),
i) FROM t"))
+ val stats = CometScalaUDFCodegen.stats()
+ assert(
+ stats.compileCount + stats.cacheHitCount >= 1,
+ s"expected the codegen dispatcher to run for this query, got $stats")
+ assert(
+ sparkErr.isDefined,
+ "expected Spark to raise on the invalid ANSI cast; the test row is
no longer a witness")
+ assert(
+ cometErr.isDefined,
+ "Comet returned a value where Spark raised: the null short-circuit
swallowed the error")
+ assert(
+ cometErr.get.getMessage.contains("CAST_INVALID_INPUT"),
+ s"expected the same CAST_INVALID_INPUT error Spark raises, got:
${cometErr.get}")
+ }
+ }
+ }
+
+ test("single-input NullIntolerant tree still short-circuits nulls") {
+ // Guards against over-correcting #5218: the single-ordinal short-circuit
is exact (Spark also
+ // evaluates nothing when that one input is null) and must be preserved.
+ withSubjects("abc", null, "xyz") {
+ assertCodegenRan {
+ checkSparkAnswerAndOperator(sql("SELECT upper(substring(s, 1, 2)) FROM
t"))
+ }
+ }
+ }
+
+ test("multi-input leaf-only NullIntolerant tree short-circuits nulls
correctly (#5218)") {
+ // The leaf-only-children shape keeps the union-of-ordinals short-circuit,
because the only
+ // code Spark runs ahead of its own null checks is `BoundReference` reads.
Exercises every
+ // null combination across two ordinals to confirm the disjunction matches
Spark's per-node
+ // left-to-right null handling row for row. Wrapped in a UDF so the
argument expression is
+ // guaranteed to route through the dispatcher rather than Comet's native
path.
+ spark.udf.register("idInt", (i: Integer) => i)
+ withTable("t") {
+ sql("CREATE TABLE t (a INT, b INT) USING parquet")
+ sql(
+ "INSERT INTO t VALUES (7, 3), (NULL, 3), (7, NULL), (NULL, NULL), " +
+ "(-7, 3), (7, -3), (0, 3)")
+ assertCodegenRan {
+ checkSparkAnswerAndOperator(sql("SELECT idInt(pmod(a, b)) FROM t"))
+ }
+ assertCodegenRan {
+ checkSparkAnswerAndOperator(sql("SELECT idInt(a + b) FROM t"))
+ }
+ }
+ }
+
+ test("leaf-only short-circuit preserves ANSI remainder-by-zero behaviour
(#5218)") {
+ // The one shape where a `NullIntolerant` root's error check is not simply
gated behind
+ // "all inputs non-null": `Pmod.doGenCode` evaluates the divisor first and
throws on a zero
+ // divisor under ANSI. It still tests the dividend's null *before* that
throw, so
+ // `pmod(NULL, 0)` returns NULL in Spark and the union short-circuit stays
exact. The
+ // (NULL, 0) row pins the non-raising case, the (7, 0) row the raising one.
+ //
+ // Spark 4.1 introduced REMAINDER_BY_ZERO; older versions raise
DIVIDE_BY_ZERO for `pmod`.
+ // The error comes from Spark's own generated code running inside the
kernel, so the class
+ // tracks the Spark version under test.
+ val expectedError = if (isSpark41Plus) "REMAINDER_BY_ZERO" else
"DIVIDE_BY_ZERO"
+ spark.udf.register("idInt", (i: Integer) => i)
+ withTable("t") {
+ sql("CREATE TABLE t (a INT, b INT) USING parquet")
+ sql("INSERT INTO t VALUES (NULL, 0)")
+ withSQLConf("spark.sql.ansi.enabled" -> "true") {
+ assertCodegenRan {
+ checkSparkAnswerAndOperator(sql("SELECT idInt(pmod(a, b)) FROM t"))
+ }
+ }
+ sql("INSERT INTO t VALUES (7, 0)")
+ withSQLConf("spark.sql.ansi.enabled" -> "true") {
+ val (sparkErr, cometErr) =
+ checkSparkAnswerMaybeThrows(sql("SELECT idInt(pmod(a, b)) FROM t"))
+ assert(
+ sparkErr.isDefined,
+ "expected Spark to raise on pmod by zero under ANSI; the test row is
no longer a witness")
+ assert(
+ cometErr.isDefined,
+ "Comet returned a value where Spark raised: the null short-circuit
swallowed the error")
+ assert(
+ cometErr.get.getMessage.contains(expectedError),
+ s"expected the same $expectedError error Spark raises, got:
${cometErr.get}")
+ assert(
+ sparkErr.get.getMessage.contains(expectedError),
+ s"expected Spark to raise $expectedError, got: ${sparkErr.get}")
+ }
+ }
+ }
+
+ test("TIME input column routes through the dispatcher (#5218)") {
+ assume(isSpark41Plus, "TimeType requires Spark 4.1+")
+ // `canHandle` accepts TIME (`isSupportedDataType`) and `emitTypedGetters`
emits a getLong
+ // case for `TimeNanoVector`, but `CometScalaUDFCodegen.specFor` used to
omit the vector class
+ // and throw `UnsupportedOperationException` at execute time -- after the
plan had already
+ // committed to the kernel, so there was no fallback. Driven through the
dispatcher directly
+ // because Spark 4.1 still rejects TIME columns in file-based data
sources, so no SQL query
+ // can produce a TIME input today.
+ val timeVec = new TimeNanoVector("tm", CometArrowAllocator)
+ val exprVec = new VarBinaryVector("expr", CometArrowAllocator)
+ var out: ValueVector = null
+ try {
+ timeVec.allocateNew()
+ timeVec.setSafe(0, 45296000000000L) // 12:34:56
+ timeVec.setNull(1)
+ timeVec.setValueCount(2)
+
+ val timeType =
org.apache.spark.sql.comet.util.Utils.fromArrowField(timeVec.getField)
+ val expr = BoundReference(0, timeType, nullable = true)
+ val serialized =
SparkEnv.get.closureSerializer.newInstance().serialize(expr)
+ val bytes = new Array[Byte](serialized.remaining())
+ serialized.get(bytes)
+ exprVec.allocateNew()
+ exprVec.setSafe(0, bytes)
+ exprVec.setValueCount(1)
+
+ out = new CometScalaUDFCodegen().evaluate(Array(exprVec, timeVec), 2)
+ val comet = CometVector.getVector(out.asInstanceOf[FieldVector], null)
+ assert(comet.getLong(0) === 45296000000000L)
+ assert(comet.isNullAt(1))
+ } finally {
+ if (out != null) out.close()
+ exprVec.close()
+ timeVec.close()
+ }
+ }
+
// Runtime coverage for nullable nested `getStruct` / `getArray` / `getMap`
element reads is
// exercised through HOFs in `CometCodegenHOFSuite`. Static emitter
assertions live in
// `CometCodegenSourceSuite`.
diff --git
a/spark/src/test/scala/org/apache/comet/codegen/CometSpecializedGettersDispatchSuite.scala
b/spark/src/test/scala/org/apache/comet/codegen/CometSpecializedGettersDispatchSuite.scala
new file mode 100644
index 0000000000..7ad93fd9fc
--- /dev/null
+++
b/spark/src/test/scala/org/apache/comet/codegen/CometSpecializedGettersDispatchSuite.scala
@@ -0,0 +1,74 @@
+/*
+ * 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.
+ */
+
+package org.apache.comet.codegen
+
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.arrow.vector.types.TimeUnit
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.types.{DataType, DateType, LongType}
+
+import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus
+
+/**
+ * Unit coverage for the generic `SpecializedGetters.get(ordinal, dataType)`
dispatch the kernel
+ * bases delegate to. Spark reaches it from `SafeProjection` (for ScalaUDF
struct arguments) and
+ * from every `CodegenFallback.eval(row)`, including higher-order functions,
so its type surface
+ * has to match what `CometBatchKernelCodegen.isSupportedDataType` admits at
plan time. A type
+ * accepted there but missing here throws at execute time, when there is no
longer a fallback.
+ */
+class CometSpecializedGettersDispatchSuite extends AnyFunSuite {
+
+ /** Minimal `CometInternalRow` over one fixed long, to drive the dispatch
through `get`. */
+ private class StubRow(value: Long, isNull: Boolean = false) extends
CometInternalRow {
+ override def numFields: Int = 1
+ override def isNullAt(ordinal: Int): Boolean = isNull
+ override def getLong(ordinal: Int): Long = value
+ override def getInt(ordinal: Int): Int = value.toInt
+ }
+
+ /** Spark `DataType` for Arrow `Time(NANOSECOND, 64)`, resolved without a
4.1 compile dep. */
+ private def timeType: DataType =
+ Utils.fromArrowField(
+ new Field(
+ "t",
+ FieldType.nullable(new ArrowType.Time(TimeUnit.NANOSECOND, 64)),
+ java.util.Collections.emptyList[Field]()))
+
+ test("get dispatches TimeType through getLong (#5218)") {
+ assume(isSpark41Plus, "TimeType requires Spark 4.1+")
+ val row = new StubRow(45296000000000L) // 12:34:56
+ assert(row.get(0, timeType) === java.lang.Long.valueOf(45296000000000L))
+ }
+
+ test("get returns null for a null TimeType slot") {
+ assume(isSpark41Plus, "TimeType requires Spark 4.1+")
+ val row = new StubRow(45296000000000L, isNull = true)
+ assert(row.get(0, timeType) == null)
+ }
+
+ test("get still dispatches the long- and int-backed temporal types") {
+ // Guards against the TimeType case being inserted ahead of a case it
would shadow.
+ val row = new StubRow(19723L) // 2024-01-31 as an epoch day
+ assert(row.get(0, LongType) === java.lang.Long.valueOf(19723L))
+ assert(row.get(0, DateType) === java.lang.Integer.valueOf(19723))
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
