andygrove commented on code in PR #5219:
URL: https://github.com/apache/datafusion-comet/pull/5219#discussion_r3705580983
##########
spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala:
##########
@@ -351,82 +350,94 @@ 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.
+ val ord = inputOrdinals.head
+ val nullCheck =
+
s"this.col$ord.${CometBatchKernelCodegenInput.nullCheckMethod(inputSchema(ord))}(i)"
+ // `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. Three conditions, all
necessary:
+ *
+ * - The tree reads exactly one input ordinal. 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
+ * short-circuit on the union of several ordinals is not equivalent: it
skips a subtree that
+ * Spark would have evaluated, and with it any error that subtree
raises. Under ANSI,
+ * `add_months(cast(s as date), i)` on `('notadate', NULL)` raises
`CAST_INVALID_INPUT` in
+ * Spark, so returning null here would be wrong. With a single ordinal
there is nothing left
+ * for Spark to evaluate ahead of that ordinal's own null check, so the
short-circuit is
+ * exact. (A literal-only subtree that raises would be the one
exception, but Catalyst's
+ * `ConstantFolding` evaluates those at analysis time and never leaves
one in the tree.)
+ * - 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.
+ */
+ private def canShortCircuitNulls(expr: Expression, inputOrdinals: Seq[Int]):
Boolean =
+ inputOrdinals.size == 1 && isNullIntolerant(expr) &&
allNullIntolerant(expr)
Review Comment:
Done in dd6ae3225 — `canShortCircuitNulls` now accepts either exactly one
input ordinal or a leaf-only-children root, and `defaultBody` emits the
disjunction over every ordinal the tree reads instead of just the first. That
restores the fast path for the no-cast multi-argument shapes.
One wrinkle worth recording, because it makes the rule correct for a
slightly different reason than stated above. The argument was that with
leaf-only children "the only place an error can occur is the root's own `eval`,
and that only runs in the `else` branch either way." That isn't quite true for
`pmod`: it is one of the expressions that deliberately reorders its children,
so its error check is *not* simply gated behind all-inputs-non-null.
`Pmod.doGenCode` (it extends `BinaryArithmetic` directly and carries its own
copy of the `DivModLike` shape) evaluates the divisor first and throws
`REMAINDER_BY_ZERO` under ANSI:
```java
${eval2.code} // divisor
if (${eval2.isNull}) {
${ev.isNull} = true;
} else {
${eval1.code} // dividend
if (${eval1.isNull}) {
${ev.isNull} = true; // <-- reached before the throw
} else {
if ($isZero) throw ...remainderByZeroError(...);
$result
}
}
```
The rule survives, but because the *dividend's null check precedes the
throw*, not because the throw is gated on all inputs being non-null. So
`pmod(NULL, 0)` returns NULL in Spark, and the union short-circuit stays exact
for it. Had the throw been hoisted above the `eval1.isNull` test, `pmod(a, b)`
with no cast on either argument would have been a swallowed-error case under
the leaf-only rule.
Since that is load-bearing and invisible from Comet's side, I documented it
on `canShortCircuitNulls` and pinned it from both directions in
`CometCodegenSuite`: `(NULL, 0)` must return NULL, `(7, 0)` must raise
`REMAINDER_BY_ZERO`. If a future Spark version reorders that throw, the second
assertion fails rather than the behaviour silently regressing.
Also added: source-shape tests for both the kept (leaf-only) and skipped
(subtree-under-root) shapes, and a runtime test over every null combination
across two ordinals. The #4554 test goes back to its original two-`setNull`
assertion, since its `MakeTimestamp` over six `BoundReference`s is leaf-only
and keeps the short-circuit; it now also asserts all six ordinals get tested. A
new sibling test puts a `Cast` under the same root to pin the single-`setNull`
shape, so the #4554 post-eval guard is still verified independently of the
short-circuit.
I confirmed the new assertions actually discriminate: reverting the gate to
the single-ordinal-only rule fails exactly the two tests that pin the restored
fast path and nothing else. 175 tests across the 5 codegen suites and 137
across the dispatcher-backed expression suites, 0 failures.
--
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]
