You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "juliuszsompolski (via GitHub)" <gi...@apache.org> on 2026/06/18 17:30:58 UTC
[PR] [SPARK-56743][SPARK-56773][SQL][CORE][TESTS] Exercise writer-stage retries in DSv2 DML metric tests and fix injection-state cleanup under AQE [spark]
juliuszsompolski opened a new pull request, #56597:
URL: https://github.com/apache/spark/pull/56597
### What changes were proposed in this pull request?
Follow-up to [SPARK-56743](https://issues.apache.org/jira/browse/SPARK-56743) (the
`INJECT_SHUFFLE_FETCH_FAILURES` injection infra and SLAM for DSv2 DML metrics). It strengthens the
DSv2 MERGE/UPDATE retry tests so they actually exercise a writer-stage retry, fixes a
DAGScheduler bug that prevented that under AQE, and closes the test-coverage gap that let the bug
through.
Three parts:
**1. DAGScheduler: don't evict the test-injection state per-stage.**
`cleanupStateForJobAndIndependentStages` removed the per-shuffle injection bookkeeping (the three
`injectShuffleFetchFailures*` maps) whenever a stage was removed. Under AQE each `Exchange` is
materialized as its *own* map-stage job, so that cleanup ran *between* the producer job and the
consumer job and dropped the pending deferred corruption before the consumer was ever submitted -
no `FetchFailed`, no retry. The maps are keyed by the globally-unique (never-reused) `shuffleId`,
only allocated under `Utils.isTesting`, and bounded by the number of shuffles in a test
`SparkContext`, so they are simply retained for its lifetime (restoring the original body of that
cleanup loop).
**2. `MetricsFailureInjectionSuite`: add AQE coverage.** The existing
`INJECT_SHUFFLE_FETCH_FAILURES` tests all run with AQE disabled (the suite mixes in
`DisableAdaptiveExecutionSuite`), so none exercised AQE's per-shuffle-materialization path - the
exact path where the eviction above suppressed the retry. New test
`Three stage metrics block failure injection with AQE` runs the same 3-stage query with
`ADAPTIVE_EXECUTION_ENABLED=true` and asserts the non-leaf stage's raw counter overcounts (a retry
actually fired) while SLAM stays stable. It fails on the pre-fix code and passes after.
**3. DSv2 MERGE/UPDATE retry tests** (`MergeIntoTableSuiteBase`, `UpdateTableSuiteBase`): the
`"metric values are stable across stage retries"` tests now run under the injection and exercise a
real retry. For the metadata MERGE variants - where the writer's `RequiresDistributionAndOrdering`
forces a re-shuffle between `MergeRowsExec` and the writer - `MergeRowsExec` sits in a non-leaf
shuffle map stage, re-runs under the injection, and its raw per-row counters overcount, while the
SLAM-aware `MergeSummary` stays correct; the test asserts both. The `noMetadata` variants skip the
overcount assertion (there `MergeRowsExec` is in the result stage and cannot be re-run by an
upstream injection). UPDATE writer-side metrics live on the result stage and single-count by
design (`ResultStage.findMissingPartitions` only re-runs not-yet-completed partitions), so that
test is regression coverage that retries don't break the SLAM-aware `UpdateSummary`. A `noMetadata`
accessor is added on `RowLevelOperationSuiteBase` so MERGE variants can branch on whether the
writer requires a re-shuffle.
### Why are the changes needed?
The DSv2 DML retry tests added in SPARK-56743 only verified that SLAM values stay correct *given*
retries happen - which is vacuously true even when no retry fires. With the merged injection infra
they did not actually trigger a writer-stage retry under AQE, because the per-stage eviction
dropped the deferred corruption between AQE's per-shuffle jobs. This PR makes the tests demand a
real retry (raw-metric overcount), fixes the infra so that retry actually happens under AQE, and
adds an infra-level AQE test so the regression is caught directly in `MetricsFailureInjectionSuite`
rather than only end-to-end.
### Does this PR introduce _any_ user-facing change?
No. The `DAGScheduler` change only affects test-only state allocated under `Utils.isTesting`; the
rest is test code.
### How was this patch tested?
- New + existing `MetricsFailureInjectionSuite` (13 tests, incl. the new AQE test) pass; the new
AQE test was confirmed to fail on the pre-fix code (eviction present) and pass after.
- `SQLLastAttemptMetricIntegrationSuite` (+ `WithStageRetries` / `WithChecksumMismatch`) and
`SQLLastAttemptMetricPlanShapesSuite` still pass (258 tests) - no regression from the
`DAGScheduler` change.
- All 4 MERGE and 4 UPDATE row-level-operation variants pass; metadata MERGE genuinely overcounts
the raw `MergeRowsExec` accumulator (`numTargetRowsUpdated=6`) while `MergeSummary` reports 2.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code, 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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56743][SPARK-56773][SQL][CORE][TESTS] Exercise writer-stage retries in DSv2 DML metric tests and fix injection-state cleanup under AQE [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #56597:
URL: https://github.com/apache/spark/pull/56597#discussion_r3462632909
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala:
##########
@@ -420,6 +420,63 @@ class MetricsFailureInjectionSuite
}
}
+ test("Three stage metrics block failure injection with AQE") {
+ // Same as the previous test but with AQE enabled. Under AQE each Exchange is materialized
+ // as its own map-stage job, which exercises a different DAGScheduler path than the
+ // AQE-disabled variant: the injection's deferred corruption must survive across those
+ // per-shuffle jobs for the downstream FetchFailed (and thus the producer re-run) to fire.
+ val stage1Metric = SQLMetrics.createMetric(spark.sparkContext, "stage 1 counter")
+ val stage2Metric = SQLMetrics.createMetric(spark.sparkContext, "stage 2 counter")
+ val stage3Metric = SQLMetrics.createMetric(spark.sparkContext, "stage 3 counter")
+ val stage1SLAMetric =
+ SQLLastAttemptMetrics.createMetric(spark.sparkContext, "stage 1 SLAM")
+ val stage2SLAMetric =
+ SQLLastAttemptMetrics.createMetric(spark.sparkContext, "stage 2 SLAM")
+ val stage3SLAMetric =
+ SQLLastAttemptMetrics.createMetric(spark.sparkContext, "stage 3 SLAM")
+
+ withTable("primary_table", "secondary_table") {
+ setUpTestTable("primary_table")
+ setUpTestTable("secondary_table")
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
+ withSparkContextConf(
+ config.Tests.INJECT_SHUFFLE_FETCH_FAILURES.key -> "true") {
+ val stage1MetricsExpr = incrementMetrics(Seq(stage1Metric, stage1SLAMetric))
+ val stage1 = spark.read.table("primary_table")
+ .filter(Column(stage1MetricsExpr))
+ val stage2MetricsExpr = incrementMetrics(Seq(stage2Metric, stage2SLAMetric))
+ val stage2 = stage1.join(
+ spark.read.table("secondary_table"),
+ usingColumn = "id",
+ joinType = "fullOuter")
+ .filter(Column(stage2MetricsExpr))
+ val stage3MetricsExpr = incrementMetrics(Seq(stage3Metric, stage3SLAMetric))
+ val stage3 = stage2
+ .groupBy("primary_table.low_cardinality_col")
+ .count()
+ .filter(Column(stage3MetricsExpr))
+ val finalDf = stage3.as[(Int, Long)]
+ val result = finalDf.collect()
+ assert(result.toMap === (0 until 5).map(v => (v, 300 / 5)).toMap)
+
+ // The non-leaf stage 2 gets its first successful attempt corrupted and re-runs, so its
+ // raw counter overcounts. SLAM reports only the last successful attempt per RDD.
Review Comment:
This comment explains only stage 2, but line 464 asserts `stage1Metric.value > 300` too — under AQE the leaf stage 1 is also a shuffle map stage (it feeds the join shuffle), so it gets corrupted and re-runs as well. The non-AQE sibling's comment covers both; matching it:
```suggestion
// Both the leaf stage 1 and the non-leaf stage 2 get their first successful attempt
// corrupted and re-run, so their raw counters overcount. SLAM reports only the last
// successful attempt per RDD.
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56743][SPARK-56773][SQL][CORE][TESTS] Exercise writer-stage retries in DSv2 DML metric tests and fix injection-state cleanup under AQE [spark]
Posted by "juliuszsompolski (via GitHub)" <gi...@apache.org>.
juliuszsompolski commented on code in PR #56597:
URL: https://github.com/apache/spark/pull/56597#discussion_r3460766021
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala:
##########
@@ -420,6 +420,58 @@ class MetricsFailureInjectionSuite
}
}
+ test("Three stage metrics block failure injection with AQE") {
+ // Same as the previous test but with AQE enabled. Under AQE each Exchange is materialized
Review Comment:
Told my agent to stop being sloppy and readd the assertion :-).
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56743][SPARK-56773][SQL][CORE][TESTS] Exercise writer-stage retries in DSv2 DML metric tests and fix injection-state cleanup under AQE [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #56597:
URL: https://github.com/apache/spark/pull/56597#discussion_r3456635771
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala:
##########
@@ -420,6 +420,58 @@ class MetricsFailureInjectionSuite
}
}
+ test("Three stage metrics block failure injection with AQE") {
+ // Same as the previous test but with AQE enabled. Under AQE each Exchange is materialized
Review Comment:
This test drops four assertions its non-AQE sibling has — `stage3Metric.value === 5` and the three `lastAttemptValueForDataset(finalDf)` checks — yet the comment calls it the same test. If they don't hold under AQE (e.g. dataset lookup through `AdaptiveSparkPlanExec`), say so here; otherwise restore them so the AQE path gets equal coverage.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala:
##########
@@ -342,12 +342,12 @@ abstract class UpdateTableSuiteBase extends RowLevelOperationSuiteBase {
}
test("metric values are stable across stage retries") {
- // Force a shuffle in the UPDATE plan via an IN-subquery (with broadcast disabled), then
- // have the DAGScheduler corrupt the first attempt of every upstream shuffle map stage.
- // Note: the current fetch-failure injection does not retry the writer stage, so this
- // test passes equally well with plain SQLMetric — it only exercises the SLAM-aware
- // read path. Follow-up #55738 will add infra to actually retry the writer stage and
- // exercise the SLAM behavior end-to-end for UPDATE.
+ // INJECT_SHUFFLE_FETCH_FAILURES corrupts the partition-0 task of the first successful
Review Comment:
Nothing here asserts a retry actually fired. With the writer single-counting by design, the test passes even if the injection silently stops retrying — the vacuous-pass gap this PR closes for metadata MERGE. Consider asserting an upstream raw-metric overcount, or note that the new infra-level AQE test is what guards retry-fires. (Same for the MERGE `if (!noMetadata)` path.)
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56743][SPARK-56773][SQL][CORE][TESTS] Exercise writer-stage retries in DSv2 DML metric tests and fix injection-state cleanup under AQE [spark]
Posted by "juliuszsompolski (via GitHub)" <gi...@apache.org>.
juliuszsompolski commented on code in PR #56597:
URL: https://github.com/apache/spark/pull/56597#discussion_r3460799081
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala:
##########
@@ -342,12 +342,12 @@ abstract class UpdateTableSuiteBase extends RowLevelOperationSuiteBase {
}
test("metric values are stable across stage retries") {
- // Force a shuffle in the UPDATE plan via an IN-subquery (with broadcast disabled), then
- // have the DAGScheduler corrupt the first attempt of every upstream shuffle map stage.
- // Note: the current fetch-failure injection does not retry the writer stage, so this
- // test passes equally well with plain SQLMetric — it only exercises the SLAM-aware
- // read path. Follow-up #55738 will add infra to actually retry the writer stage and
- // exercise the SLAM behavior end-to-end for UPDATE.
+ // INJECT_SHUFFLE_FETCH_FAILURES corrupts the partition-0 task of the first successful
Review Comment:
yes, unfortunately it's still impossible to assert for overcounts in result stage metrics, because Spark does not support any retries in this stage. But the infra PR at least added more "interesting" scenarios of restarts, and it added coverage for Merge with metadata.
The new infra-level AQE tests checks that retry fires - it asserts overcount in stage1 and stage2, but not stage3 (again, because there can't be restarts in result stage).
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56743][SPARK-56773][SQL][CORE][TESTS] Exercise writer-stage retries in DSv2 DML metric tests and fix injection-state cleanup under AQE [spark]
Posted by "juliuszsompolski (via GitHub)" <gi...@apache.org>.
juliuszsompolski commented on code in PR #56597:
URL: https://github.com/apache/spark/pull/56597#discussion_r3460752118
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala:
##########
@@ -342,12 +342,12 @@ abstract class UpdateTableSuiteBase extends RowLevelOperationSuiteBase {
}
test("metric values are stable across stage retries") {
- // Force a shuffle in the UPDATE plan via an IN-subquery (with broadcast disabled), then
- // have the DAGScheduler corrupt the first attempt of every upstream shuffle map stage.
- // Note: the current fetch-failure injection does not retry the writer stage, so this
- // test passes equally well with plain SQLMetric — it only exercises the SLAM-aware
- // read path. Follow-up #55738 will add infra to actually retry the writer stage and
- // exercise the SLAM behavior end-to-end for UPDATE.
+ // INJECT_SHUFFLE_FETCH_FAILURES corrupts the partition-0 task of the first successful
Review Comment:
UPDATE writer metrics single-count on the result stage, so this test can't observe an overcount and doesn't independently assert a retry fired. The retry-actually-fires behavior under AQE is guarded by MetricsFailureInjectionSuite's "Three stage metrics block failure injection with AQE" test, so I think there's no point in testing here on an unrelated upstream metric.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56743][SPARK-56773][SQL][CORE][TESTS] Exercise writer-stage retries in DSv2 DML metric tests and fix injection-state cleanup under AQE [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on PR #56597:
URL: https://github.com/apache/spark/pull/56597#issuecomment-4785876389
thanks, merging to master/4.x
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56743][SPARK-56773][SQL][CORE][TESTS] Exercise writer-stage retries in DSv2 DML metric tests and fix injection-state cleanup under AQE [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan closed pull request #56597: [SPARK-56743][SPARK-56773][SQL][CORE][TESTS] Exercise writer-stage retries in DSv2 DML metric tests and fix injection-state cleanup under AQE
URL: https://github.com/apache/spark/pull/56597
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56743][SPARK-56773][SQL][CORE][TESTS] Exercise writer-stage retries in DSv2 DML metric tests and fix injection-state cleanup under AQE [spark]
Posted by "juliuszsompolski (via GitHub)" <gi...@apache.org>.
juliuszsompolski commented on code in PR #56597:
URL: https://github.com/apache/spark/pull/56597#discussion_r3460752118
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala:
##########
@@ -342,12 +342,12 @@ abstract class UpdateTableSuiteBase extends RowLevelOperationSuiteBase {
}
test("metric values are stable across stage retries") {
- // Force a shuffle in the UPDATE plan via an IN-subquery (with broadcast disabled), then
- // have the DAGScheduler corrupt the first attempt of every upstream shuffle map stage.
- // Note: the current fetch-failure injection does not retry the writer stage, so this
- // test passes equally well with plain SQLMetric — it only exercises the SLAM-aware
- // read path. Follow-up #55738 will add infra to actually retry the writer stage and
- // exercise the SLAM behavior end-to-end for UPDATE.
+ // INJECT_SHUFFLE_FETCH_FAILURES corrupts the partition-0 task of the first successful
Review Comment:
UPDATE writer metrics single-count on the result stage, so this test can't observe an overcount and doesn't independently assert a retry fired. The retry-actually-fires behavior under AQE is guarded by MetricsFailureInjectionSuite's "Three stage metrics block failure injection with AQE" test, so I think there's no point in testing here on an unrelated upstream metric.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56743][SPARK-56773][SQL][CORE][TESTS] Exercise writer-stage retries in DSv2 DML metric tests and fix injection-state cleanup under AQE [spark]
Posted by "juliuszsompolski (via GitHub)" <gi...@apache.org>.
juliuszsompolski commented on PR #56597:
URL: https://github.com/apache/spark/pull/56597#issuecomment-4769704622
@cloud-fan here's a PR to use the infra from https://github.com/apache/spark/pull/55738 for additional metrics tests, and also a small fix to the infra - it did not work in AQE, because AQE submits each stage separately.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org