perf: cache parsed plan data across a stage tasks

[dwsmith1983]

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 bounded per-executor LRU caches, 16 entries each. The parsed base plan is cached keyed on byte content, so an executor parses a stage's tree once instead of once per task; the parsed NativeScanCommon is cached the same way the Iceberg injector already caches its common; and the source key gets a memo that rides protobuf's reference-identity fast path once the base plan instance is shared. 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 two threads racing a cold key both end up holding the same instance, first insert wins.

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.

How are these changes tested?

Eight new tests in PlanDataInjectorSuite (hit and miss behavior, distinct plans staying separate, eviction plus rerun, eight-thread concurrency, cold-key race adopting one instance across 200 barrier-synchronized trials, shared-common reference equality with per-partition file isolation, and the memo matching a fresh derivation), alongside the existing six. The end-to-end paths run through CometScanWithPlanDataSuite (5), CometNativeReaderSuite (54), CometExecSuite (142), and CometNativeShuffleSuite (40), all green. Spotless and scalastyle clean.

[sunchao]

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]

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.

[sunchao]

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.

[dwsmith1983]

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.

[sunchao]

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.

[dwsmith1983]

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.

[sunchao]

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.

[dwsmith1983]

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.

[sunchao]

Sorry for back and forth @dwsmith1983 . I just added a few more instructions to my Comet PR review skill especially for perf PRs. Please check if these comments make sense.

[sunchao]

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.

[dwsmith1983]

Took your preferred direction and it worked out well, thanks for the concrete numbers, they made the case obvious. Pushed in 36fa94d:

• The scan key is now derived once on the driver and carried in the NativeScan proto (new source_key field, native ignores it). The executor reads it directly, so the keyCache LRU is deleted along with derivation after eviction. The shuffle writer path benefits too since the key rides inside spec.childNativeOp.
• Prepared per-scan data now lives inside the base plan's own cache entry (CachedPlanData holding the parsed plan plus its prepared commons), so plans and their scans share one eviction unit and the standalone common cache is gone. The shuffle path gets an equivalent store scoped by shuffleId.
• The base plan cache keys on a PlanKey that stores its hash, computed once per task outside the monitor. Driver transport for that one was not practical since the plan bytes are the task binary itself, so this is your middle option.
• On your identity caveat: sourceKey does not cover the scalar subquery filters finalized after planning, so each prepared entry pins the exact finalized bytes it was parsed from and a hit is honored only on byte equality; a changed common under the same key re-prepares and replaces, never serves stale. There is a test proving the replacement.
• The memo test now compares getKey with sourceKey(common) directly, per your correction.

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:

threads	uncached parse	previous	now
1	111.6	29.9	28.2
8	164.1	242.6	36.5
32	500.0	2348.2	117.8

Your churn shapes, warm pass, prepared commons reused:

shape	previous	now
1 plan x 17 scans	0/17	17/17
9 plans x 2 scans	0/18	18/18

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]

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.

[sunchao]

[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.

[dwsmith1983]

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]

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.

[sunchao]

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]

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]

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.