[andygrove]

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:

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 numeric arithmetic outside ANSI mode. Under ANSI mode, which is the default from Spark 4, a cast or arithmetic that can fail still counts as able to raise, so the filter above keeps falling back there.
• 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, interval arithmetic (it raises on overflow in every mode), and a decimal expression after promotion (which wraps it in CheckOverflow).
• 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 and a shared struct field can be skipped. A shared element_at(arr, 0), a shared interval addition, and an expression that cannot raise but holds a raising one still block 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 the AtLeastNNonNulls tests of CometExpressionSuite (8 tests) pass on Spark 3.5. QueryPlanSerdeSuite and the AtLeastNNonNulls tests also pass on Spark 3.4, and the main and test code compile on Spark 3.4, 4.0 and the default profile. I have not run the tests on 4.x.
• The example above runs as a Spark Filter on current main and as a CometFilter with this change.

Related

#7322 has a similar fail-closed check, private to CometArrayFilter, for whether a lambda body can raise. If both land they can share one classifier; I kept them apart so neither PR waits on the other.

[comphead]

Thanks, the fail-closed list reads well and the tests pin both directions. One gap in the Cast arm.

[comphead]

A non-ANSI cast can still raise. Cast from DATE to TIMESTAMP or TIMESTAMP_NTZ goes through daysToMicros and instantToMicros, which use Math.multiplyExact, so a date past about year 294247 (dates go up to 5881580) throws ArithmeticException in every eval mode. With such a cast shared between na.drop children, Spark raises on every row while native evaluation would skip it. The comment on keyCastCannotFail in CometCast.scala already calls out these date and timestamp conversions. Could this arm also exclude casts to or from date, timestamp and interval types, so the motivating string to int cast keeps the skip and the list stays closed for the conversions that overflow? I have not run it.