andygrove commented on code in PR #6715:
URL: https://github.com/apache/datafusion-comet/pull/6715#discussion_r4210898156
##########
spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala:
##########
@@ -58,6 +58,40 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
override def convert(expr: ScalaUDF, inputs: Seq[Attribute], binding:
Boolean): Option[Expr] =
emitJvmCodegenDispatch(expr, inputs, binding)
+ /**
+ * Emits Spark's null guard around a Scala UDF, together with the UDF, as
one dispatcher call.
+ * Returns `None` when `expr` is not that guard or the dispatcher cannot
take it, and records
+ * nothing on `expr` in that case: the caller converts `expr` as usual,
which sends the UDF to
+ * the dispatcher on its own, and the UDF records its own fallback reason if
it cannot go.
+ *
+ * Spark's `HandleNullInputsForUDF` rule wraps a UDF with a primitive
parameter over a nullable
+ * input as `if (isnull(a) or isnull(b)) null else f(knownnotnull(a),
knownnotnull(b))`, one
+ * `IsNull` per such argument. Converted natively, the `If` is a `CASE` that
splits the batch by
+ * the predicate and merges the two branches back together, which costs more
than the call it
+ * protects. In the kernel it is a branch per row.
+ *
+ * Only that shape matches: the predicate must be the one the rule builds
from the UDF's guarded
+ * arguments. The kernel then runs no expression it would not run for the
UDF alone, apart from
+ * the guard's own `if`, `or` and `isnull`, and it runs Spark's code for all
of them. In a
+ * filter or join condition, and in the predicate of a conditional,
+ * `ReplaceNullWithFalseInPredicate` replaces the null with `false`, so that
form matches too.
+ */
+ def emitNullGuardDispatch(expr: If, inputs: Seq[Attribute], binding:
Boolean): Option[Expr] =
+ if (isNullGuard(expr)) tryJvmCodegenDispatch(expr, inputs,
binding).toOption else None
Review Comment:
Fixed in 25f9046eec. `emitNullGuardDispatch` now declines the guard when
converting it natively would hit a disabled expression: the UDF itself, or any
node of the predicate or the true branch. The native `If` then records the same
fallback reason as main. Checking the whole predicate rather than only `IsNull`
and `Or` matters when an argument is an expression. On main, `plusOne(a * 2)`
with `spark.comet.expression.Multiply.enabled=false` falls back because of the
predicate's copy of `a * 2`, and checking only the three guard nodes would have
kept it in Comet.
The new test disables `ScalaUDF`, `IsNull`, `Or` and `Multiply` one at a
time and expects the release's fallback reason with no `CometProject`. It fails
without the check, it fails on the `Multiply` case when only those three nodes
are checked, and it passes against main's serde.
@mbutrovich yes, I filed #6756 for the expressions inside a dispatched tree.
I left the UDF's own arguments unchecked here, as they are on main, because the
fix belongs in `emitJvmCodegenDispatch` for every dispatched expression rather
than in this guard.
--
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]
