andygrove opened a new pull request, #7529:
URL: https://github.com/apache/datafusion-comet/pull/7529
## Which issue does this PR close?
N/A. Follow-up to #6180 (native `AtLeastNNonNulls`).
## Rationale for this change
`AtLeastNNonNulls` (what `df.na.drop(n)` plans) falls back to Spark with
`AtLeastNNonNulls would skip a subexpression that Spark evaluates for every
row` whenever two of its children share a subexpression, even when that
subexpression cannot fail. A common case is `na.drop` over columns derived from
one input:
```scala
df.withColumn("prefix", substr($"s", 1, 4))
.withColumn("is_recent", substr($"s", 1, 4).cast("int") >= 16)
.withColumn("is_latest", substr($"s", 1, 4).cast("int") >= 18)
.na.drop(1)
```
The filter runs on Spark, with a transition on each side, although every
expression in it is native.
The guard exists because Spark evaluates a shared subexpression for every
row, so a shared expression that fails (an ANSI cast of a malformed string,
say) raises even on a row where `AtLeastNNonNulls` has already reached n.
Native evaluation that skips the later children would not raise. If the shared
subexpression cannot fail, skipping it shows nothing, so the guard has no
reason to apply. The same check (`canSkipEvaluation`) also decides the null
short-circuit of function arguments, so this applies to them too.
## What changes are included in this PR?
- `SharedSubexpressions.common` no longer contains a shared subexpression
that cannot raise. A new `cannotRaise` check lists what cannot fail whatever
the row holds: literals and column references, struct field reads, comparisons
and boolean operators, `coalesce`, `if` and `CASE`, `substr`, `upper`, `lower`,
`length`, `startswith`, `endswith` and `contains`, and casts and arithmetic
outside ANSI mode.
- The list fails closed. Anything it does not name counts as able to raise,
including `element_at` (the index 0 raises in every mode), an ANSI cast or
arithmetic, and a decimal expression after promotion.
- The two `QueryPlanSerdeSuite` tests that used a string-to-int cast as the
shared expression now use `element_at(arr, 0)`, which raises in every mode, so
they still pin the mechanism. The cast is no longer a case where the guard
applies outside ANSI mode.
## How are these changes tested?
- `QueryPlanSerdeSuite`: a shared `length`, a shared struct field, and a
shared expression with a raising node below an expression that cannot raise,
each checked with `canShortCircuitNulls`. A shared `element_at(arr, 0)` still
blocks the skip, also when other shared subexpressions that cannot raise are
present.
- `CometExpressionSuite`: two `AtLeastNNonNulls` queries that share
`upper(s)`, and a struct field, with subexpression elimination on. Both now run
natively and match Spark. The two existing tests, which share an ANSI cast that
fails and an ANSI divide by zero, still fall back and still raise.
- `CometSqlFileTestSuite` (769 tests) and `CometExpressionSuite`
`AtLeastNNonNulls` (8 tests) pass on Spark 3.5. I have not run the other Spark
profiles.
--
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]