Diveyam-Mishra opened a new pull request, #58476:
URL: https://github.com/apache/spark/pull/58476
### What changes were proposed in this pull request?
This PR resolves an overhead issue with state-store instance metric
reporting in `StateStoreWriter`.
Currently, `StateStoreWriter` creates separate `SQLMetric` accumulators for
every `(partitionId, metric, storeName)` combination upfront on the physical
plan. On executors, every task deserializes and registers all of those
accumulators in its `TaskContext`, and includes all of them in the task result
payload even though only a single partition was processed.
This PR replaces the upfront $O(\text{numPartitions})$ `SQLMetric` instances
with a single `instanceMetricsAccumulator:
CollectionAccumulator[(StateStoreInstanceMetric, Long)]` on `StateStoreWriter`
(following the pattern of `checkpointInfoAccumulator`). Each task appends only
its own active partition's metric updates. The driver then aggregates, sorts,
and formats the top-$K$ instance metrics during `getProgress()`.
### Why are the changes needed?
For queries with high partition counts (e.g. 20,000 shuffle partitions),
each task was shipping ~3 MiB of unused metric updates, totaling ~60 GiB across
tasks. In terminal `ResultStage`s (such as RDD checkpoint materialization or
direct file writes), this causes `spark.driver.maxResultSize` errors, severe
driver heap pressure, and network congestion.
Fixes https://github.com/apache/spark/issues/58394 and
https://issues.apache.org/jira/browse/SPARK-59174.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
- Added unit test in `StateStoreInstanceMetricSuite`: `"SPARK-59174:
StateStoreWriter uses single accumulator for instance metrics"`.
- Ran full `StateStoreInstanceMetricSuite` (all 12 tests passed).
Metric | Before Fix | After Fix
Accumulator objects created | O(N × metrics × stores) — e.g. 20,000–60,000
accumulators | 1 accumulator
Task result payload (per task) | ~3 MiB / task | ~100 bytes / task
### Was this patch authored or co-authored using generative AI tooling?
Yes
Generated-by: Gemini 3.7 Flash
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
