sunchao commented on code in PR #6715:
URL: https://github.com/apache/datafusion-comet/pull/6715#discussion_r4196427044
##########
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:
[P2] Please preserve the expression-enable checks before fusing the guard.
Register `plusOne` as `(x: Long) => x + 1`, query a nullable Parquet `BIGINT`
column, and set `spark.comet.expression.ScalaUDF.enabled=false` while leaving
codegen dispatch enabled. Previously, serializing the UDF through
`exprToProtoInternal` rejected it and the projection fell back to Spark. This
direct dispatch now skips that check and retains a `CometProject` executing the
disabled UDF. `spark.comet.expression.IsNull.enabled=false` is similarly
ignored. This breaks the documented per-expression fallback control used to
avoid problematic execution paths. Decline fusion when an expression-enable
gate that native conversion would visit is disabled, then let normal conversion
report the fallback.
Evidence: Reproduced on Spark 4.1.3/JDK 17 with exact-head JVM classes and
native artifact 11419676187, whose workflow head SHA matches this review.
Disposable `PR6715ConfigSuite` queried Parquet rows `(1), (NULL), (-3)` under
each disabled setting. Both tests failed at head: no fallback reason,
`CometProject`, and four dispatched expressions. Both passed after recompiling
only `CometScalaUDF.scala` and `conditional.scala` from parent
9dc8c3ca89962506078831f7064a25c75da83883. A separate serde probe confirmed
`FALLBACK` at the parent versus `JVM_SCALAR_UDF` at head. Sources and logs are
under `/tmp/pr6715-validation/`.
--
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]
