andygrove commented on code in PR #6715:
URL: https://github.com/apache/datafusion-comet/pull/6715#discussion_r4222146075
##########
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:
Applied in 02aa2bbe90, and the scaladoc now says the predicate checks each
argument once. `BooleanSimplification` drops repeats the same way from 3.4
through 4.2, and `Or` canonicalizes by flattening and sorting its operands, so
the balanced `Or` it builds `semanticEquals` the left-deep one rebuilt here.
Both sides compare canonical forms, so `addBoth(a + 1, 1 + a)` matches too.
The first test now covers `addBoth(a, a)`, `addThree(a, b, a)` and
`addBoth(a, 1L)`. Against the old match it fails on `addBoth(a, a)`, with `if`
and `isnull` native. `CometCodegenSuite` passes on 3.4, 3.5, 4.0, 4.1 and 4.2.
--
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]
