mbutrovich commented on code in PR #6715:
URL: https://github.com/apache/datafusion-comet/pull/6715#discussion_r4197956212
##########
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:
> Please preserve the expression-enable checks before fusing the guard.
> `spark.comet.expression.IsNull.enabled=false` is similarly ignored.
I can reproduce both at `02d248a` on Spark 4.1. With
`spark.comet.expression.ScalaUDF.enabled=false`, `SELECT plusOne(a) FROM t`
plans a `CometProject` that dispatches `if`, `isnull`, `knownnotnull` and
`plusone`, and records no fallback reason. Setting
`spark.comet.expression.IsNull.enabled=false` gives the same plan.
This might help decide how far the fix should go. The dispatcher already
ignores these configs for the expressions inside a UDF's arguments. With
`spark.comet.expression.Multiply.enabled=false`, `SELECT plusOneBoxed(a * 2)
FROM t` (a `java.lang.Long` parameter, so Spark adds no guard) stays a
`CometProject` with `multiply` dispatched in the UDF's kernel. That happens
because the dispatcher binds the whole tree and only calls
`exprToProtoInternal` on the attributes and the payload literal
([`branch-1.1`](https://github.com/apache/datafusion-comet/blob/e9efd9f764ee0a59b7898ff028d6985d4a7a28e1/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala#L107-L108)).
This PR doesn't change that behavior.
The nodes this PR adds to the kernel are different. Before this change,
`ScalaUDF`, `IsNull` and `Or` each went through the [enable check in
`exprToProtoInternal`](https://github.com/apache/datafusion-comet/blob/8c783aa88104616dcf0f7876b4f7a9ba71e2bf31/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala#L1012-L1018),
and a disabled one sent the projection to Spark. What do you think about
having `emitNullGuardDispatch` return `None` when `CometConf.isExprEnabled` is
false for any of those nodes in the guard? The existing native conversion would
then record the same fallback reason as the release does. A test with each
config disabled would cover it. Should we also open an issue for the configs of
expressions inside a dispatched argument tree, since that gap is already in the
release?
--
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]
