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.

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

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

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