andygrove opened a new pull request, #5219:
URL: https://github.com/apache/datafusion-comet/pull/5219
## Which issue does this PR close?
Closes #5218.
## Rationale for this change
`CometBatchKernelCodegen.defaultBody` applied an "any input null implies
null output" short-circuit on the **union of input ordinals** whenever every
node in the bound tree was `NullIntolerant`:
```java
if (this.col0.isNull(i) || this.col1.isNullAt(i)) { output.setNull(i); }
else { ev.code; write }
```
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. A short-circuit over
the union of ordinals therefore skips a subtree Spark would have evaluated, and
with it any error that subtree raises.
Under ANSI, `add_months(CAST(s AS DATE), i)` on the row `('notadate', NULL)`
raises `CAST_INVALID_INPUT` in Spark. Comet returned `NULL`. ANSI is on by
default in Spark 4, and roughly 70 built-in expressions route through
`CometCodegenDispatch` (math, arrays, maps, strings, json, csv, xpath,
datetime), many of them `NullIntolerant` with error paths: `AddMonths`,
`MonthsBetween`, `MakeTimestamp`, `GetTimestamp`, `MakeDTInterval`, `ToNumber`,
`Conv`, `WidthBucket`, `Pmod`.
Values were never wrong, only errors were lost: a `NullIntolerant` node that
ignored one of its inputs would already be a Spark bug. I checked `nanvl`
specifically, since it ignores its second argument when the first is non-NaN,
and it is correctly *not* `NullIntolerant`.
## What changes are included in this PR?
**The short-circuit fix.** `canShortCircuitNulls` now requires the tree to
read **exactly one** input ordinal, in addition to the existing
`NullIntolerant` root and `allNullIntolerant` tree conditions. With a single
input there is nothing left for Spark to evaluate ahead of that ordinal's own
null check, so the short-circuit is exact and the common unary shapes (`upper`,
`length`, `date_format`) keep the fast path. Multi-input trees fall back to the
plain `ev.code` plus `ev.isNull` body, which is what Spark's own generated code
does.
**Two TIME-type gaps in the same dispatcher**, where the plan-time gate
accepted a type the runtime rejected. Because `canHandle` greenlights the
expression before the plan commits, both were execute-time failures with no
fallback:
- `CometScalaUDFCodegen.specFor` omitted `TimeNanoVector` and fell into
`case other => throw new UnsupportedOperationException`, even though
`isSupportedDataType` accepts time types, `primitiveArrowClasses` includes the
vector class, and `emitTypedGetters` emits a `getLong` case for it.
- `CometSpecializedGettersDispatch.get` had no time-type branch, so
`get(ordinal, TimeType)` threw. That is the path `SafeProjection` (for ScalaUDF
struct arguments) and every `CodegenFallback.eval(row)` use, including
higher-order functions. `emitSpecializedGetterExpr` and `elementGetterCall`
already routed time types to `getLong`, so only the generic dispatch was
missing.
Both are latent today: Spark 4.1.3 rejects TIME columns in file-based data
sources (`UNSUPPORTED_TIME_TYPE`), so no SQL query can produce a TIME input to
the dispatcher yet. They become live as soon as Spark supports TIME in Parquet
or a native operator emits a TIME column. Doc comments were added on `specFor`
and the dispatch object noting that their type surface has to stay in step with
`isSupportedDataType`.
## How are these changes tested?
Four new tests, each verified to fail against the pre-fix behaviour and pass
after:
- `CometCodegenSuite`: "multi-input NullIntolerant tree does not swallow an
ANSI error (#5218)" runs the `add_months(CAST(s AS DATE), i)` query through
`checkSparkAnswerMaybeThrows` under ANSI and asserts Comet raises the same
`CAST_INVALID_INPUT` Spark does, plus asserts the dispatcher actually ran for
the query. Pre-fix this returned `cometErr=None`.
- `CometCodegenSuite`: "single-input NullIntolerant tree still
short-circuits nulls" guards against over-correcting.
- `CometCodegenSourceSuite`: "NullIntolerant short-circuit skipped for a
multi-input tree (#5218)" pins the emitted source shape.
- `CometCodegenSuite`: "TIME input column routes through the dispatcher
(#5218)" drives `CometScalaUDFCodegen.evaluate` with a `TimeNanoVector`
directly, since no SQL query can reach it. Gated on `isSpark41Plus`.
- New
`spark/src/test/scala/org/apache/comet/codegen/CometSpecializedGettersDispatchSuite.scala`
covers `get` for `TimeType`, the null slot, and the neighbouring
long/int-backed temporal types so a future reordering that shadows a case is
caught.
An existing test needed updating: "nullable NullIntolerant root keeps
post-eval isNull guard (#4554)" asserted two `setNull` sites, one of which was
the multi-input short-circuit this PR removes. It now asserts exactly one site
(the post-eval `ev.isNull` guard, which is the actual #4554 regression) plus
the absence of the short-circuit, and a new sibling test covers the
single-input shape where both sites still appear.
Suite runs on the default profile (Spark 4.1 / Scala 2.13, JDK 17):
- `CometCodegenSuite`, `CometCodegenSourceSuite`, `CometCodegenHOFSuite`,
`CometCodegenFuzzSuite`, `CometSpecializedGettersDispatchSuite`: 171 tests, 0
failures.
- Because the change touches the shared code path for every
`CometCodegenDispatch` expression, also ran `CometTemporalExpressionSuite`,
`CometStringExpressionSuite`, `CometMathExpressionSuite`,
`CometArrayExpressionSuite`, `CometMapExpressionSuite`,
`CometJsonExpressionSuite`, `CometCsvExpressionSuite`: 137 tests, 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]