sunchao commented on code in PR #5615:
URL: https://github.com/apache/datafusion-comet/pull/5615#discussion_r3917090591
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -336,7 +412,10 @@ private[comet] object NativeScanPlanDataInjector extends
PlanDataInjector {
op.getNativeScan.hasCommon &&
!op.getNativeScan.hasFilePartition
- override def getKey(op: Operator): Option[String] =
Some(sourceKey(op.getNativeScan.getCommon))
+ override def getKey(op: Operator): Option[String] = {
+ val common = op.getNativeScan.getCommon
+ Some(PlanDataInjector.cachedOrCompute(keyCache, common)(sourceKey(common)))
Review Comment:
Could we carry the existing driver-computed `sourceKey` in the serialized
`NativeScan` and read it directly here? [The driver already derives
it](https://github.com/apache/datafusion-comet/blob/843845940b4d8d3c4ab9efe9ec2a2851e2a33a04/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala#L365-L369).
Transporting that same key would preserve current matching semantics and the
injector interface while removing this LRU, repeated derivation after eviction,
and the dependency on sharing one protobuf instance to make lookup cheap.
It would also cover the native-shuffle path: [the writer builds its unified
plan from `spec.childNativeOp` and calls injection
directly](https://github.com/apache/datafusion-comet/blob/843845940b4d8d3c4ab9efe9ec2a2851e2a33a04/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala#L141-L155),
bypassing `parseBasePlan`. That child arrives through task dependency
deserialization, so a warm key-cache hit there still has to hash a fresh
protobuf and compare it structurally to the retained one. It avoids
stringification, but does not get the shared-instance fast path described above.
This is a proposed simplification, not a measured end-to-end alternative. It
needs the usual Java/Rust protobuf regeneration and a round-trip check
preserving key matching across query-context interning and scans with different
filters/projections. There is no need to change native injection or the contrib
SPI for this approach.
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -95,6 +95,56 @@ private[comet] trait PlanDataInjector {
* Registry and utilities for injecting per-partition planning data into
operator trees.
*/
private[comet] object PlanDataInjector extends Logging {
+ import java.nio.ByteBuffer
+ import java.util.{LinkedHashMap, Map => JMap}
+
+ private[comet] final val maxCachedBasePlans = 16
+
+ // Every task of a stage deserializes its own byte-identical copy of the
base plan, so
+ // without a cache an executor re-parses the same operator tree once per
task. Parsed
+ // Operators are immutable, so one instance is safely shared across
concurrent tasks.
+ // Keyed by content (ByteBuffer hashes/compares the bytes) since the arrays
are distinct.
+ //
+ // Entries are whole parsed plan trees, so the entry count is what bounds
executor memory:
+ // at most 16 recent stages' plans stay live, LRU-evicted as stages turn
over. A stage
+ // rerun that misses after eviction simply re-parses.
+ private val basePlanCache = java.util.Collections.synchronizedMap(
+ new LinkedHashMap[ByteBuffer, Operator](4, 0.75f, true) {
+ override def removeEldestEntry(eldest: JMap.Entry[ByteBuffer,
Operator]): Boolean = {
+ size() > maxCachedBasePlans
+ }
+ })
+
+ /**
+ * Look up `key`, computing and inserting the value on a miss. The
computation runs outside any
+ * lock so unrelated misses never serialize behind each other; when two
threads race the same
+ * cold key, the first insert wins and the loser adopts it, keeping the
cached value
+ * reference-shared (which the sourceKey memo's identity fast path relies
on).
+ */
+ private[comet] def cachedOrCompute[K, V](cache: JMap[K, V], key: K)(compute:
=> V): V = {
+ val cached = cache.get(key)
+ if (cached != null) {
+ cached
+ } else {
+ val computed = compute
+ cache.synchronized {
+ val winner = cache.get(key)
+ if (winner != null) {
+ winner
+ } else {
+ cache.put(key, computed)
+ computed
+ }
+ }
+ }
+ }
+
+ /**
+ * Parse a stage's base plan bytes, sharing the parsed tree across the
executor's tasks. Falls
+ * back to a plain parse on eviction, so a stage rerun is always correct.
+ */
+ def parseBasePlan(bytes: Array[Byte]): Operator =
+ cachedOrCompute(basePlanCache,
ByteBuffer.wrap(bytes))(Operator.parseFrom(bytes))
Review Comment:
Could we give the content key a stored hash, calculated before entering the
cache monitor and preferably once on the driver? `ByteBuffer.hashCode()` scans
every byte on every lookup. `synchronizedMap.get()` computes that hash while
holding the executor-wide lock, so even warm hits serialize work proportional
to plan size. The same cost applies to the new common-data cache. A miss in an
already populated cache can hash the same bytes again for the second lookup and
insertion.
I measured a warmed, single-entry cache-hit operation using a 45,697-byte
protobuf plan with 1,000 Long columns (required/data schemas, fields and
projection), with each worker holding a distinct equal byte array. At 8
threads, repeated measurements gave:
| Key implementation | Aggregate elapsed microseconds per successful lookup |
| --- | ---: |
| Current ByteBuffer key | 52-53 |
| Hash computed outside the lock | 5.7 |
| Previously computed hash carried with the bytes | 1.6 |
The alternatives still perform full content equality on hits and collisions.
These are short component measurements on JDK 17 with a shared 16-CPU host, not
individual task latency or whole-query speedups. The precomputed-hash case
excludes hash preparation because the proposal performs it once before task
execution. This demonstrates avoidable lookup overhead, without claiming that
the PR is slower overall than its uncached base.
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -328,6 +378,32 @@ private[comet] object IcebergPlanDataInjector extends
PlanDataInjector {
* Injector for NativeScan operators.
*/
private[comet] object NativeScanPlanDataInjector extends PlanDataInjector {
+ import java.nio.ByteBuffer
+ import java.util.{LinkedHashMap, Map => JMap}
+
+ private final val maxCacheEntries = 16
+
+ // Same rationale as IcebergPlanDataInjector's commonCache: the common bytes
are identical
+ // for every partition of a stage, and parsing them dominates inject() for
wide schemas.
+ private val commonCache = java.util.Collections.synchronizedMap(
+ new LinkedHashMap[ByteBuffer, OperatorOuterClass.NativeScanCommon](4,
0.75f, true) {
+ override def removeEldestEntry(
+ eldest: JMap.Entry[ByteBuffer,
OperatorOuterClass.NativeScanCommon]): Boolean = {
+ size() > maxCacheEntries
Review Comment:
Could the prepared scan data share the base plan's ownership/eviction unit?
The base cache holds 16 plans, but this cache and `keyCache` each hold only 16
scans. A single still-cached plan can therefore exceed both scan caches and
repeatedly evict everything needed by the next partition.
Using the exact cache/injector code in a component harness, traversing the
same distinct scans in the same order gave:
- One plan with 16 scans: the next pass reused 16/16 key strings and 16/16
parsed commons.
- One plan with 17 scans: the base plan was reused, but the next pass reused
0/17 keys and 0/17 commons.
- Nine plans with two distinct scans each: the next pass reused 9/9 base
plans, but 0/18 keys and 0/18 commons.
Thus schema-to-string key derivation and common parsing keep running even
while the relevant base plans are all cached. This is a conditional loss of the
intended reuse, not a demonstrated total regression versus the base. A prepared
entry owning the plan's keys and finalized common metadata would avoid
independent scan eviction. If preparation includes common data, its identity
must cover that finalized data or the execution, since resolved scalar-subquery
filters can differ for identical base-plan bytes. Please cover this scan-count
case in the performance 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]
