mbutrovich commented on code in PR #6715:
URL: https://github.com/apache/datafusion-comet/pull/6715#discussion_r4221802584
##########
spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala:
##########
@@ -58,6 +58,63 @@ 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, when an expression in it is
disabled, or when
+ * 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) && isEnabled(expr)) {
+ tryJvmCodegenDispatch(expr, inputs, binding).toOption
+ } else {
+ None
+ }
+
+ /**
+ * Whether converting `guard` natively would pass every
`spark.comet.expression.<name>.enabled`
+ * check it makes: on the UDF, and on each node of the predicate and the
true branch, the
+ * guarded arguments included. The dispatcher makes none of these checks, so
a guard with a
+ * disabled expression stays native, where the conversion records the
fallback reason. The UDF's
+ * arguments are not checked, because the dispatcher takes them unchecked
with or without the
+ * guard.
+ */
+ private def isEnabled(guard: If): Boolean = {
+ def disabled(e: Expression): Boolean =
+ QueryPlanSerde.exprSerdeMap.get(e.getClass).exists { serde =>
+ !CometConf.isExprEnabled(
+
serde.asInstanceOf[CometExpressionSerde[Expression]].getExprConfigName(e))
+ }
+ !disabled(guard.falseValue) && !guard.predicate.exists(disabled) &&
+ !guard.trueValue.exists(disabled)
+ }
+
+ private def isNullGuard(expr: If): Boolean = expr match {
+ case If(predicate, Literal(null, _) | Literal.FalseLiteral, udf: ScalaUDF)
=>
+ val guarded = udf.inputPrimitives.zip(udf.children).collect {
+ case (true, KnownNotNull(arg)) => arg
+ }
+ // `semanticEquals` is false for a nondeterministic argument, so that
guard stays native.
+ guarded
+ .map(IsNull(_))
+ .reduceLeftOption[Expression](Or)
+ .exists(_.semanticEquals(predicate))
Review Comment:
What do you think about also matching a UDF that gets the same column at
more than one primitive position?
[`HandleNullInputsForUDF`](https://github.com/apache/spark/blob/e221b56be7b6d9e48e107fc4d1cf0c15f02700f8/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala#L3625-L3660)
adds one `isnull` per position, and
[`BooleanSimplification`](https://github.com/apache/spark/blob/e221b56be7b6d9e48e107fc4d1cf0c15f02700f8/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala#L385)
then drops the repeats. With `addBoth = (a: Long, b: Long) => a + b` on Spark
4.1, `SELECT addBoth(a, a) FROM t` optimizes to `if (isnull(a)) null else
addBoth(knownnotnull(a), knownnotnull(a))`, and `addThree(a, b, a)` optimizes
to `if ((isnull(a) OR isnull(b))) null else addThree(...)`. The predicate
rebuilt here keeps every repeat, so neither one matches, and extended explain
lists `if` and `isnull` as native for both.
Removing repeated arguments before building the predicate would cover these.
[`ExpressionSet.add`](https://github.com/apache/spark/blob/e221b56be7b6d9e48e107fc4d1cf0c15f02700f8/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExpressionSet.scala#L97-L104)
keeps the first of each group of semantically equal deterministic expressions
and keeps every nondeterministic one, so a nondeterministic argument still
fails `semanticEquals` and its guard stays native. This also needs
`ExpressionSet` in the import on line 25.
Could the first test also cover `addBoth(a, a)` and `addThree(a, b, a)`? A
case where one primitive argument isn't nullable, such as `addBoth(a, 1L)`,
would also check that an argument without `knownnotnull` gets no `isnull` in
the predicate.
```suggestion
// The optimizer drops a repeated `isnull`, as in `f(a, a)`, so check
each argument once.
// `semanticEquals` is false for a nondeterministic argument, so that
guard stays native.
ExpressionSet(guarded).toSeq
.map(IsNull(_))
.reduceLeftOption[Expression](Or)
.exists(_.semanticEquals(predicate))
```
--
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]
