You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2026/05/06 04:58:24 UTC
(spark) branch branch-4.x updated: [SPARK-56647][SQL] Optimize storage of SparkSQL Last Attempt Metrics
This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new aed6ac3be829 [SPARK-56647][SQL] Optimize storage of SparkSQL Last Attempt Metrics
aed6ac3be829 is described below
commit aed6ac3be829115fc293cb739e6667b6dc8f4289
Author: Juliusz Sompolski <25...@users.noreply.github.com>
AuthorDate: Wed May 6 12:57:44 2026 +0800
[SPARK-56647][SQL] Optimize storage of SparkSQL Last Attempt Metrics
### What changes were proposed in this pull request?
This PR optimizes the storage of per-partition tracking state in `LastAttemptRDDVals`, which backs `SQLLastAttemptMetric` (SLAM) introduced in [SPARK-56509](https://issues.apache.org/jira/browse/SPARK-56509).
In the original design, every `LastAttemptRDDVals` allocated three `Array[Int]` of size `numPartitions` to track `(stageId, stageAttemptId, taskAttemptNumber)` for each partition. In a typical execution without retries every entry of those arrays held the same value, and the three arrays together took ~12·N bytes per RDD that uses a SLAM accumulator.
The three int arrays are replaced by a **common attempt** plus three **lazily-allocated per-component override arrays**:
- `commonStageId`, `commonStageAttemptId`, `commonTaskAttemptNumber`: 3 ints set on the first `update()` from the first task that completes. Partitions whose attempt matches the common values carry no per-partition stage/attempt state at all.
- `computedBitmap` (`Array[Long]`): one bit per partition, packed into longs. Tracks which partitions have been computed; replaces the previous `EMPTY_ID` sentinel that was stored in all three int arrays.
- `overrideStageIds` / `overrideStageAttemptIds` / `overrideTaskAttemptNumbers` (`Array[Int]`, sized `numPartitions`): each is **null** until some partition's value for *that* component diverges from the common. Once allocated, an entry equal to `EMPTY_ID` (-1) means "fall back to the common value", any other value is the per-partition override. The three arrays are independent — a retry that only changes `taskAttemptNumber` allocates only `overrideTaskAttemptNumbers`, leaving the oth [...]
Once an override array exists, every subsequent update for that partition writes its value into the slot — even when the value happens to equal the common; `lookupComponent` returns it correctly either way. So a slot that still holds `EMPTY_ID` after the array is allocated only happens when the slot was never written (the partition's value matched the common when the array was first allocated for a *different* partition's divergence).
`partialValueAt`, `isEmptyAt`, and the merge logic in `mergeLastAttempt` preserve their existing semantics, so `DAGScheduler.updateAccumulators` and the SLAM merge path are unchanged. `toString` reconstructs the original per-partition `stageIds` / `stageAttemptIds` / `taskAttemptNumbers` view for debug logs by walking `partialValueAt` for every partition; the compact internal representation isn't separately printed because the reconstructed view shows the same effective values.
The single-writer/multi-reader concurrency contract is preserved. Each override array reference is `volatile`. On the first divergence for a component the writer allocates the array, fills it with `EMPTY_ID` via `Array.fill`, writes the override entry, and only then assigns the field — so a concurrent reader either sees null (and falls back to the common value) or sees a fully initialized array. Subsequent in-place element writes are plain ints, eventually consistent, matching the loo [...]
Memory comparison for a 200-partition RDD. The `partitionPartialVals` array of size N (e.g. `Array[Long](200)` ≈ 1.6 KB for `SQLLastAttemptMetric`) is required to hold the per-partition partial values and is unchanged; the numbers below are the additional attempt-tracking overhead per RDD on top of that array:
| State | Attempt-tracking overhead |
|-----------------------------------------|---------------------------|
| Before | ~2.4 KB (three `Array[Int](200)` + headers) |
| After, no retries | ~50 B (3 common ints + ~25 B computed bitmap; all three override arrays `null`) |
| After, pure stage-attempt retry | ~850 B (above + one `Array[Int](200)` for `overrideStageAttemptIds`) |
| After, mid-stage executor-loss retry | ~850 B (above + one `Array[Int](200)` for `overrideTaskAttemptNumbers`) |
| After, full cross-stage retry (worst) | ~2.4 KB (above + all three `Array[Int](200)`) |
The worst case matches the original design's footprint; the typical no-retry case improves by ~50x and partial retries by ~3x, with the size paid only when the corresponding component actually changes.
### Why are the changes needed?
Follow-up to the SLAM PR ([SPARK-56509](https://issues.apache.org/jira/browse/SPARK-56509) / [#55371](https://github.com/apache/spark/pull/55371)), addressing a memory-efficiency concern raised in review: per-partition storage of `(stageId, stageAttemptId, taskAttemptNumber)` is redundant in the common case, and even under retries the components diverge independently — usually only one of the three actually changes.
If SLAM accumulators see broader use — for example as the basis for accurate Spark UI metrics under retries — many `SQLMetric` instances throughout a query plan would each carry a `LastAttemptRDDVals` per RDD that updated them. Reducing the per-RDD overhead from 3·N ints to a handful of ints + a small bitmap (with each retry component paid for only on actual divergence) is a meaningful win at scale.
### Does this PR introduce _any_ user-facing change?
No. This is an internal storage optimization for `LastAttemptRDDVals`. The public `LastAttemptAccumulator` / `SQLLastAttemptMetric` APIs, the values they return, and their behavior under stage retries are unchanged.
### How was this patch tested?
- All existing SLAM test suites pass unchanged, exercising the same code paths as before:
- `SQLLastAttemptMetricUnitSuite`
- `SQLLastAttemptMetricIntegrationSuite`
- `SQLLastAttemptMetricIntegrationSuiteWithStageRetries`
- `SQLLastAttemptMetricPlanShapesSuite`
- `MetricsFailureInjectionSuite`
175 tests, all green.
- New unit test `SQLLastAttemptMetricUnitSuite."compact storage: per-component override arrays allocated only when component diverges"` reflects on the internal fields to verify:
- With no retries: all three override arrays stay null; the common attempt fields are populated.
- Pure stage-attempt retry of a partition: only `overrideStageAttemptIds` is allocated; the other two stay null.
- Mid-stage retry (different `taskAttemptNumber`, same stage): only `overrideTaskAttemptNumbers` is allocated.
- Cross-stage retry: `overrideStageIds` is allocated alongside the others.
- When a later update brings a previously-overridden component back to the common value, the slot holds the new value (which equals the common); `lookupComponent` returns it correctly either way.
- The aggregated metric values across all of the above transitions are correct.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (claude-opus-4-7)
Closes #55585 from juliuszsompolski/slam-improve.
Lead-authored-by: Juliusz Sompolski <25...@users.noreply.github.com>
Co-authored-by: Juliusz Sompolski <Juliusz Sompolski>
Signed-off-by: Wenchen Fan <we...@databricks.com>
(cherry picked from commit 5d2c576d5d427ec70c9ff4a40fd6d1f153195732)
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../apache/spark/util/LastAttemptAccumulator.scala | 183 +++++++++++++++++----
.../metric/SQLLastAttemptMetricUnitSuite.scala | 102 ++++++++++++
2 files changed, 252 insertions(+), 33 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/util/LastAttemptAccumulator.scala b/core/src/main/scala/org/apache/spark/util/LastAttemptAccumulator.scala
index cd17476a0cc4..a262750a8cd3 100644
--- a/core/src/main/scala/org/apache/spark/util/LastAttemptAccumulator.scala
+++ b/core/src/main/scala/org/apache/spark/util/LastAttemptAccumulator.scala
@@ -126,8 +126,7 @@ import org.apache.spark.scheduler.TaskInfo
private class LastAttemptRDDVals[@specialized T](
val rddId: Int,
val rddScopeId: Option[String],
- // Arrays of partial metric values, and the corresponding stage, stage attempt and task attempt,
- // with index representing RDD partition id.
+ // Array of partial metric values, indexed by RDD partition id.
// Metric updates to a given RDD partition can come from different stageAttempts if a retry
// happens while a Job with the Stage is running (a downstream Stage within a Job detects
// missing blocks and triggers recompute), or from different Stages, if a retry happens later
@@ -144,18 +143,19 @@ private class LastAttemptRDDVals[@specialized T](
// in take/limit), or AQE task coalescing may be visible as an update of the partition id of
// the first partition of the coalesced range. AQE guarantees that if these are retried, they
// will be coalesced in the same ranges, so update the same values.
- // Not computed partitions should have EMPTY_ID in all the Int Arrays.
+ // Whether a partition has been computed is tracked by [[computedBitmap]] below; the value at
+ // its slot in [[partitionPartialVals]] is undefined (typically the zero of T) for uncomputed
+ // partitions.
//
// Arrays of primitive types are more memory efficient than an array of objects due to
// references, object headers and paddings overheads.
// The `@specialized` annotation should make scala specialize it to use primitive array instead
// of boxed objects.
- val partitionPartialVals: Array[T],
- val stageIds: Array[Int],
- val stageAttemptIds: Array[Int],
- val taskAttemptNumbers: Array[Int])
+ val partitionPartialVals: Array[T])
{
+ import LastAttemptRDDVals.EMPTY_ID
+
// In a case of repeated execution of the same QueryExecution and reuse of the SparkPlan
// (for example multiple `collect()` on the same Dataset), a new RDD may be executed in the same
// RDDOperationScope for the new execution. Hence, we can have multiple RDDs with the same
@@ -171,67 +171,184 @@ private class LastAttemptRDDVals[@specialized T](
// however should not happen in practice and would likely produce other unexpected effects.
var lastSqlExecutionId: Option[Long] = None
- def numPartitions: Int = stageIds.length
+ // Common (stageId, stageAttemptId, taskAttemptNumber) shared by the majority of computed
+ // partitions. In the common case (no stage retries), every computed partition has the same
+ // attempt tuple, so we store it once at the RDD level instead of allocating three N-sized int
+ // arrays. The values are set on the first update and never changed; partitions whose attempt
+ // differs are recorded in the override arrays below.
+ // EMPTY_ID until the first update.
+ private var commonStageId: Int = EMPTY_ID
+ private var commonStageAttemptId: Int = EMPTY_ID
+ private var commonTaskAttemptNumber: Int = EMPTY_ID
+
+ // Bitmap of partitions that have been computed, one bit per partition packed into longs.
+ // A bit is set when a partition receives its first update; a partition with a clear bit has not
+ // been computed (e.g. early stop in take/limit, AQE task coalescing).
+ // Reads of an individual long are atomic on 64-bit JVMs, matching the loose concurrency
+ // semantics of the original per-partition int arrays.
+ private val computedBitmap: Array[Long] =
+ new Array[Long]((partitionPartialVals.length + 63) >>> 6)
+
+ // Per-partition override arrays for each component of the attempt tuple. Each is allocated
+ // lazily and independently the first time some partition's value for that component diverges
+ // from the common value; until then the field is null and no per-partition state is kept for
+ // that component. Once allocated, an array is sized [[numPartitions]]: entries equal to
+ // EMPTY_ID mean "match the common value" and any other value is the per-partition override.
+ // This way:
+ // - RDDs without retries pay zero per-partition allocations (all three fields stay null).
+ // - A pure stage retry (new stageAttemptId, same stageId, taskAttemptNumber resets to 0)
+ // allocates only [[overrideStageAttemptIds]].
+ // - A mid-stage retry (executor lost, some tasks restart with a higher taskAttemptNumber)
+ // allocates only [[overrideTaskAttemptNumbers]].
+ // - Whole-stage cross-Stage retry (new stageId) allocates [[overrideStageIds]] too.
+ //
+ // Concurrency: update() is called only from the DAGScheduler scheduler loop. Some readers of
+ // the state can run concurrently (e.g. logAccumulatorState formatting). The fields are
+ // declared @volatile, and the new array is fully populated before the field is assigned, so a
+ // reader either sees null (use common) or sees an array whose Array.fill initialization is
+ // visible. In-place element writes for subsequent overrides are plain ints; readers may see
+ // them eventually, matching the loose semantics of the original per-partition int arrays.
+ @volatile private var overrideStageIds: Array[Int] = null
+ @volatile private var overrideStageAttemptIds: Array[Int] = null
+ @volatile private var overrideTaskAttemptNumbers: Array[Int] = null
+
+ def numPartitions: Int = partitionPartialVals.length
def isEmptyAt(partitionId: Int): Boolean = {
- if (stageIds(partitionId) == LastAttemptRDDVals.EMPTY_ID) {
- assert(stageAttemptIds(partitionId) == LastAttemptRDDVals.EMPTY_ID)
- assert(taskAttemptNumbers(partitionId) == LastAttemptRDDVals.EMPTY_ID)
- true
+ val word = computedBitmap(partitionId >>> 6)
+ ((word >>> (partitionId & 63)) & 1L) == 0L
+ }
+
+ private def setComputedBit(partitionId: Int): Unit = {
+ val idx = partitionId >>> 6
+ computedBitmap(idx) = computedBitmap(idx) | (1L << (partitionId & 63))
+ }
+
+ /**
+ * Records a new value for one component (stageId / stageAttemptId / taskAttemptNumber) of the
+ * attempt tuple at `partitionId`, allocating the override array on first divergence. Returns
+ * the array reference the caller should write back to the @volatile field - either a freshly
+ * allocated and populated array (first override for this component) or the existing array
+ * after an in-place update. Once the array exists, the value is always written, even when it
+ * matches the common value - lookupComponent returns it correctly either way.
+ */
+ private def setOverrideComponent(
+ array: Array[Int],
+ partitionId: Int,
+ value: Int,
+ common: Int): Array[Int] = {
+ if (array == null) {
+ if (value == common) null
+ else {
+ val newArr = Array.fill(partitionPartialVals.length)(EMPTY_ID)
+ newArr(partitionId) = value
+ newArr
+ }
} else {
- false
+ array(partitionId) = value
+ array
+ }
+ }
+
+ /** Reads one component's value at `partitionId`, falling back to `common` when the override
+ * array is null or the entry is still EMPTY_ID (the slot was either not yet written, or was
+ * initialized to EMPTY_ID and never overwritten because the partition's value matched the
+ * common when the array was first allocated for a different partition). */
+ private def lookupComponent(array: Array[Int], partitionId: Int, common: Int): Int = {
+ if (array == null) common
+ else {
+ val v = array(partitionId)
+ if (v == EMPTY_ID) common else v
}
}
def update(partialValue: AccumulatorPartialVal[T]): Unit = {
- partitionPartialVals(partialValue.rddPartitionId) = partialValue.partialMergeVal
- stageIds(partialValue.rddPartitionId) = partialValue.stageId
- stageAttemptIds(partialValue.rddPartitionId) = partialValue.stageAttemptId
- taskAttemptNumbers(partialValue.rddPartitionId) = partialValue.taskAttemptNumber
+ val partId = partialValue.rddPartitionId
+ if (commonStageId == EMPTY_ID) {
+ commonStageId = partialValue.stageId
+ commonStageAttemptId = partialValue.stageAttemptId
+ commonTaskAttemptNumber = partialValue.taskAttemptNumber
+ }
+ partitionPartialVals(partId) = partialValue.partialMergeVal
+ setComputedBit(partId)
+ overrideStageIds = setOverrideComponent(
+ overrideStageIds, partId, partialValue.stageId, commonStageId)
+ overrideStageAttemptIds = setOverrideComponent(
+ overrideStageAttemptIds, partId, partialValue.stageAttemptId, commonStageAttemptId)
+ overrideTaskAttemptNumbers = setOverrideComponent(
+ overrideTaskAttemptNumbers, partId, partialValue.taskAttemptNumber, commonTaskAttemptNumber)
lastSqlExecutionId = partialValue.sqlExecutionId
}
def partialValueAt(partId: Int): AccumulatorPartialVal[T] = {
+ var sId = EMPTY_ID
+ var saId = EMPTY_ID
+ var tan = EMPTY_ID
+ if (!isEmptyAt(partId)) {
+ sId = lookupComponent(overrideStageIds, partId, commonStageId)
+ saId = lookupComponent(overrideStageAttemptIds, partId, commonStageAttemptId)
+ tan = lookupComponent(overrideTaskAttemptNumbers, partId, commonTaskAttemptNumber)
+ }
AccumulatorPartialVal(
partialMergeVal = partitionPartialVals(partId),
rddId = rddId,
rddPartitionId = partId,
- rddNumPartitions = stageIds.length,
+ rddNumPartitions = partitionPartialVals.length,
rddScopeId = rddScopeId,
- stageId = stageIds(partId),
- stageAttemptId = stageAttemptIds(partId),
- taskAttemptNumber = taskAttemptNumbers(partId),
+ stageId = sId,
+ stageAttemptId = saId,
+ taskAttemptNumber = tan,
sqlExecutionId = lastSqlExecutionId)
}
override def toString: String = {
+ val n = numPartitions
+ val partVals = new StringBuilder("[")
+ val sIds = new StringBuilder("[")
+ val saIds = new StringBuilder("[")
+ val tans = new StringBuilder("[")
+ var i = 0
+ while (i < n) {
+ if (i > 0) {
+ partVals.append(',')
+ sIds.append(',')
+ saIds.append(',')
+ tans.append(',')
+ }
+ partVals.append(partitionPartialVals(i))
+ val pv = partialValueAt(i)
+ sIds.append(pv.stageId)
+ saIds.append(pv.stageAttemptId)
+ tans.append(pv.taskAttemptNumber)
+ i += 1
+ }
+ partVals.append(']')
+ sIds.append(']')
+ saIds.append(']')
+ tans.append(']')
s"""LastAttemptVal(
| rddId=$rddId,
| rddScopeId=$rddScopeId,
| lastSqlExecutionId=$lastSqlExecutionId,
- | partitionPartialVals=${partitionPartialVals.mkString("[", ",", "]")},
- | stageIds=${stageIds.mkString("[", ",", "]")},
- | stageAttemptIds=${stageAttemptIds.mkString("[", ",", "]")},
- | taskAttemptNumbers=${taskAttemptNumbers.mkString("[", ",", "]")}
+ | partitionPartialVals=$partVals,
+ | stageIds=$sIds,
+ | stageAttemptIds=$saIds,
+ | taskAttemptNumbers=$tans
|)""".stripMargin
}
}
private object LastAttemptRDDVals {
- // EMPTY_ID in stageId means that the partition was not computed.
+ // EMPTY_ID indicates "no attempt recorded": used as the initial value of the common
+ // (stageId, stageAttemptId, taskAttemptNumber) before any update, and as the value returned
+ // by partialValueAt for partitions that have not been computed.
val EMPTY_ID: Int = -1
def apply[@specialized T](
rddId: Int,
rddScopeId: Option[String],
numPartitions: Int)(implicit ct: ClassTag[T]): LastAttemptRDDVals[T] = {
- new LastAttemptRDDVals[T](
- rddId,
- rddScopeId,
- new Array[T](numPartitions),
- Array.fill(numPartitions)(LastAttemptRDDVals.EMPTY_ID),
- Array.fill(numPartitions)(LastAttemptRDDVals.EMPTY_ID),
- Array.fill(numPartitions)(LastAttemptRDDVals.EMPTY_ID))
+ new LastAttemptRDDVals[T](rddId, rddScopeId, new Array[T](numPartitions))
}
def createFromFirstUpdate[@specialized T](
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricUnitSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricUnitSuite.scala
index f3a696f81450..45dff0fe4aa9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricUnitSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricUnitSuite.scala
@@ -185,4 +185,106 @@ class SQLLastAttemptMetricUnitSuite extends SparkFunSuite with SharedSparkContex
assert(slam.lastAttemptValueForRDDId(2) === Some(42)) // new RDD added
assert(slam.lastAttemptValueForAllRDDs() === Some(60))
}
+
+ test("compact storage: per-component override arrays allocated only when component diverges") {
+ val slam = SQLLastAttemptMetrics.createMetric(sc, "test SLAM")
+ val acc = slam.copy()
+
+ val rddsMap = lastAttemptRddsMapField.get(slam)
+ val mapGetMethod = rddsMap.getClass.getMethod("get", classOf[Object])
+
+ // Establish a common attempt across all 5 mock partitions.
+ for (partId <- 0 until 5) {
+ setMockAttempt(rddId = 1, partitionId = partId)
+ acc.set(10)
+ slam.mergeLastAttempt(acc, mockRdd, mockTaskInfo, 10, 10, mockProperties)
+ }
+ val rddVals = mapGetMethod.invoke(rddsMap, Int.box(1)).asInstanceOf[Option[Object]].get
+ val rddValsClass = rddVals.getClass
+ // @specialized var fields are emitted with `org$apache$spark$util$LastAttemptRDDVals$$`
+ // prefix so specialized subclasses can override them.
+ val fieldPrefix = "org$apache$spark$util$LastAttemptRDDVals$$"
+ val commonStageIdFld = rddValsClass.getDeclaredField(fieldPrefix + "commonStageId")
+ commonStageIdFld.setAccessible(true)
+ val commonStageAttemptIdFld =
+ rddValsClass.getDeclaredField(fieldPrefix + "commonStageAttemptId")
+ commonStageAttemptIdFld.setAccessible(true)
+ val commonTaskAttemptNumberFld =
+ rddValsClass.getDeclaredField(fieldPrefix + "commonTaskAttemptNumber")
+ commonTaskAttemptNumberFld.setAccessible(true)
+ val overrideStageIdsFld = rddValsClass.getDeclaredField(fieldPrefix + "overrideStageIds")
+ overrideStageIdsFld.setAccessible(true)
+ val overrideStageAttemptIdsFld =
+ rddValsClass.getDeclaredField(fieldPrefix + "overrideStageAttemptIds")
+ overrideStageAttemptIdsFld.setAccessible(true)
+ val overrideTaskAttemptNumbersFld =
+ rddValsClass.getDeclaredField(fieldPrefix + "overrideTaskAttemptNumbers")
+ overrideTaskAttemptNumbersFld.setAccessible(true)
+
+ // No retries: common is set, none of the override arrays are allocated.
+ assert(commonStageIdFld.getInt(rddVals) === 10)
+ assert(commonStageAttemptIdFld.getInt(rddVals) === 10)
+ assert(commonTaskAttemptNumberFld.getInt(rddVals) === 0)
+ assert(overrideStageIdsFld.get(rddVals) === null)
+ assert(overrideStageAttemptIdsFld.get(rddVals) === null)
+ assert(overrideTaskAttemptNumbersFld.get(rddVals) === null)
+ assert(slam.lastAttemptValueForRDDId(1) === Some(50))
+
+ // Pure stage-attempt retry of partition 0: only stageAttemptId diverges from the common.
+ setMockAttempt(rddId = 1, partitionId = 0)
+ acc.set(20)
+ slam.mergeLastAttempt(acc, mockRdd, mockTaskInfo, 10, 11, mockProperties)
+ assert(overrideStageIdsFld.get(rddVals) === null,
+ "stageId still matches common; its override array should not be allocated")
+ assert(overrideStageAttemptIdsFld.get(rddVals) != null,
+ "stageAttemptId diverged; its override array should be allocated")
+ assert(overrideTaskAttemptNumbersFld.get(rddVals) === null,
+ "taskAttemptNumber still matches common; its override array should not be allocated")
+ val saIds1 = overrideStageAttemptIdsFld.get(rddVals).asInstanceOf[Array[Int]]
+ assert(saIds1.length === 5)
+ assert(saIds1(0) === 11)
+ assert(saIds1(1) === -1, "Untouched partitions should hold EMPTY_ID (= -1) sentinel")
+ assert(slam.lastAttemptValueForRDDId(1) === Some(60)) // 20 + 10*4
+
+ // Mid-stage retry of partition 1: only taskAttemptNumber diverges.
+ when(mockTaskInfo.attemptNumber).thenReturn(1)
+ when(mockRdd.id).thenReturn(1)
+ when(mockTaskInfo.partitionId).thenReturn(1)
+ acc.set(15)
+ slam.mergeLastAttempt(acc, mockRdd, mockTaskInfo, 10, 10, mockProperties)
+ assert(overrideStageIdsFld.get(rddVals) === null)
+ assert(overrideTaskAttemptNumbersFld.get(rddVals) != null,
+ "taskAttemptNumber diverged; its override array should be allocated")
+ val tans2 = overrideTaskAttemptNumbersFld.get(rddVals).asInstanceOf[Array[Int]]
+ assert(tans2(1) === 1)
+ assert(tans2(0) === -1)
+ assert(slam.lastAttemptValueForRDDId(1) === Some(65)) // 20 + 15 + 10*3
+
+ // Cross-stage retry of partition 2 (new stageId). Now stageId also diverges.
+ when(mockTaskInfo.attemptNumber).thenReturn(0)
+ when(mockTaskInfo.partitionId).thenReturn(2)
+ acc.set(30)
+ slam.mergeLastAttempt(acc, mockRdd, mockTaskInfo, 11, 0, mockProperties)
+ assert(overrideStageIdsFld.get(rddVals) != null,
+ "stageId now diverges; its override array should be allocated")
+ val sIds3 = overrideStageIdsFld.get(rddVals).asInstanceOf[Array[Int]]
+ assert(sIds3(2) === 11)
+ assert(sIds3(0) === -1)
+ assert(slam.lastAttemptValueForRDDId(1) === Some(85)) // 20 + 15 + 30 + 10*2
+
+ // Re-update partition 0 with a value that brings stageAttemptId back to common (10) while
+ // diverging stageId (12). Once an override array exists, every update writes its value into
+ // the slot - even when the value equals the common - so partition 0's stageAttemptId entry
+ // becomes 10 (rather than being cleared to EMPTY_ID).
+ when(mockTaskInfo.partitionId).thenReturn(0)
+ acc.set(40)
+ slam.mergeLastAttempt(acc, mockRdd, mockTaskInfo, 12, 10, mockProperties)
+ val saIds4 = overrideStageAttemptIdsFld.get(rddVals).asInstanceOf[Array[Int]]
+ assert(saIds4(0) === 10,
+ "Partition 0's stageAttemptId entry should hold the new value (which happens to equal " +
+ "the common), not EMPTY_ID")
+ val sIds4 = overrideStageIdsFld.get(rddVals).asInstanceOf[Array[Int]]
+ assert(sIds4(0) === 12)
+ assert(slam.lastAttemptValueForRDDId(1) === Some(105)) // 40 + 15 + 30 + 10*2
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org