perf: cache parsed plan data across a stage tasks - #5615
Conversation
Every task of a stage deserializes byte-identical plan bytes, yet each one parsed the full operator tree, re-derived the scan source key by stringifying its schema and filter lists, and re-parsed the scan common message before injecting its partition data. Three bounded per-executor caches now share that work: the parsed base plan keyed on content, the parsed NativeScanCommon, and a source-key memo that hits protobuf reference-identity fast path once the base plan is shared. Injection itself stays per task, since partition data genuinely differs, and the injected tree is never cached. Per-task overhead drops from roughly 300us to 60us on a 100-column scan plan and about 3x on a 1000-column plan. Cache misses compute outside any lock so unrelated stages never serialize behind one parse, and racing threads on a cold key adopt a single instance so reference sharing holds.
2ed7ece to
8438459
Compare
sunchao
left a comment
There was a problem hiding this comment.
Reviewed 843845940b4d8d3c4ab9efe9ec2a2851e2a33a04 against 8729f6e6adf7091e18a48670e790d4ba8fd41e51. No actionable P1/P2 findings. The caches use full content equality, retain immutable parsed messages, and leave partition-file injection and executable task state separate. Concurrent cold misses may duplicate computation, but the synchronized second lookup adopts an already-retained value.
For the performance claim, could you add a matched BASE/HEAD microbenchmark with 1/8/32 concurrent task threads, cold and warm caches, small and wide plans, and more than 16 interleaved plans? Please include throughput, tail latency, allocations/retained heap, and equal injected results with distinct partition file lists. Byte hashing/equality still runs inside synchronized lookups, and the 16-entry limit is not a byte limit; the reported warm-loop timing does not establish contention or churn behavior. I have not measured a regression.
This was a source review, including the added tests; I did not execute tests or benchmarks. The three current-head workflows require action and no head check runs are available, so CI is not independently validated.
sunchao
left a comment
There was a problem hiding this comment.
Reviewed 843845940b4d8d3c4ab9efe9ec2a2851e2a33a04 against 8729f6e6adf7091e18a48670e790d4ba8fd41e51. Sharing immutable plan metadata while keeping partition files and executable state task-local is a sound boundary. cachedOrCompute is a small, useful helper, and preserving the existing injector SPI keeps this change contained. I did not reproduce a new data-correctness failure.
Following up on the earlier benchmark discussion, the inline comments add concrete evidence for two avoidable costs: full byte hashing under the shared cache monitor, and scan-cache churn while all relevant base plans remain cached. They also suggest transporting the scan key already computed on the driver, which could remove the third cache.
For a broader simplification, the natural ownership unit is an immutable prepared plan for one execution, including its scan bindings and finalized common metadata. That would let related objects share an eviction/lifetime boundary. It must include finalized commonByKey in its identity, or use an execution-specific identity: scalar-subquery filters are added after initial planning, so equal base-plan bytes do not guarantee equal final scan metadata. Simply reusing the common object embedded in the base plan would lose that distinction.
An explicit broadcast of a serializable holder with lazily prepared executor state is another candidate to benchmark. A plain lazy field on the task-deserialized RDD would not provide the same sharing. The broadcast alternative adds setup and cleanup responsibilities, so it is not automatically the cheapest design. Moving all injection native would expand the SPI/JNI scope considerably. I would evaluate stored content hashes and a transported scan key before introducing a broader preparation framework.
Please include total scan count as well as plan count in the matched base/head benchmark: one retained plan with 17 distinct scans, and nine retained plans with two distinct scans each, both churn the new scan caches. Warm/cold cases, 1/8/32 task threads, small/wide plans, native shuffle, allocation/retained heap, and equal injected results with different partition files would make the performance claim easier to assess. The cache limit is an entry count, not a byte budget, and the new caches retain both serialized arrays and parsed object graphs. There is also a small coverage correction in the memo test: compare getKey directly with sourceKey(common), since the current “fresh” call uses an equal protobuf and reads the same memo entry.
Validation: five independent review scopes, plus a fresh review of the unchanged diff. All 14 PlanDataInjectorSuite tests passed again in a component harness using protobuf 3.25.5 schemas generated from this head and the extracted cache/injection implementations. The harness substitutes the built-in injector registry and an exception class, omits the Spark-dependent plan-data discovery method and logging, and does not run Spark or native execution. Additional component controls passed for malformed-input recovery and 256 tasks across eight threads sharing one base plan while injecting four runtime-common variants and 256 distinct file paths, including serialization/readback. I separately reproduced the scan-cache reuse counts and repeated the lookup benchmark. The lookup timings isolate cache access and do not establish a whole-query speedup or a total regression against the uncached base. No full project build or end-to-end base/head task benchmark was run locally.
Current CI: 64 successful checks and 9 skipped, with no failed or pending checks in the snapshot. The PR benchmark job is skipped. This updates the earlier review's CI snapshot.
| * 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)) |
There was a problem hiding this 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.
There was a problem hiding this comment.
Done in 36fa94d. The base plan cache now keys on a PlanKey that stores its hash, computed once per task before the monitor; equals is identity then Arrays.equals. Driver transport was not practical for this one since the plan bytes are the task binary itself, so this is your measured middle option. The other two caches are gone entirely, see the main comment.
| new LinkedHashMap[ByteBuffer, OperatorOuterClass.NativeScanCommon](4, 0.75f, true) { | ||
| override def removeEldestEntry( | ||
| eldest: JMap.Entry[ByteBuffer, OperatorOuterClass.NativeScanCommon]): Boolean = { | ||
| size() > maxCacheEntries |
There was a problem hiding this 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.
There was a problem hiding this comment.
Done. The cache value is now a holder with the parsed plan plus its prepared per scan commons, so everything for a plan lives and dies with its single entry. Reran your churn shapes: 1 plan with 17 scans goes from 0/17 reused to 17/17, nine plans with two scans each from 0/18 to 18/18, and the warm pass drops from 14.7ms to 1.0ms. Numbers in the main comment.
| 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))) |
There was a problem hiding this comment.
Could we carry the existing driver-computed sourceKey in the serialized NativeScan and read it directly here? The driver already derives it. 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, 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.
There was a problem hiding this comment.
This was the right call, thanks. source_key now rides in the NativeScan proto, derived once on the driver, and the keyCache is deleted. Reading the transported key measures 0.07us against 0.4 to 0.55ms per derivation under churn. It reaches the shuffle writer too via childNativeOp, and the executor keeps a derivation fallback only for plans built without the field. The scalar subquery caveat is handled by pinning the finalized bytes on each prepared entry and honoring hits only on byte equality.
|
Sorry for back and forth @dwsmith1983 . I just added a few more instructions to my Comet PR review skill especially for |
sunchao
left a comment
There was a problem hiding this comment.
Rechecked 4d103e4c against 10537e14. The PR-only patch is unchanged, and the head update exactly matches the base update. I checked the cache/injection integration and found no additional P1/P2.
The existing hashing, cache-churn, key-transport feedback and benchmark requests remain applicable. This base sync does not address them. The existing approval is unchanged. No tests were rerun, and current-head workflows still await approval.
…ase plan entry The executor derived each native scan's lookup key by stringifying its schema and filter lists, memoized in an LRU probed under the map monitor, and parsed scan commons through a second per-scan LRU that a single plan with 17 scans churned to zero reuse. The driver already derives the key once in CometNativeScanExec, so carry it inside the NativeScan proto and read it back on every injection path, including the native shuffle writer's, which previously missed the fast paths entirely. Prepared commons now live inside the base plan's own cache entry (scoped to the shuffleId on the shuffle-write path), so a plan and its scans form one eviction unit. Entries pin the finalized common bytes because scalar-subquery data filters resolve per execution: equal base plan bytes do not guarantee equal finalized commons, and a stale entry is replaced rather than served. The base plan cache keys on a stored hash computed once per task outside the monitor instead of a raw ByteBuffer that rescanned the bytes on every probe.
|
Took your preferred direction and it worked out well, thanks for the concrete numbers, they made the case obvious. Pushed in 36fa94d:
Benchmark with your requested matrix (component harness against the exact cache/injector code at three commits: main uncached, the previous design, this one; Apple M5, JDK 17, one forked JVM per cell, reuse counts observed by reference identity): Warm single-entry lookup, 42KB plan, us per lookup:
Your churn shapes, warm pass, prepared commons reused:
Warm pass time on the second shape drops 14.7ms to 1.0ms and allocation 61MB to 0.8MB per pass. Transported key read is 0.07us vs 0.4 to 0.55ms per derivation under churn. Shuffle steady state improves 1654 to 1354us per call with 1 prepare instead of 64 on the uncached base. Injected outputs are byte identical across all three commits for equal inputs and differ only in partition fields across different file sets. Retained heap per entry is unchanged between the two designs (about 241KB parsed plan, 567KB with a prepared wide common), the difference is eviction shape, not weight. One honest note: on a cold start with 8 threads racing the very first shuffle calls, concurrent callers can each prepare the same common once before the store converges (8 then 1 thereafter). Transient duplicate work only, the hot path stays lock free on hits. And yes, the new review format works well from this side: the measured tables made it unambiguous what to fix and in what order. |
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed 36fa94d3ae84210da6ff0d3f9df5a31334f4b641. The earlier hashing, scan-cache churn and key-transport concerns are addressed. I found one new P2 concerning the lifetime of the shuffle cache, detailed inline.
The 25 isolated component tests passed. The retention finding is supported by source tracing and a cache-level probe, not a full SparkContext restart or heap measurement. CI is still running.
| // never pass through parseBasePlan. Prepared commons for that path are scoped to the | ||
| // shuffleId instead: one shuffle stage's scans still share a single eviction unit. | ||
| private val shufflePreparedCommons = java.util.Collections.synchronizedMap( | ||
| new LinkedHashMap[Integer, ConcurrentHashMap[String, PreparedCommon]](4, 0.75f, true) { |
There was a problem hiding this comment.
[P2] Scope prepared shuffle data to the shuffle-manager lifetime
For local or embedded callers that stop and recreate SparkContext in the same JVM and Comet classloader, each context restarts shuffle IDs at zero. This singleton survives, and neither Comet shuffle manager clears it in stop(). When successive contexts perform native shuffles with new scan keys, those keys accumulate in the same inner map if the reused IDs stay within the 16-entry limit. An exact-source component probe retained 128 prepared commons under ID 0 in one outer entry. The finalized-byte guard prevents stale reads but does not remove old keys. Could this store be owned by the manager lifetime or explicitly cleared on stop, with a recreated-context regression test? Otherwise successive contexts can keep retaining more commons despite the outer LRU bound.
There was a problem hiding this comment.
Fixed in f24e4fc, at both boundaries you named. unregisterShuffle now releases that shuffle's prepared commons in both managers, which is the precise lifetime (verified the ContextCleaner path in the 3.5.9 bytecode: doCleanupShuffle to BlockManagerStorageEndpoint RemoveShuffle to ShuffleManager.unregisterShuffle, and neither Comet manager delegates that call away), and stop() clears the whole store as the safety net for recreated contexts. Your reused-id scenario is the regression: two manager instances with a stop between them, shuffle 0 holds exactly the second context's keys instead of four. There is also a real SparkContext suite that runs a native shuffle, proves the cleaner path releases the entry, stops the context, recreates a session, and asserts the store holds only the new context's keys. The base plan cache is left alone deliberately since it is keyed by plan bytes rather than a per-context counter, so a recreated context either hits or evicts through the existing bound. One residual note: a straggling map task calling injection after unregister would re-insert an empty inner map for that id, bounded by the outer LRU and cleared on stop, so I did not add a tombstone.
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed 1eecb29e4b58f281ea36dc66523283201c92ba80 against ef62b463. [P2] The existing cache-lifetime finding remains. The singleton store and both shuffle managers' cleanup paths are unchanged. Successive local or embedded SparkContexts sharing one JVM and Comet classloader can reuse shuffle IDs while retaining new scan keys in the same inner maps. The finalized-byte check protects values but does not release those old keys. Could you scope the store to its manager/environment lifetime or clear it at that lifecycle boundary, with a recreated-context regression?
This re-review reused the source-verified component evidence. No new Spark restart, native or end-to-end execution was added. The three current-head workflows still require authorization.
…the manager stops The shuffle-scoped prepared-commons store is a JVM singleton keyed by shuffleId, and shuffle ids restart at zero for every SparkContext, so a local or embedded caller that stops and recreates its context kept stacking new scan keys under ids the previous context had used. Both CometShuffleManager and CometCelebornShuffleManager now drop a shuffle's entry in unregisterShuffle (reached from Spark's ContextCleaner via BlockManagerStorageEndpoint) and clear the store in stop(), so a recreated context starts empty and long-lived contexts release each shuffle's data with the shuffle.
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed f24e4fcc against ef62b463. The existing cache-lifetime P2 is addressed for the reported normal context-restart case: both managers now release per-shuffle entries and clear the store on stop. The earlier hashing, cache-ownership and key-transport fixes remain intact. I found no remaining P1/P2.
All 32 isolated component tests passed, and I independently reran the seven lifecycle cases. These use extracted cleanup methods and stand-ins, not a real SparkContext or native shuffle. The new real-context tests and their Linux/macOS registration were inspected but not run. Late injection can repopulate a nonempty map, so this does not establish permanent emptiness after shutdown.
The three current-head workflows remain action_required, with no head or merge check results.
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed ea876124 against 55ae4f20. The feature patch is unchanged through this merge, and the manager cleanup paths and lifecycle-suite registration remain intact. I found no new P1/P2. The existing approval stands.
This was a source and integration-boundary recheck. No runtime tests or benchmarks were rerun, and the prior 32 component passes remain historical evidence. At 04:50 UTC, CI had 42 successful checks, 22 running and 7 skipped, with no failures reported.
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Re-reviewed 4ac9d46a against 2da32915, including the merge since ea876124. The feature changes are unchanged, and the head update matches the base update. I found no new P1/P2. The existing approval can stand.
I rechecked the lifecycle against the maintained Spark 3.5 and 4.0 sources. Shuffle IDs are local to each SparkContext, normal shuffle removal reaches unregisterShuffle, and environment shutdown calls the shuffle manager's stop. Both Comet managers retain the per-shuffle release and stop-time cleanup for the existing lifetime finding. Late injection can still repopulate a nonempty map. This does not establish permanent emptiness after shutdown.
The finalized-common-byte check remains intact. It also covers the newly merged has_data_filters field, so changing that flag under an existing key causes re-preparation. Shared parsed metadata remains separate from each partition's file data and executable task state. The transported source_key and derivation fallback remain present. The cache feature's expression type, null, overflow, ANSI-mode and fallback behavior is unchanged.
Validation and CI
This round used source inspection, commit/tree checks, identical normalized feature patches and matching stable patch IDs. No runtime tests or benchmarks were rerun. The prior 32 component passes and seven lifecycle cases are historical evidence, not fresh validation of this head. The lifecycle suite and its Linux/macOS registration remain present.
At 2026-09-04T18:36:38Z, the three current-head workflows were action_required, with no head or merge check results. Current-head CI is not validated.
Performance
The stored plan hash is still computed outside the cache monitor, and prepared native-scan metadata still shares its plan's eviction unit. The transported key avoids repeated derivation on both RDD and native-shuffle paths. The merge adds no new cache lookup, copy or locking mechanism. Concurrent cold preparation can still duplicate work, and the 16-entry limit bounds plans or shuffle IDs rather than bytes. No new performance issue or fresh performance claim was identified.
Design
The ownership boundary remains unchanged: cached immutable plan/common metadata, task-local partition injection, and shuffle cleanup through the manager lifecycle. The upstream scan-filter flag integrates through the existing finalized-byte guard without requiring another cache or identity mechanism.
Abstraction & complexity
This merge introduces no new cache abstraction or helper. The existing plan holder, stored-hash key and private preparation helper keep ownership and reuse in one place. No additional abstraction concern was identified in this update.
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Re-reviewed 27269632 against 719cba11. The full feature patch is byte-identical to the previously reviewed pair, and all 13 feature files are unchanged. The entire head update matches the base update and adds only CometExplodeBenchmark.scala. I found no new P1/P2. The existing approval can stand.
The maintained Spark 3.5 and 4.0 source comparison still supports the existing lifecycle fix: shuffle IDs belong to a SparkContext, normal removal reaches unregisterShuffle, and environment shutdown calls the manager's stop. Both Comet managers retain their cleanup hooks. The finalized-common-byte guard still separates scalar-subquery results across executions, including has_data_filters, while partition files and executable task state remain task-specific. No expression type, null, boundary, overflow, ANSI-mode, error or fallback behavior changed in this update. Maintained Spark 3.4 and 4.1 branches were unavailable, so those versions are not newly qualified.
Validation and CI
Source and Git identity checks passed. The dependency manifests and native lockfile are unchanged. No component, JVM, native or benchmark execution was rerun. Earlier component passes and author timing tables remain historical evidence. At 2026-09-05T05:42:43Z, all three current-head workflows were action_required, with no head or merge check results. Current-head CI is not validated.
Performance
This update adds no cache lookup, synchronization or copying cost. The stored plan hash, plan-owned prepared commons and transported scan key are unchanged. Concurrent cold preparation can still duplicate work, and the 16-entry limit still bounds entries rather than bytes. No new performance finding or measured gain is claimed.
Design
The ownership and lifetime boundaries are unchanged. The added upstream benchmark does not alter cache identity, subquery finalization, partition injection or manager cleanup. No new design concern was identified.
Abstraction & complexity
No cache abstraction or helper changed. The existing holder and preparation helper retain the same responsibilities, with no additional complexity introduced into this feature.
andygrove
left a comment
There was a problem hiding this comment.
I checked this out locally and ran it, so most of what follows is measurement rather than reading. No correctness problem found. I built the case I thought most likely to break the design, the same query text run twice with a different scalar-subquery result so the finalized common changes under an unchanged transported key, and Comet matched Spark exactly on both runs. All 58 tests across PlanDataInjectorSuite, PlanDataInjectorShuffleLifecycleSuite, CometScanWithPlanDataSuite and CometCelebornShuffleManagerSuite pass on Spark 4.1. I also ran the new lifecycle suite ahead of two other Spark suites in one JVM, since it stops and recreates the context and CI runs it mid-list, and it does not poison anything after it. I reproduced your ~4-8x on the real repo code rather than an extracted harness, so the cache is clearly earning its place.
Two things came out of instrumenting it that I think are worth acting on.
The first is that the base plan cache does not fire in the default configuration. I went looking for a cache hit in a real query and could not get one. parseBasePlan is only reachable from CometExecRDD.compute, and with native shuffle on the scan fuses into the shuffle writer's per-task plan and takes injectPlanDataForShuffle instead. On a 300-column table with 16 files, taking a snapshot of both stores immediately after each action:
| shape | basePlanCache |
shuffle store |
|---|---|---|
count(), partial agg into native shuffle |
0 | 1 |
| noop write, map-only | 1 | 0 |
collect() of 300 columns, map-only |
1 | 0 |
count() with spark.comet.shuffle.mode=jvm |
1 | 0 |
count() with spark.comet.shuffle.enabled=false |
1 | 0 |
So the transported key and the shuffle-scoped prepared commons carry the shuffle path, and the base plan cache carries map-only stages and the JVM-shuffle configuration. Both are real wins, and your CometNativeShuffleWriter comment already explains why there is no entry on that path. But since every aggregate and every join has the fused shape, the description reads as though the base plan cache is the main event for most stages when it is really the other two pieces doing that work. Could you say which stage shapes each piece covers, and add an end-to-end assertion that the base plan cache is actually hit? PlanDataInjectorShuffleLifecycleSuite already has the pattern with preparedShuffleSnapshot, so an equivalent view over basePlanCache would make it a few lines. For what it is worth, retained bytes for one 300-column entry measured 39746 as a serialized-size proxy, so the 16-entry cap is a small number and I would not treat memory as a concern.
The second is about where the remaining cost sits. I expected to confirm the byte-hashing-under-the-monitor concern and it turns out not to be there at all. Taking the monitor and doing the full Arrays.equals under it costs 1.02x to 1.13x of the hash alone, and hit throughput scales close to linearly, 24.9K to 49.7K to 98.4K to 168.8K to 229.3K ops/s over 1, 2, 4, 8 and 16 threads on 16 cores. Your fix for that one worked, and I think the 1/8/32-thread latency table was showing CPU oversubscription rather than contention. My own first harness had the same defect and I nearly reported it.
What is left is Arrays.hashCode in PlanKey, which is about 31% of the whole cached path:
1000-column plan, 47900 plan bytes, 16 cores, JDK 17
hashOnly = 40.19 us/op (1 thread)
cacheHit = 41.28 us/op (1 thread)
full cached path = 130.46 us/op vs uncached 1054.85 us/op
On JDK 17 that is a scalar byte-at-a-time loop, because the vectorized hashCode intrinsic only landed in JDK 21 and we still support 11 and 17, so it is a full pass over the plan bytes on every task. Could the driver hand that over the same way it now hands over source_key? The bytes themselves are the task binary so they cannot be transported, but a planFingerprint: Long computed once in doExecuteColumnar and passed as another CometExecRDD constructor field would ride along for free. PlanKey.hashCode could then just return it while equals keeps Arrays.equals for collision safety. That should take the 1000-column case from 8.1x to somewhere around 11x.
A handful of smaller things, none of them blocking.
Now that injectPlanDataForShuffle supplies a memo, is IcebergPlanDataInjector's own commonCache still buying anything? Both production call sites pass a non-null preparedCommons, so prepareCommon already runs at most once per entry and key and byte set, and the inner cache adds a second retained copy of each Iceberg common plus a global monitor held across the parse. The comment above it still says it serves the native shuffle writer's per-task plans, which was true before this PR but is not any more. Relatedly, injectPlanData(op, commonByKey, partitionByKey) and the preparedCommons == null branch in prepareShared look like they have no production caller left. Would you mind folding the tests onto the two real entry points, or marking the overload as test-only? As it stands a future call site can pick the un-memoized overload by accident and nothing will complain.
QueryContextInterner's class comment says the plan bytes are "re-parsed and re-serialized per task by CometExecRDD.compute". Since this change the re-parse is once per stage on that path. The re-serialize and the native decode still make the interner pay back, so it is just the word "re-parsed" that needs to come out.
The stop() hook clears shufflePreparedCommons but not basePlanCache, so a stopped context's plan trees and prepared commons sit in the static field until 16 new plans evict them. That is bounded and self-healing, so it is not the shuffle-id problem over again, but since the hook already exists, is there a reason not to clear both there?
Two things about the new SPI shape. prepareCommon returning AnyRef and inject casting it means nothing checks that an implementor's two halves agree on a type, and the stub in CometScanWithPlanDataSuite is already an example that compiles without agreeing. Would an abstract type Prepared <: AnyRef on the trait be worth the existential at the map boundary? Separately, preparedCommons is keyed on the source key alone. The registry comment points out that every contrib scan arrives as the same CONTRIB_SCAN envelope, so two contrib injectors in one plan that agreed on a key and on the common bytes would hand one of them the other's prepared object. That needs a key collision plus byte-identical commons across two different formats, so I do not think it is reachable, but prefixing the memo key with the injector class would close it for free.
shufflePreparedCommons bounds itself with maxCachedBasePlans even though its entries are shuffle ids. Same number today, different thing being counted, so a separate maxCachedShuffles would read correctly and let the two move apart later.
cachedOrCompute is correct because both callers pass a Collections.synchronizedMap, whose mutex is the wrapper itself, so cache.synchronized takes the same monitor the map's own methods take. The parameter is typed JMap[K, V] though, so handing it a ConcurrentHashMap would compile and quietly drop both the atomicity and the LRU bookkeeping. Worth a line in the scaladoc saying that is the contract.
source_key is s"${common.getSource}_${hash}", and common.source is already in the same NativeScan message that canInject requires to be present. Would shipping just the int32 hash and rebuilding the string in getKey be better? Same O(1) executor read, and it stops re-shipping the simpleStringWithNodeId text the plan already carries, which seems worth having given #5200 was about plan size.
Finally, a couple of lines that look like they can go. PlanKey.equals opens with (this eq that), but HashMap.getNode does the reference check before calling equals, so from the only caller that branch is unreachable. And in the recreated-context test the second new CometShuffleManager(...) is built and thrown away, and the assertion holds without it, so dropping it would make it clearer that stop() is the thing under test.
Things I did not cover: Celeborn and Iceberg end to end, the Spark 3.4, 3.5 and 4.0 profiles, and CometNativeShuffleSuite and CometExecSuite which you report green.
…cache by it Every task of a stage built its PlanKey by hashing the full plan bytes with Arrays.hashCode, a scalar loop that took about a quarter of the cached path on JDK 17 for a wide plan. The driver now computes one XXH64 fingerprint where it serializes the plan and carries it on CometExecRDD, so the executor probe hashes nothing. Equality still compares the bytes, so a fingerprint collision cannot serve another stage's plan. The reference check in PlanKey.equals goes too, since HashMap already performs it before calling equals.
…on stop Type the injector's two halves through an abstract Prepared member so an implementation's prepareCommon and inject must agree, and move the single remaining cast to the memo boundary. Prefix memo keys with the injector class: every contrib scan arrives as the same CONTRIB_SCAN envelope, so two injectors agreeing on a key and byte-identical commons would otherwise share one prepared object. Drop the Iceberg injector's inner cache and the injectPlanData overload with no production caller, since both real entry points now memoize. The shuffle store gets its own maxCachedShuffles bound, the shuffle managers clear the base plan cache together with the shuffle store on stop, and the lifecycle suite asserts which stage shapes land in which store: a map-only stage in the base plan cache, a scan fused into a native shuffle under the shuffle id.
The transported key was the scan's source followed by a hash, and the source already sits in the common of the same message. The proto now carries the int32 hash alone as an optional field, so presence is explicit and a zero hash is still a transported value, and the executor rebuilds the key from the source next to it. The field is unreleased, so it changes in place.
QueryContextInterner described the plan as re-parsed per task by CometExecRDD.compute; parsing now happens once per stage on that path, while the per-task re-serialization it also mentions still holds.
…ey's classloader assumption
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Re-reviewed 2f5d128e against 8e684685, including the full contribution and the update since the published review at 6c5dbc80. All 15 feature files and the complete feature patch are unchanged. The entire 34-file increment comes from the base update. I found no new or remaining P1/P2 issue, and my existing approval stands.
I checked the newly inherited Variant normalization and concat_ws paths against the cache boundary. Variant conversion operates on incoming batches and creates its rewrite state locally. Runtime scalar subqueries still resolve through each native execution context, and the new concatenation adapter returns a scalar for batch broadcasting without retaining its value in the JVM metadata memo. Native planning, task bindings, full-byte cache admission and manager cleanup remain unchanged. The earlier remote-decoder lifetime analysis also remains applicable through exact source identity.
Maintained Spark 3.5/4.0 sources confirm the relevant concatenation rules: ordered string/array flattening, null-separator propagation and skipping null inputs. Spark 4.0's Variant writer and lookup use Java string ordering. These checks address interaction with this cache contribution, not full qualification of the inherited features. The cache update adds no expression-type, null, ANSI/overflow or fallback change. Maintained Spark 3.4/4.1 sources remain unavailable.
Current pagination and body comparisons found no new or edited discussion. All four source-addressed threads remain unresolved. Four diff whitespace checks passed. Synthetic merge 6dc4c190 has the exact assigned base/head parents and a tree equal to HEAD. CI has now started. At 2026-09-08 16:56:38 UTC, the head had 13 successful checks, six running, seven skipped and one queued, with no failed checks. This is not a completed suite or an executed-checkout provenance claim. No local build, runtime test or benchmark was run, and earlier execution reports are not current-head results.
Performance
The feature's lookup, preparation and injection paths are unchanged. Full-byte comparisons, possible duplicate cold preparation, entry-count retention limits and per-task serialization/native decoding still apply. The inherited Variant work and concatenation optimization operate during execution. Their cost or benefit cannot be inferred from plan-cache timings. The new hash benchmark measures native SQL/shuffle kernels rather than the driver plan fingerprint. I make no new measured speedup claim.
Design
The added base paths preserve the ownership split: shared entries contain plan/common metadata, while physical expressions, subquery values and batch conversion remain within native execution. Neither integration requires another cache owner or changes the existing shutdown hooks.
Abstraction & complexity
The typed injector SPI and cache helper are unchanged. Dependency versions and the lockfile are unchanged, with only benchmark registrations added to the expression manifest. The inherited helper implementations introduce no new abstraction in this contribution and expose no actionable simplification of its cache design.
|
@andygrove everything from your review is in, including the driver fingerprint. Ready for a look when you have time. |
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Re-reviewed 7d65a6b6 against fefee03d, including the full contribution and the update since review 5144681707 at 2f5d128e. All added and deleted lines in the 15-file contribution are unchanged in a zero-context diff, excluding only blob IDs and hunk positions. The 64-file increment consists entirely of inherited base changes. Eleven feature files remain byte-identical, and the four overlapping files contain the base's changes alongside the unchanged contribution. The complete injector/cache implementation is byte-identical. I found no new or remaining P1/P2. My existing approval stands.
The new dynamic-filter setting is written into the HashJoin protobuf, so changing it changes the plan bytes used for cache identity. Its mutable build domain is created inside native execution, after per-task injection and JNI decoding. DynamicFilterJoinExec creates a fresh predicate and resets its producer for each stream. Neither the predicate nor the join's build state enters the shared JVM metadata cache. The inherited test covers simultaneous executions, dropping one while the other remains live, reset and child replacement. I inspected that test. It was not run locally.
The broader QueryStageExec input boundary includes table-cache stages without moving cached batches into PlanDataInjector. Maintained Spark 3.5/4.0 sources confirm the table-cache stage and InMemoryTableScanLike contracts used here. The runtime daemon-thread change preserves the existing shuffle-manager cleanup path. Both managers still clear both metadata stores on stop, and shuffle removal still releases its prepared commons. Exact common-byte checks, task-specific file injection and subquery registration are unchanged. These cache interactions add no new null, type, ANSI/overflow or fallback-semantic change.
All current discussion was checked, including the author's readiness note. It adds no new implementation or measurement. Three diff checks passed. Merge 6fb3e52e has the assigned base/head parents and a tree equal to HEAD. At 2026-09-09T04:32:35.996689+00:00, only the labeling check had succeeded. CI, CodeQL and the Delta build gate required approval. No executed build/test checkout, local runtime result or native artifact is credited. Maintained Spark 3.4/4.1 source coverage remains unavailable.
Performance
The cache's lookup, preparation, eviction and serialization work is unchanged. Full-byte equality remains on hits, cold races can duplicate preparation, and the two 16-entry limits bound plans/shuffle IDs rather than retained bytes. Evicted entries can remain live while tasks use them. The inherited dynamic-filter and expression optimizations do not add work to the metadata memo and do not validate its performance.
The reported timings still describe earlier component measurements. No matched benchmark of this exact head/base, native artifact or current output reconciliation was supplied or run in this review. JVM setup savings remain separate from per-task serialization, JNI decoding, native planning and whole-query time. I make no new speedup claim.
Design
The base update preserves the ownership boundary: immutable operator/common metadata may be shared, while inputs, native plans, runtime predicates and execution contexts belong to each task. The inherited join wrapper additionally scopes its mutable producer/consumer to a stream. AQE cache materialization remains with Spark's cache builder, and manager shutdown remains the normal metadata cleanup boundary.
Abstraction & complexity
The typed Prepared SPI, class-prefixed memo key and synchronized-map helper are unchanged. No additional cache abstraction or lifecycle owner is introduced by this update. The inherited dependency change adds the already-locked base64 version to the Iceberg writer. It does not change the injector contract or establish out-of-tree Delta compatibility. I found no actionable simplification beyond the changes already addressed in the earlier discussion.
andygrove
left a comment
There was a problem hiding this comment.
prepareShared in operators.scala (used by both the base-plan cache and the shuffle-scoped store, around line 378) does a plain get then, on a miss, computes prepareCommon and calls put unconditionally. That is different from cachedOrCompute (around line 188), which re-checks for a winner after computing and adopts it if one already landed. Two task threads racing the first access to a given scan's common, which is the normal case at the start of a stage since an executor's task slots all start around the same time, can each parse their own Prepared instance and return it to their own caller, and whichever thread's put runs last is what stays in the map. The byte comparison on every hit still protects correctness, since a losing instance is content-equal to whatever ends up cached, so I do not think this produces wrong results. But the single-shared-instance guarantee the PR body describes for racing threads holds for parseBasePlan and not for this per-scan memo, and PlanDataInjectorSuite's 200-trial barrier test covers that guarantee for parseBasePlan while I don't see an equivalent for prepareShared. Is that gap intentional, or is it worth giving prepareShared the same adopt-the-winner check?
The rest of what I looked at holds up. The base plan cache key compares full bytes rather than just the fingerprint, so a fingerprint collision cannot serve the wrong plan, and it is keyed by content rather than anything like stage id that could alias across jobs or SQL executions. Eviction is scoped to the whole plan entry now, so a wide plan's scans cannot be evicted out from under it individually. The shuffle-scoped store is released from unregisterShuffle and stop in both shuffle managers, and the local-fallback branch in CometCelebornShuffleManager.unregisterShuffle is covered too since it delegates to CometShuffleManager's own release. Nothing mutates a shared common or plan after publication, inject only ever calls toBuilder/setCommon to build a new message on top of it. Speculative execution and task retries do not need special handling here either, since the cache keys are pure functions of content rather than task or attempt identity.
…n key prepareShared put its freshly prepared common unconditionally, so tasks racing the first access to a scan could each keep their own instance. It now merges into the memo atomically and returns whichever equal-bytes entry landed first, on both the base-plan and shuffle-scoped paths. Barrier tests cover both.
|
Not intentional, the per-scan memo just never got the same treatment. In 0b2a0b8 |
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Re-reviewed 0b2a0b8cfbb8d499113c5e14ab4d0afec07591c3 against 5627ab8c01473b186db2fbe73c91c599b1387f9b, including the full contribution and the update since review 5149785454 at 7d65a6b6. Thirteen of the fifteen feature files remain byte-identical. The update adds the two-file winner-adoption fix on top of a sixteen-file base merge, with no merge-specific feature edits. I found no new or remaining P1/P2.
The change addresses the prepareShared race comment. Both preparation paths now return an equal-bytes entry that won the atomic merge, rather than retaining their own losing parse. If an overlapping execution supplies different finalized bytes, the callback installs and returns that caller's fresh value. A previously matched local hit also remains valid if another caller later replaces the slot. Preparation still happens before publication, so a parsing failure does not install a failed replacement. This corrects the sharing guarantee without treating the old content-equivalent race as a wrong-result bug. The atomic selection follows the Java ConcurrentHashMap.merge contract.
The finalized-byte guard continues to distinguish execution-specific scalar-subquery filters. Maintained Spark 3.5/4.0 still resolve subqueries before execution, and the Comet bridge preserves that order. Partition files, iterators and native state remain task-specific. The earlier shuffle-lifetime fix remains intact in both managers and matches the maintained removal/shutdown paths. The inherited JNI, cached-string statistics, Parquet conversion and Iceberg writer changes do not move execution objects into the metadata memo. This update adds no expression type, null, ANSI/overflow or fallback behavior. Maintained Spark 3.4/4.1 sources remain unavailable.
Validation and CI
The two new tests each run four callers over 200 cold-key trials, checking common-object identity through injectPlanData and injectPlanDataForShuffle. They cover both paths that the old parseBasePlan race test did not. They do not count preparation calls, force every caller to observe the miss, or exercise concurrent different-byte replacements. The existing stale-common test is serial, and the existing partition-isolation test checks distinct file lists. I inspected these tests but did not execute them. The author's before/after split-instance counts remain reported evidence.
Three diff checks and source/tree identity checks passed. Merge b5666bea07c134baf67b193b242f48ffe5bd6afd has the assigned base/head parents and the same tree as HEAD. At the final check on September 9 at 15:36 UTC, only the labeling check had succeeded. CI, CodeQL, the Delta gate and title check require authorization, each with zero jobs. No current-head product execution, local runtime test, native artifact or benchmark is credited. An independent metadata refresh at 15:36 UTC still reports mergeable=false, dirty and no current merge commit, while the assigned head/base remain unchanged and Git confirms the base is an ancestor of the head. The fetched merge establishes source identity only. This source review does not establish that GitHub currently considers the PR mergeable.
Performance
The warm path still performs a lookup and exact-byte comparison. Cold races may still parse more than once because preparation remains outside the atomic operation. The improvement is that losing prepared objects are discarded and callers share the winning object. The merge callback adds a byte comparison when an entry already exists and performs no parsing under the map's update lock. I found no new performance issue and make no fresh speedup claim.
The existing entry-count limits, active references surviving eviction, per-task serialization and native decoding remain unchanged. The earlier component timing tables do not measure this winner-adoption change or establish whole-query benefit.
Design
Using one atomic map operation keeps publication and winner selection together while preserving the existing exact-byte replacement rule. Both injection paths use the same helper. Sharing is scoped to a live memo entry: eviction, shutdown or intervening different-byte replacements can create another generation. The fix introduces no new cache owner or cleanup obligation.
Abstraction & complexity
The change stays inside prepareShared, with one test helper covering the two production entry points. The typed injector SPI and cast at the memo boundary are unchanged. No additional synchronization abstraction or update-specific simplification is needed.
# Conflicts: # spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Re-reviewed ca46057284413358a1b0bee7d32385f403bd6154 against 424c31aa79d13fddf743ffa29bae3c6f146e6c5e, including the full contribution and the update since review 5156527271 at 0b2a0b8c. The six-file head increment matches the base increment. All additions and deletions in the 15-file contribution are unchanged after removing diff metadata. Fourteen feature files remain byte-identical, and the overlapping operators.scala change comes entirely from the base. The complete cache/injector section is byte-identical. I found no new or remaining P1/P2.
The prepared-common race fix remains intact in both injection paths. Atomic winner selection still checks the exact finalized bytes, and preparation failures cannot publish a replacement. The earlier shuffle-lifetime fix also remains intact: both managers release individual shuffles and clear both stores on stop. Maintained Spark 3.5/4.0 still assign shuffle IDs per context and route normal cleanup and shutdown through these hooks. Late injection can repopulate an entry, so permanent emptiness after shutdown is not claimed.
I checked the inherited FIRST/LAST PartialMerge enablement at the cache boundary. Maintained Spark 3.5/4.0 use nullable value/value-set buffers and document order-dependent results after shuffle. Comet still serializes aggregate modes and ignoreNulls into the plan. Each task gets a new native execution context and constructs its own aggregate operators and stream. Sharing the immutable protobuf does not share aggregation buffers or evaluated values. Finalized scan-common bytes and per-partition files remain separate, including execution-specific subquery filters. These are integration checks, not fresh runtime qualification of the inherited aggregate change. The cache contribution adds no expression-type, null, ANSI/overflow or fallback-semantic change in this update.
Validation and CI
Read all current discussion and refreshed 34 prior source comparisons. Only operators.scala changed, outside the cache/injector section. The five checked dependency manifests/lockfile are unchanged. Three diff checks and source/tree checks passed. Merge d49db7091c94717fcd07d8a7719652b1715268c6 has the assigned base/head parents and the same tree as HEAD.
At 2026-09-09T18:25:56.934824+00:00, CI, CodeQL and the Delta build gate require approval, each with zero jobs. The successful labeling job checked out the base commit, so it provides no changed-head compile or runtime coverage. No local build, runtime test or benchmark was run. The existing race/lifecycle tests were inspected, and earlier component passes and author measurements remain historical evidence. Maintained Spark 3.4/4.1 sources remain unavailable.
Performance
The update adds no cache lookup, copy, synchronization or preparation step. Driver fingerprinting, exact-byte comparisons, winner adoption and both eviction boundaries are unchanged. Cold races can still duplicate preparation, active callers can retain evicted objects, and the 16-entry limits count plans/shuffle IDs rather than bytes. Per-task serialization and native decoding remain. The inherited aggregate enablement changes which plans can execute natively, but earlier cache timings do not measure that effect. I found no new performance issue and make no fresh speedup claim.
Design
The conflict resolution preserves the previously reviewed ownership boundary. Cache entries retain immutable plan/common metadata, while aggregate state, partition inputs and native execution belong to each task. The broader aggregation eligibility requires no new cache identity or invalidation mechanism: serialized plan changes pass through the existing full-byte equality check. Both cleanup paths and the tests for map-only and fused-shuffle stages are unchanged.
Abstraction & complexity
No cache abstraction, injector SPI or preparation helper changed. The typed Prepared contract, injector-specific memo key and atomic selection remain in the same place. The inherited removal of FIRST/LAST fallback checks does not add another cache owner or synchronization layer. No update-specific simplification is required before merge.
andygrove
left a comment
There was a problem hiding this comment.
Everything from both of my earlier rounds is addressed, and I checked each one against the head rather than the summary.
prepareShared now merges instead of putting, and ConcurrentHashMap.merge is the right primitive: the remapping function runs atomically for the key, adopts an equal-bytes entry that landed meanwhile, and installs the fresh one over a stale or absent slot. prepareCommon stays outside the map's lock, so racing threads still each pay one parse but only one instance survives, which is the tradeoff I wanted. The asymmetry with cachedOrCompute, which adopts the winner without comparing content, is correct rather than an oversight: PlanKey.equals already compares the full bytes, so an equal key implies equal content, whereas prepareShared's key is a scan-key string and needs the byte compare.
I did not want to take the new tests on faith, so I reverted only the prepareShared hunk back to the unconditional put, kept the tests, and ran them:
- injectPlanData gives racing tasks on a cold scan key the same prepared common *** FAILED ***
113 of 200 trials prepared more than one common
- injectPlanDataForShuffle gives racing map tasks on a cold scan key the same prepared common *** FAILED ***
173 of 200 trials prepared more than one common
With the fix restored, all 27 PlanDataInjectorSuite tests pass. So these are not tests that happen to pass, they fail loudly on the exact defect and on both paths. That is the strongest form of the guarantee the PR body claims, and it now matches parseBasePlan's 200-trial barrier test rather than being weaker than it.
The rest of the first round is in too. planFingerprint is computed on the driver and carried, PlanKey.hashCode returns it while equals keeps Arrays.equals, so the Arrays.hashCode pass that was 31% of the cached path is gone from the executor. source_key_hash ships as an optional int32 with the string rebuilt in getKey, which stops re-shipping the simpleStringWithNodeId text. type Prepared <: AnyRef on the trait means an implementor's two halves have to agree, and preparedKey prefixes the injector class so two contrib injectors agreeing on a scan key cannot swap prepared objects. IcebergPlanDataInjector's own commonCache is gone, the un-memoized injectPlanData overload is private now, stop() clears both stores, maxCachedShuffles is its own constant, and cachedOrCompute's scaladoc states the Collections.synchronizedMap contract that its JMap parameter cannot enforce.
The two coverage asks landed as well, and they are the ones I care most about because they turn my instrumentation into something CI keeps honest. a map-only stage hits the base plan cache on every task after the first asserts the end-to-end hit I could not get in the default configuration, and a scan fused into a native shuffle prepares under the shuffle id, not the base cache pins the stage-shape split so nobody later reads the base plan cache as covering aggregates and joins. All three suites are registered in both pr_build_linux.yml and pr_build_macos.yml.
Approving. Nice piece of work, and thanks for taking the long way round on the race rather than arguing that the byte compare made it benign.
# Conflicts: # spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
|
Thanks for reverting the |
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed 5c22391e since review 5158279081 at ca460572. No new or remaining P1/P2 findings. The three merges from main preserve the cache contribution and its tests. The 87-file increment matches the base changes after accounting for the retained XXH64 import. The prepared-common race fix, exact finalized-byte guard, per-partition file isolation, and both shuffle-manager cleanup paths are unchanged.
I checked the inherited aggregate and object-store changes at the cache boundary. Parsed protobufs remain shared metadata, while each task builds its own native context and execution state. The new maintainer report describes both race tests failing with the fix reverted and all 27 tests passing after restoration at ca460572. That is prior-head reported execution, not a fresh run of this head.
At 2026-09-11T06:43:51.427800+00:00, CI, CodeQL and the Delta build gate require approval, each with zero jobs. The successful labeling job checked out base 5e302d99. Merge 9a4ff9e1 has the assigned base/head parents and the same tree as HEAD, but has no product-test coverage. Source comparisons and three diff checks passed. No local product tests or benchmarks were run. Canonical Spark checks used maintained 3.5/4.0 sources only. 3.4/4.1 remain unavailable.
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed f99475fb after review 5175689710, whose body covered 5c22391e. No new or remaining P1/P2 findings. All 15 contribution files are byte-identical to that revision, and the complete 35-file increment exactly matches the base update. The cache race fix, finalized-common byte guard, per-partition file isolation and shuffle cleanup paths remain unchanged. GitHub reports the existing approval on this head; preserving it.
I checked the inherited collect/min-max aggregate and Parquet I/O-metric changes at the cache boundary. Native execution contexts, aggregate functions and scan counters are still constructed per task; the shared cache retains immutable protobuf metadata. Generic query-context traversal covers the new aggregate messages. No cache hot path changed, and no new runtime or benchmark evidence was reported.
At 2026-09-11T13:26:20.610577+00:00, CI, CodeQL and the Delta build gate require approval with zero jobs each. The successful labeling job checked out base 8320ae48. Merge 924b9934 has base/head parents and HEAD's tree, but no product execution. Source comparisons and three diff checks passed; no local product tests or benchmarks were run. Canonical Spark checks used maintained 3.5/4.0 sources; 3.4/4.1 remain unavailable.
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed f99475fb → 38c16318 against base db790673 since review 5179221562. No new or remaining verified P1/P2 findings. The 45-file increment matches the base update after removing diff index and hunk-location metadata. Twelve of the 15 contribution files are byte-identical; the other three contain inherited workflow/protobuf changes. The cache race fix, exact finalized-common byte guard, per-partition file isolation and both shuffle-manager cleanup paths remain unchanged. Preserving the existing approval.
I checked the new Iceberg deletion-vector pools through the cache boundary: the whole finalized common is retained, pool changes invalidate the prepared entry, and each task receives its own file-task list. Native delete-file objects and Parquet adapters are constructed independently of the JVM cache. The inherited DataFusion 55.1 update and map kernel do not add execution state to cached protobufs.
At 2026-09-12 19:19:19 UTC, CI has 53 successful, 10 skipped and 3 failed checks. Both underlying failures stopped in Java setup with ECONNRESET: Celeborn 0.7 and Spark SQL core shard 3. Required Checks consequently failed. On verified merge 22d568de, with the assigned base/head parents and HEAD’s tree, Spark 4.1 exec passed 869 tests, including all 27 injector and 4 lifecycle cases; scans passed 509, including both deletion-vector cases; shuffle passed 489. CodeQL and the Delta gate succeeded. Local source, suite-registration and CI-configuration checks passed. No local product build/tests or benchmark ran. Maintained Spark 3.4/4.1 source gaps remain; CI execution does not close those source gaps.
Which issue does this PR close?
No dedicated issue. #5200 fixed the size of the serialized plan; this addresses the per-task work done on those bytes.
Rationale for this change
The serialized plan bytes are identical for every partition of a stage, yet every task parsed the full operator tree from bytes, re-derived the scan's source key (stringifying the schema and filter lists, which turns out to be the single most expensive step for wide schemas), and re-parsed the scan's common message, all before injecting its own partition data. That cost scales with plan size times partition count and lands hardest on large scan plans.
What changes are included in this PR?
Three pieces, split by stage shape. The parsed base plan is cached per executor in a 16-entry LRU keyed by the plan bytes; a 64-bit fingerprint of those bytes is computed once on the driver and shipped with the RDD, so executors never hash the plan bytes themselves, while equality still compares the bytes. The scan's source key hash is computed once on the driver and carried in the
NativeScanproto, so executors rebuild the key from the source string already in the message instead of re-deriving it from the schema and filter lists. For scans fused into the native shuffle writer's per-task plan, a shuffle-scoped store (also 16 entries) keeps the prepared common per shuffle id, released inunregisterShuffleand, together with the base plan cache, when the shuffle manager stops.The two caches split the stage shapes between them. With native shuffle enabled the scan fuses into the shuffle writer's per-task plan, so every map task takes
injectPlanDataForShuffle: the transported source key hash avoids re-hashing the common and the shuffle-scoped prepared commons share one parsed common across the shuffle's map tasks, while the base plan cache is never touched. Map-only stages (collect, noop writes) and every stage underspark.comet.shuffle.mode=jvmor with Comet shuffle disabled go throughCometExecRDD.compute, where the base plan cache parses the plan once per stage and every later task hits it.PlanDataInjectorShuffleLifecycleSuiteasserts both shapes end to end.Injection itself stays per task, since partition data genuinely differs, and the injected tree is never cached, so per-partition file lists cannot leak across tasks (there is a test asserting the shared common is reference-equal while the file lists diverge). Cache misses compute outside any lock, and threads racing a cold key end up holding the same instance, first insert wins; this holds for the base plan entry and for each scan's prepared common in both the base-plan memo and the shuffle-scoped store, so a stage's tasks share one prepared common whichever path they enter through. Plan-data injectors now declare the type of their prepared common (
type Prepared), and the memo key carries the injector class so two contrib injectors agreeing on a scan key stay separate. The Iceberg injector's own common cache is gone, since the shared memo already runsprepareCommonat most once per entry.A larger follow-up was considered and set aside: shipping the base plan to native once per executor and merging partition data there would also remove the per-task reserialize and native decode, but injection is a ServiceLoader SPI implemented by out-of-tree modules, so moving the merge native would break that extension point. Noted for later rather than folded in here.
Measured per-task cost (parse plus key derivation plus reserialize, 5000 iterations after warmup): a 100-column scan plan goes from roughly 274-380us to 44-73us, and a 1000-column plan from roughly 2.0-2.5ms to 0.55-0.93ms. Moving the plan-bytes hash to the driver takes a further 21 percent off the cached path on a 1000-column plan (77.6us to 61.3us per task, single thread, JDK 17), since
Arrays.hashCodeover the bytes was a quarter of it on JDK 17's scalar loop.How are these changes tested?
PlanDataInjectorSuite (27: hit and miss behavior, distinct plans staying separate, eviction for both stores, eight-thread concurrency, cold-key races adopting one instance across 200 barrier-synchronized trials for the parsed plan and for the prepared common on both injection paths, shared-common reference equality with per-partition file isolation, the transported key matching a fresh derivation, two injectors sharing a key, and a stop that clears both stores), CometScanWithPlanDataSuite (4), PlanDataInjectorShuffleLifecycleSuite (4: the map-only stage hits the base plan cache, the fused shuffle shape fills the shuffle store and not the base cache, per-shuffle release, and a recreated context starting empty), CometCelebornShuffleManagerSuite (29), CometNativeShuffleSuite and CometExecSuite (182 together), all green on Spark 3.5. Spotless and scalastyle clean.