szehon-ho opened a new pull request, #57323:
URL: https://github.com/apache/spark/pull/57323
### What changes were proposed in this pull request?
Observation metrics (`Dataset.observe` / `CollectMetrics`, surfaced via
`CollectMetricsExec`) are aggregated on the driver through an
`AggregatingAccumulator`. When a task, stage, or RDD is recomputed
(speculation, fetch-failure retries, AQE stage cancellation/replanning), the
accumulator receives updates from *every* attempt and merges them together, so
the reported metric can reflect superseded attempts (e.g. double-counted rows
or an over-sized `collect_list`).
This PR makes observation metrics report only the values produced by the
last successful attempt:
- Adds `LastAttemptAggregatingAccumulator`, a subclass of
`AggregatingAccumulator` that mixes in the core `LastAttemptAccumulator` trait.
It tracks per-RDD / per-stage-attempt / per-partition contributions and can
return the aggregate for the highest RDD id (`lastAttemptValueForHighestRDDId`).
- `partialMergeVal` snapshots the (mutable) aggregate buffer via
`buffer.copy()`; `partialMerge` merges a snapshot back in.
- `isMergeable` guards merges by comparing an eagerly-captured
`outputSchema` (captured before task serialization drops the transient result
expressions on task-side copies).
- `add` mirrors direct driver-side updates into the last-attempt state.
- `AggregatingAccumulator` gains a `lastAttempt(...)` factory (alongside
refactoring to allow subclassing).
- `CollectMetricsExec` now creates the last-attempt accumulator,
initializes/resets its last-attempt state, and reads the collected metrics from
`lastAttemptValueForHighestRDDId()` (falling back to `value`).
### Why are the changes needed?
Observation metrics are supposed to describe the data that actually flowed
through the query. Under retries/replans the existing accumulator over-counts
by merging superseded attempts, which is incorrect and non-deterministic. This
is especially visible for non-idempotent aggregate buffers (e.g.
`collect_list`), and for AQE, where a replanned stage produces a brand-new RDD
whose earlier attempt should not be counted.
### Does this PR introduce _any_ user-facing change?
Yes. When a stage/RDD backing an observed operator is recomputed, the
observed metric values now reflect only the last successful attempt instead of
a merge across all attempts. Without recomputation, the reported values are
unchanged.
### How was this patch tested?
- `AggregatingAccumulatorSuite`: existing cases parameterized over
`lastAttempt = {true, false}`, plus new last-attempt-specific cases
(schema-mismatch invalidation, mixed driver/task updates, highest-RDD-id
selection, skipped uncomputed partitions, latest stage attempt per partition).
- `MetricsFailureInjectionSuite`: new end-to-end tests under
shuffle-fetch-failure and AQE-cancellation injection that observe a mix of
declarative / imperative / typed-imperative aggregates.
`size(collect_list(...)) == 300` deterministically proves last-attempt
selection for a non-idempotent typed buffer. Also adds coverage for observation
on a non-leaf (post-shuffle) stage, multiple observation points in one query,
and last-attempt state reset across reuse of the same `Dataset`.
- `DataFrameCallbackSuite`: additional observation coverage.
All of `MetricsFailureInjectionSuite` (20 tests) passes locally.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Cursor (Opus 4.8)
--
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]
