Skip to content

[SPARK-59174][SS] State-store instance metrics are shipped in every task result - #58476

Open
Diveyam-Mishra wants to merge 1 commit into
apache:masterfrom
Diveyam-Mishra:SPARK-59174
Open

[SPARK-59174][SS] State-store instance metrics are shipped in every task result#58476
Diveyam-Mishra wants to merge 1 commit into
apache:masterfrom
Diveyam-Mishra:SPARK-59174

Conversation

@Diveyam-Mishra

Copy link
Copy Markdown

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 ResultStages (such as RDD checkpoint materialization or direct file writes), this causes spark.driver.maxResultSize errors, severe driver heap pressure, and network congestion.

Fixes #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

@Diveyam-Mishra

Copy link
Copy Markdown
Author

Hey its my first time contributing not knowing how to ask for a review so tagging let me know if there is a different way. (I Read the contributing guide )
@dongjoon-hyun

*/
lazy val instanceMetrics: Map[StateStoreInstanceMetric, SQLMetric] =
stateStoreInstanceMetrics
val instanceMetricsAccumulator: CollectionAccumulator[(StateStoreInstanceMetric, Long)] = {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we use or extend PartitionKeyedAccumulator here? Unlike CollectionAccumulator, it bounds driver-side state by partition when retries or speculative attempts produce duplicate updates. We would need to preserve StateStoreInstanceMetric.combine semantics rather than its default last-write-wins behavior, but that may avoid retaining every successful attempt’s raw metrics on the driver.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[SQL] State-store instance metrics are shipped in every task result

2 participants