apache / apache/spark

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

Open
#58,394 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Scala
Stars
44k
Forks
29.4k
PR merge metrics
No merged PRs in 30d

Description

## Summary

Stateful operators currently create one state-store instance metric accumulator for every combination of:

```text
state partition
x supported provider metric
x state-store name
```

Every task then receives and registers all of those accumulators, even though it processes only one state partition. On task completion, all external accumulators are included in the task result.

For a job with 20,000 shuffle partitions, RocksDB state, and one default store, this can result in approximately 20,000 state-instance metric updates per task.

## Reproduction

A representative workload is:

```text
spark.sql.shuffle.partitions = 20000
stateful operator using RocksDB state
checkpoint or write the stateful result
```

Observed behavior:

```text
per-task result size: approximately 3 MiB
number of tasks: approximately 20,000
terminal stage: ResultStage
```

The cumulative task-result data can therefore be approximately 60 GiB. The job may fail with `spark.driver.maxResultSize` even though the actual rows are written or checkpointed by executors.

The same issue can occur with direct file writes because the final write stage is submitted through `SparkContext.runJob` and uses `ResultTask`s. Each task returns a commit result plus accumulator updates.

## Root cause

Relevant code paths on `master` include:

- `StateStoreWriter.stateStoreInstanceMetricsWithIds` creates metrics for all state partitions, supported metrics, and store names.
- `StateStoreWriter.setStoreInstanceMetrics` updates the corresponding per-instance `SQLMetric`.
- `AccumulatorV2.readObject` registers all deserialized accumulators with the task.
- `Task.collectAccumulatorUpdates` sends all external accumulators on successful task completion.
- `Executor` serializes those accumulator updates into `DirectTaskResult`.
- `TaskSetManager.canFetchMoreResults` applies `spark.driver.maxResultSize` to `ResultStage`s.

`ignoreIfUnchanged` and `numStateStoreInstanceMetricsToReport` are applied while generating state-operator progress after the driver has received and merged the updates. They do not reduce the task-result payload.

## Why checkpointing and direct writes are affected

RDD checkpoint materialization calls `SparkContext.runJob`, making the checkpointing stage a `ResultStage`.

Direct file writes also use `runJob` and return per-task write results or commit messages to the driver. Those task results contain the write result plus accumulator updates, so state-store metrics can contribute to the cumulative result size.

The `spark.driver.maxResultSize` check is not enforced for upstream `ShuffleMapStage`s, but it is enforced for these terminal `ResultStage`s.

## Impact

This behavior causes:

- inflated task-result serialization and network traffic;
- increased executor and driver deserialization work;
- large `CompletionEvent` payloads on the driver;
- driver heap pressure if result processing falls behind task completion;
- possible `spark.driver.maxResultSize` failures;
- increased per-task metric metadata retained by scheduler or UI structures.

## Scope

This is not specific to `TransformWithState`. It can affect any stateful operator using `StateStoreWriter` and a provider that exposes state-store instance metrics.

Multiple state stores, such as those used by streaming joins, multiply the number of metrics included in each task result.

Contributor guide

Open the contributing guide

Research direction

Start with StateStoreWriter.stateStoreInstanceMetricsWithIds and setStoreInstanceMetrics, then trace accumulator handling through AccumulatorV2.readObject, Task.collectAccumulatorUpdates, Executor, and TaskSetManager.canFetchMoreResults. Reproduce the described high-partition ResultStage workload and inspect task-result sizes. Done should prevent irrelevant state-store instance metrics from inflating terminal task results while preserving the required state-operator progress metrics.

Written by the indexing model from the issue text.

Assessment

Tech stack
scala, sql
Domain
data-engineering, distributed-systems
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.