mbutrovich commented on PR #6697: URL: https://github.com/apache/datafusion-comet/pull/6697#issuecomment-6006552028
Thanks @andygrove. I wanted to know how much of the gap between a row-level UDF on the codegen dispatcher (`CometScalaUDFCodegen`) and a vectorized UDF comes from the dispatcher itself, because that part could be closed without users rewriting anything. I wrote a benchmark that splits the dispatcher's cost into layers and ran it against this PR at 7ab1aa8. > Above the no-UDF floor, the vectorized UDF costs about 1.5 to 2 ns per row and the dispatched UDF about 10 ns per row. Neither form slows down from `x + 1` to the mix, so at this size the cost is in calling the function rather than in its arithmetic. That per-row call is what a vectorized UDF saves. As I read the numbers below, most of that gap is in the dispatch layer around the call. C2 (HotSpot's optimizing JIT compiler) inlines the call itself into the kernel's loop. ### End to end The query is `SELECT max(f(c))` over 4M rows of a Parquet `bigint` column with a tenth of the rows null, as in `CometJvmUdfBenchmark`. Each number is milliseconds above `max(c)` with no function, best of at least 5 runs. | Form | batch 1024 | batch 8192 | batch 65536 | | --- | --- | --- | --- | | native `c + 1` | 2 | 0 | 0 | | vectorized UDF (`AddOneUdf` from this PR) | 26 | 7 | 6 | | vectorized UDF under `IF(c IS NULL, NULL, f(c))` | 45 | 20 | 19 | | dispatcher, `(x: Long) => x + 1` | 82 | 28 | 23 | | dispatcher, `(x: java.lang.Long) => ...` | 123 | 70 | 65 | | floor, `max(c)` (absolute ms) | 36 | 35 | 33 | At batch size 8192, the 21 ms between the dispatched `(x: Long)` UDF and the vectorized UDF splits three ways. 1. About 13 ms is Spark's null guard. Spark plans a UDF with a primitive parameter over a nullable column as `if (isnull(c)) null else prim_add_one(knownnotnull(c))`. The serde sends only the `ScalaUDF` to the dispatcher. The `If` runs natively as a DataFusion `CASE` ([`if_expr.rs`](https://github.com/apache/datafusion-comet/blob/7ab1aa87012a895f66575f70419fc840f3cd008e/native/spark-expr/src/conditional_funcs/if_expr.rs#L79-L81)), and DataFusion 55.1.0 evaluates a one-branch `CASE` with an `ELSE` by filtering the batch for each branch and merging the results ([`case.rs`](https://github.com/apache/datafusion/blob/7d3835c71f30cbd3c3ae4041732267f1f453097a/datafusion/physical-expr/src/expressions/case.rs#L998-L1009)). Putting the vectorized UDF under the same `IF` raises its cost from 7 ms to 20 ms. 2. About 4 ms is a fixed cost of roughly 8 microseconds per batch. The batch-size sweep shows it: the dispatched UDF minus the vectorized UDF under `IF` is 37, 8 and 4 ms at 4096, 512 and 64 batches. Calling `CometScalaUDFCodegen.evaluate` directly costs about 6 microseconds per batch more than calling its generated kernel. Each batch copies the 7,347-byte serialized closure ([L99](https://github.com/apache/datafusion-comet/blob/7ab1aa87012a895f66575f70419fc840f3cd008e/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala#L99)) and hashes it into the cache key ([L118](https://github.com/apache/datafusion-comet/blob/7ab1aa87012a895f66575f70419fc840f3cd008e/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala#L118)), which the `perf-cache-key` TODO already describes ([L311-L317](https://github.com/apache/datafusion-comet/blob/7ab1aa87012a895f66575f70419fc840f3cd008e/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen. scala#L311-L317)). I haven't found where the other 2 microseconds per batch go. 3. About 4 ms, under 1 ns per row, is the generated code. ### JVM layers alone These cases call each layer directly on one in-memory 8192-row batch, with no Spark query and no native code, so they leave out the null guard. Numbers are ns per row. | Layer | 10% nulls | no nulls | | --- | --- | --- | | vectorized UDF (`AddOneUdf`) | 1.2 | 1.2 | | hand-written row loop with an `isNull` branch and `BigIntVector.set` | 1.2 | 1.7 | | generated kernel, `(x: Long)` | 1.9 | 2.0 | | `CometScalaUDFCodegen.evaluate`, `(x: Long)` | 2.6 | 2.7 | | generated kernel, `(x: java.lang.Long)` | 15.6 | 19.2 | | `CometScalaUDFCodegen.evaluate`, `(x: java.lang.Long)` | 15.8 | 20.1 | The first four rows take 5 to 11 ms per iteration and the harness reports whole milliseconds, so differences of a few tenths of a ns per row between them are within its resolution. With `-XX:CompileCommand=option,*SpecificCometBatchKernel::process,PrintInlining`, C2 reports every call in the `(x: Long)` kernel's loop as `inline (hot)`. That includes the user's lambda through `JFunction1$mcJJ$sp::apply`, `Long.valueOf`, and `BigIntVector.set`. The kernel runs at about the speed of the hand-written row loop, which suggests C2 also removes the boxing. I couldn't confirm that from the assembly because this JDK has no hsdis disassembler. What remains per row is setting one validity bit at a time (`BitVectorHelper.setBit`, with an `ArrowBuf` bounds check), where the vectorized UDF copies the validity buffer once per batch. Neither JVM form looks SIMD-vectorized. Native `c + 1` costs nothing measurable above the scan, while the vectorized UDF costs 1.2 ns per row on its own. The `java.lang.Long` parameter is the slow case. Spark's `ScalaUDF` converts a boxed parameter through its input encoder ([L156-L175](https://github.com/apache/spark/blob/77bbf77e86ad48f58b5dfbc6ac882b3e70cf1989/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala#L156-L175), [L1139-L1147](https://github.com/apache/spark/blob/77bbf77e86ad48f58b5dfbc6ac882b3e70cf1989/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala#L1139-L1147)). C2 leaves the deserializer's `Projection::apply` as a virtual call, so each row runs a projection into a `GenericInternalRow`. The kernel uses Spark's own `ScalaUDF.doGenCode`, so Spark's whole-stage codegen generates the same per-row conversion. ### Questions These concern the dispatcher, so they would be separate PRs, but they bear on what this PR's API is for. - Should the serde send Spark's null guard to the dispatcher together with the `ScalaUDF`, so the null check is a branch in the kernel instead of a `CASE` over the batch? - Should we take one of the options in the `perf-cache-key` TODO, such as a hash computed on the driver, so that each batch doesn't rehash the closure? - What do you think about converting boxed primitive parameters without going through the encoder? If those changes close most of the gap for simple functions, rewriting such a function in vectorized form gains little. The rewrite @comphead suggested an agent skill for in #6694 would then be one the dispatcher can do itself. The vectorized API would be for what codegen can't derive from a row-level function: setup done once per batch, a different value representation (such as working on UTF-8 bytes instead of `String`), or one call per batch into a batch-oriented library. What do you think about framing the user guide around those cases? ### `CometJvmUdfBenchmark` fails at 7ab1aa8 On the default Spark 4.1 profile, it throws `add_one: Comet, codegen dispatch ran ShuffleQueryStage in Spark` although every operator in the plan is a Comet operator. `stripAQEPlan` removes only the top-level adaptive node ([L123-L124](https://github.com/apache/datafusion-comet/blob/7ab1aa87012a895f66575f70419fc840f3cd008e/spark/src/test/scala/org/apache/spark/sql/benchmark/CometJvmUdfBenchmark.scala#L123-L124)). `findFirstNonCometOperator` then reports the `ShuffleQueryStage` (a leaf node that wraps an already-executed part of the plan under adaptive query execution) as a non-Comet operator ([`CometPlanChecker.scala`](https://github.com/apache/datafusion-comet/blob/7ab1aa87012a895f66575f70419fc840f3cd008e/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanChecker.scala#L41-L59)). Should the check look inside each query stage's plan, as below? ```scala val plan = stripAQEPlan(df.queryExecution.executedPlan) // A query stage is a leaf, so check the plan inside each one separately. val stagePlans = collect(plan) { case s: QueryStageExec => s.plan } (plan +: stagePlans) .flatMap(findFirstNonCometOperator(_, classOf[QueryStageExec])) .headOption .foreach { op => throw new IllegalStateException(s"$fn: ${form.label} ran ${op.nodeName} in Spark:\n$plan") } ``` ### Setup The machine was an Apple M5 Max on AC power, with JDK 17.0.20.1 and Spark 4.1.3 (the default profile). Comet was built with `make release` and the benchmark run with `make benchmark-org.apache.spark.sql.benchmark.CometUdfDispatchCostBenchmark`. The benchmark source is below. It reuses `AddOneUdf` from this PR's `CometJvmUdfBenchmarkUdfs`. <details> <summary><code>CometUdfDispatchCostBenchmark.scala</code></summary> ```scala /* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ package org.apache.spark.sql.benchmark import org.apache.arrow.vector.{BigIntVector, ValueVector, VarBinaryVector} import org.apache.spark.SparkEnv import org.apache.spark.benchmark.Benchmark import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSeq, BindReferences, Expression, ScalaUDF} import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.spark.sql.types.LongType import org.apache.comet.{CometArrowAllocator, CometConf} import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec import org.apache.comet.udf.CometJvmUDF import org.apache.comet.udf.codegen.CometScalaUDFCodegen /** * Splits the cost of a row-level Scala UDF on the codegen dispatcher into its layers, to find * where it loses to a vectorized UDF (`CometJvmUdfBenchmark`). * * Part 1 runs `SELECT max(f(c))` end to end at several batch sizes. A cost paid once per batch * shrinks as batches grow, and a cost paid per row does not. It also separates the `IF(c IS NULL, * NULL, f(c))` wrapper Spark puts around a UDF with a primitive parameter, by running a * boxed-parameter UDF (no wrapper) and the vectorized UDF under the same wrapper. * * Part 2 calls each JVM layer directly on one in-memory 8192-row batch, with no Spark query and * no native code: the vectorized UDF, a hand-written row loop, the generated kernel alone, and * the full `CometScalaUDFCodegen.evaluate` (kernel plus per-batch cache lookup). * * To run: * {{{ * make benchmark-org.apache.spark.sql.benchmark.CometUdfDispatchCostBenchmark * }}} */ object CometUdfDispatchCostBenchmark extends CometBenchmarkBase { import CometJvmUdfBenchmarkUdfs._ private val Rows = 4 * 1024 * 1024 private val BatchRows = 8192 private val BatchSizes = Seq(1024, 8192, 65536) override def runCometBenchmark(mainArgs: Array[String]): Unit = { spark.udf.register("prim_add_one", (x: Long) => x + 1) spark.udf.register( "boxed_add_one", (x: java.lang.Long) => if (x == null) null else java.lang.Long.valueOf(x + 1)) CometJvmUDF.register(spark, "jvm_add_one", classOf[AddOneUdf], Seq(LongType), LongType) withTempPath { dir => withTempTable(tbl, "parquetV1Table") { spark.range(Rows).createOrReplaceTempView(tbl) prepareTable(dir, spark.sql(s"SELECT IF(id % 10 = 0, NULL, id * 7919) AS c FROM $tbl")) printPlans() BatchSizes.foreach(endToEnd) Seq(true, false).foreach(jvmLayers) } } } private val endToEndCases = Seq( "dispatch, primitive param" -> ("prim_add_one(c)", true), "dispatch, boxed param" -> ("boxed_add_one(c)", true), "vectorized UDF" -> ("jvm_add_one(c)", false), "vectorized UDF under IF(c IS NULL)" -> ("IF(c IS NULL, NULL, jvm_add_one(c))", false), "native c + 1" -> ("c + 1", false), "no function, max(c)" -> ("c", false)) private def configs(dispatch: Boolean, 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 -> dispatch.toString) private def query(f: String): String = s"SELECT max($f) FROM parquetV1Table" /** Shows the expression each case hands to Comet and checks that Comet runs every operator. */ private def printPlans(): Unit = endToEndCases.foreach { case (label, (f, dispatch)) => withSQLConf(configs(dispatch, BatchRows): _*) { val df = spark.sql(query(f)) val result = df.collect().head val plan = stripAQEPlan(df.queryExecution.executedPlan) // A query stage is a leaf, so check the plan inside each one separately. val stagePlans = collect(plan) { case s: QueryStageExec => s.plan } (plan +: stagePlans) .flatMap(findFirstNonCometOperator(_, classOf[QueryStageExec])) .headOption .foreach { op => throw new IllegalStateException(s"$label ran ${op.nodeName} in Spark:\n$plan") } val aggInput = df.queryExecution.optimizedPlan.expressions.headOption // scalastyle:off println println(s"$label: result=$result optimized=${aggInput.getOrElse("")}") // scalastyle:on println } } private def endToEnd(batchSize: Int): Unit = { val benchmark = new Benchmark(s"max(f(c)), batch size $batchSize", Rows, minNumIters = 5, output = output) endToEndCases.foreach { case (label, (f, dispatch)) => benchmark.addCase(label) { _ => withSQLConf(configs(dispatch, batchSize): _*) { spark.sql(query(f)).noop() } } } benchmark.run() } /** The bound UDF subtree the dispatcher serde would ship, and its closure-serialized bytes. */ private def boundUdf(f: String): (Expression, Array[Byte]) = { val plan = spark.sql(s"SELECT $f FROM parquetV1Table").queryExecution.optimizedPlan val udf = plan.expressions.flatMap(_.collect { case u: ScalaUDF => u }).head val attrs = udf.collect { case a: AttributeReference => a }.distinct val bound = BindReferences.bindReference(udf: Expression, AttributeSeq(attrs)) val buffer = SparkEnv.get.closureSerializer.newInstance().serialize(bound) val bytes = new Array[Byte](buffer.remaining()) buffer.get(bytes) (bound, bytes) } private def jvmLayers(withNulls: Boolean): Unit = { val in = new BigIntVector("c", CometArrowAllocator) in.allocateNew(BatchRows) var i = 0 while (i < BatchRows) { if (withNulls && i % 10 == 0) in.setNull(i) else in.set(i, i * 7919L) i += 1 } in.setValueCount(BatchRows) val inputs: Array[ValueVector] = Array(in) val batchesPerIter = Rows / BatchRows val shape = if (withNulls) "10% nulls" else "no nulls" val benchmark = new Benchmark( s"JVM layers, one $BatchRows-row batch ($shape)", Rows, minNumIters = 5, output = output) def addBatchCase(label: String)(evaluate: () => ValueVector): Unit = benchmark.addCase(label) { _ => var b = 0 while (b < batchesPerIter) { evaluate().close() b += 1 } } val vectorized = new AddOneUdf addBatchCase("vectorized UDF") { () => vectorized.evaluate(inputs, BatchRows) } addBatchCase("row loop, null branch, no boxing") { () => val out = new BigIntVector("result", CometArrowAllocator) out.allocateNew(BatchRows) var r = 0 while (r < BatchRows) { if (in.isNull(r)) out.setNull(r) else out.set(r, in.get(r) + 1) r += 1 } out.setValueCount(BatchRows) out } val toClose = Seq("prim_add_one(c)" -> "primitive param", "boxed_add_one(c)" -> "boxed param") .map { case (f, param) => val (bound, bytes) = boundUdf(f) val specs = IndexedSeq(ArrowColumnSpec(classOf[BigIntVector], nullable = true)) // scalastyle:off println println(s"$param: bound=$bound serializedBytes=${bytes.length}") println(CometBatchKernelCodegen.generateSource(bound, specs).body) // scalastyle:on println val kernel = CometBatchKernelCodegen.compile(bound, specs).newInstance() kernel.init(0) val field = CometBatchKernelCodegen.toFfiArrowField("r", bound.dataType, bound.nullable) addBatchCase(s"generated kernel only, $param") { () => val out = CometBatchKernelCodegen.allocateOutput(field, BatchRows, 0) kernel.process(inputs, out, BatchRows) out.setValueCount(BatchRows) out } val exprVec = new VarBinaryVector("expr", CometArrowAllocator) exprVec.allocateNew() exprVec.setSafe(0, bytes) exprVec.setValueCount(1) val dispatcher = new CometScalaUDFCodegen val dispatchInputs: Array[ValueVector] = Array(exprVec, in) addBatchCase(s"dispatcher evaluate, $param") { () => dispatcher.evaluate(dispatchInputs, BatchRows) } exprVec } try benchmark.run() finally { toClose.foreach(_.close()) in.close() } } } ``` </details> -- 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]
