sunchao commented on code in PR #6876:
URL: https://github.com/apache/datafusion-comet/pull/6876#discussion_r4242437206
##########
spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala:
##########
@@ -144,9 +169,14 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
private def lookupOrCompile(
key: CometScalaUDFCodegen.CacheKey,
bytes: Array[Byte],
- specs: IndexedSeq[ArrowColumnSpec]): CometScalaUDFCodegen.CacheEntry = {
+ specs: IndexedSeq[ArrowColumnSpec],
+ partitionIndex: Int): CometScalaUDFCodegen.CacheEntry = {
assert(Thread.holdsLock(this), "lookupOrCompile must run under
this.synchronized")
- kernelCache.get(key) match {
+ // A deterministic kernel never reads the partition index, so one instance
under the planless
+ // key serves every plan in the task. Only a kernel with a
nondeterministic node is stored per
+ // plan.
+ val sharedKey = key.copy(planId = CometScalaUDFCodegen.NoPlan)
+ kernelCache.get(sharedKey).orElse(kernelCache.get(key)) match {
Review Comment:
[P2] Reuse the closure hash across these two lookups. For a nondeterministic
dispatched UDF with a large captured array, `get(sharedKey)` always misses
before `get(key)` succeeds. Each case-class hash calls `ByteBuffer.hashCode()`,
scanning the entire serialized closure again on every batch. The base performs
one scan. With a captured 1 MiB lookup array, the current dispatcher took 2.28
ms per 4,096-row batch versus 1.26 ms with the base dispatcher, despite
identical results within one unchanged partition. Precompute the closure/schema
hash once and reuse it for both keys while preserving plan isolation.
Evidence: Reproduced with freshly compiled head and base dispatchers on
Spark 4.1.3, Scala 2.13.17, and JDK 21. Define serializable `Lookup(val data:
Array[Byte]) extends (Int => Int)` with `apply(i) = data(Math.floorMod(i,
data.length)).toInt`. Serialize `ScalaUDF(new
Lookup(Array.fill[Byte](1048576)(7)), IntegerType, Seq(SparkPartitionID()))`
into the dispatcher’s argument-zero `VarBinaryVector`, then repeatedly evaluate
4,096 rows at partition 0, plan 10. After 150 warmups per implementation, nine
alternating samples of 200 batches produced medians of 1.25946982 ms for base
and 2.281142645 ms for head. The serialized expression was 1,051,614 bytes.
Both used the same current codegen and supporting classes. A separate
extracted-cache probe corroborated the extra hashing cost. Source and results:
`/tmp/pr6876-review-6bdf8232/DirectProbe.scala`, `direct.log`,
`CacheProbe.scala`, and `validation.json`. This is a dispatcher benchmark, not
an end-to-end query timing.
--
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]
