andygrove opened a new pull request, #6715:
URL: https://github.com/apache/datafusion-comet/pull/6715
## Which issue does this PR close?
Closes #6704.
## Rationale for this change
Spark's `HandleNullInputsForUDF` rule wraps a Scala UDF that has a primitive
parameter over a nullable input as `if (isnull(a) or ...) null else
f(knownnotnull(a), ...)`. The serde sent only the `ScalaUDF` to the codegen
dispatcher and ran the `If` natively, as a `CASE` that filters the batch for
each branch and merges the results. @mbutrovich measured on #6697 that this
guard is most of the gap between a dispatched UDF and a vectorized one.
## What changes are included in this PR?
- `CometIf` offers each `If` to a new `CometScalaUDF.emitNullGuardDispatch`
first. It matches only the shape the rule builds: the predicate has to equal
the `Or` of `IsNull` checks rebuilt from the UDF's `KnownNotNull` arguments at
primitive positions, and the true branch has to be `null`, or `false`, which
`ReplaceNullWithFalseInPredicate` puts there in filters, join conditions and
the predicates of conditionals. The whole `If` then goes to the dispatcher, and
the null check becomes a branch in the kernel's loop. The predicate only tests
the UDF's own arguments, so the kernel runs nothing it would not run for the
UDF alone apart from the guard's `if`, `or` and `isnull`, and those run as
Spark's own code.
- A nondeterministic argument does not match, because `semanticEquals` is
false for it, so that guard stays native.
- `emitJvmCodegenDispatch` is split into a private `tryJvmCodegenDispatch`
that returns the reason it cannot dispatch instead of recording it. When the
dispatcher refuses the whole guard, `CometIf` converts it natively as before
and the UDF is dispatched on its own. A refused guard therefore neither takes
the operator to Spark nor records an `if:` fallback reason for an expression
the query never wrote. The one case I found where the guard is refused and the
UDF alone is not is `spark.sql.codegen.maxFields`, because `canHandle` counts
the guarded column twice.
Some guards Spark produces still run natively, and could be follow-ups: the
outer guard of nested primitive UDFs, whose predicate the optimizer rewrites to
`isnull(a)`; `f(c) + 1`, where the optimizer pushes `+ 1` into the branch; and
guards users write around UDFs with boxed parameters.
I measured `SELECT max(f(c))` and `SELECT count(*) ... WHERE g(c)` over 4M
rows of a Parquet `bigint` column with a tenth of the rows null, with `f = (x:
Long) => x + 1` and `g = (x: Long) => x % 3 == 0`. The numbers are milliseconds
above `SELECT max(c)` in the same run, each the best of at least 10 iterations,
averaged over two interleaved runs of each build (before, after, after, before):
| Batch size | `max(f(c))` before | after | `WHERE g(c)` before | after |
| --- | --- | --- | --- | --- |
| 1024 | 93.5 | 72.5 | 85.5 | 72 |
| 8192 | 42 | 26.5 | 34 | 20.5 |
| 65536 | 32 | 18 | 25.5 | 10.5 |
That is Spark 4.1 with a release build on an Apple M3 Max. The benchmark is
not part of this PR; its source is below.
<details>
<summary>Benchmark source</summary>
```scala
package org.apache.spark.sql.benchmark
import org.apache.spark.benchmark.Benchmark
import org.apache.comet.{CometConf, ExtendedExplainInfo}
object UdfNullGuardBenchmark extends CometBenchmarkBase {
private val Rows = 4 * 1024 * 1024
private val cases = Seq(
"max(c), no function" -> "SELECT max(c) FROM parquetV1Table",
"max(c + 1), native" -> "SELECT max(c + 1) FROM parquetV1Table",
"max(prim_add_one(c)), dispatched" -> "SELECT max(prim_add_one(c)) FROM
parquetV1Table",
"count where prim_mult3(c), dispatched" ->
"SELECT count(*) FROM parquetV1Table WHERE prim_mult3(c)")
override def runCometBenchmark(mainArgs: Array[String]): Unit = {
spark.udf.register("prim_add_one", (x: Long) => x + 1)
spark.udf.register("prim_mult3", (x: Long) => x % 3 == 0)
withTempPath { dir =>
withTempTable(tbl) {
spark.range(Rows).createOrReplaceTempView(tbl)
prepareTable(dir, spark.sql(s"SELECT IF(id % 10 = 0, NULL, id *
7919) AS c FROM $tbl"))
cases.foreach { case (label, q) =>
withSQLConf(configs(8192): _*) {
val df = spark.sql(q)
val result = df.collect().head
val plan = df.queryExecution.executedPlan
val info = new ExtendedExplainInfo()
// scalastyle:off println
println(
s"$label: result=$result
native=${info.getNativeExpressions(plan)} " +
s"dispatched=${info.getCodegenDispatchExpressions(plan)} " +
s"fallback=${info.getFallbackReasons(plan)}")
// scalastyle:on println
}
}
Seq(1024, 8192, 65536).foreach { batchSize =>
val benchmark = new Benchmark(
s"null guard, batch size $batchSize",
Rows,
minNumIters = 10,
output = output)
cases.foreach { case (label, q) =>
benchmark.addCase(label) { _ =>
withSQLConf(configs(batchSize): _*) {
spark.sql(q).noop()
}
}
}
benchmark.run()
}
}
}
}
private def configs(batchSize: Int): Seq[(String, String)] = Seq(
CometConf.COMET_ENABLED.key -> "true",
CometConf.COMET_EXEC_ENABLED.key -> "true",
CometConf.COMET_BATCH_SIZE.key -> batchSize.toString,
CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true")
}
```
</details>
## How are these changes tested?
New tests in `CometCodegenSuite`:
- Single-argument, multi-argument (`or`), expression-argument and
aggregate-input guards run in the kernel, checked with
`checkSparkAnswerAndImpl` so that `if` and `isnull` must not be native.
- The `false` form, in a filter and in the predicate of an `IF` the query
wrote.
- A nondeterministic UDF that counts its calls is called only for non-null
rows, matching Spark.
- With `maxFields=2` the guard stays a native `if` around the dispatched UDF
and the projection stays in Comet. With the dispatcher off, the fallback reason
names the UDF and none is recorded for `if`.
- Shapes that do not match stay native: a predicate on another column, an
extra check, a non-null branch, and a nondeterministic argument.
Many existing dispatcher tests use primitive UDFs over nullable columns and
now take this path too.
I checked that the tests fail without the change, and against two simpler
designs. Dispatching the guard through the tagging `emitJvmCodegenDispatch`
with no native retry makes the projection fall back under `maxFields=2`.
Keeping the retry but recording the failed attempt adds an `if: ...=false`
reason with the dispatcher off. Dropping the `false` form, or comparing with
`==` instead of `semanticEquals`, also fails a test.
Local runs:
- `CometCodegenSuite` passes on Spark 3.4, 3.5 and 4.1.
- `CometCodegenFuzzSuite`, `CometCodegenHOFSuite`,
`CometCodegenSourceSuite`, `CometScalaUDFClassLoaderSuite`,
`CometUdfBridgeSuite`, `CometGenerateExecSuite`,
`CometNativePositionalRoundRobinSuite` and
`RevertNativeForTransitionHeavyStagesSuite` pass on 4.1.
- Semantic scalafix, scalastyle and spotless on 3.5 are clean.
I did not run the 4.0 and 4.2 profiles, `CometNativeUdfSuite`, which needs a
library that wasn't built locally, or Spark's SQL tests.
--
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]
