This is an automated email from the ASF dual-hosted git repository.
szehon-ho pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 911f372d388d [SPARK-58183][SQL] Report observation metrics from the
last successful task attempt
911f372d388d is described below
commit 911f372d388d259aa120a8672fc5a6cd8effd74d
Author: Szehon Ho <[email protected]>
AuthorDate: Thu Jul 23 17:11:08 2026 -0700
[SPARK-58183][SQL] Report observation metrics from the last successful task
attempt
### 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)
Closes #57323 from szehon-ho/SPARK-collect-metrics-last-attempt.
Authored-by: Szehon Ho <[email protected]>
Signed-off-by: Szehon Ho <[email protected]>
---
docs/sql-migration-guide.md | 1 +
.../org/apache/spark/sql/internal/SQLConf.scala | 10 +
.../sql/execution/AggregatingAccumulator.scala | 138 ++++--
.../spark/sql/execution/CollectMetricsExec.scala | 23 +-
.../LastAttemptAggregatingAccumulator.scala | 91 ++++
.../execution/AggregatingAccumulatorSuite.scala | 512 ++++++++++++++++-----
.../metric/MetricsFailureInjectionSuite.scala | 274 ++++++++++-
.../spark/sql/util/DataFrameCallbackSuite.scala | 37 ++
8 files changed, 924 insertions(+), 162 deletions(-)
diff --git a/docs/sql-migration-guide.md b/docs/sql-migration-guide.md
index 9b4db592b0b3..c0cc988f113f 100644
--- a/docs/sql-migration-guide.md
+++ b/docs/sql-migration-guide.md
@@ -26,6 +26,7 @@ license: |
- Since Spark 4.3, [ASOF JOIN](sql-ref-syntax-qry-select-asof-join.html) is
available as an opt-in SQL feature gated by `spark.sql.join.asofJoin.enabled`
(default `false`). When disabled, `ASOF JOIN` fails at parse time with
`UNSUPPORTED_FEATURE.ASOF_JOIN`.
- Since Spark 4.3, zero-length files are skipped during Parquet schema
inference instead of failing with a `FAILED_READ_FILE.CANNOT_READ_FILE_FOOTER`
error.
+- Since Spark 4.3, metrics produced by `Dataset.observe` include values from
only the last successful task attempts instead of aggregating values from all
attempts. To restore the previous behavior, set
`spark.sql.legacy.observeMetricsAggregateAllAttempts` to `true`.
- Since Spark 4.3, the configuration key
`spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabled` has
been renamed to
`spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` to
reflect that it now applies to storage-partitioned joins, aggregates, and
windows. The old key continues to work as an alias.
- Since Spark 4.3, the Spark Thrift Server rejects setting JVM system
properties through the `set:system:` session configuration overlay (for
example, in a JDBC connection string). To restore the previous behavior, set
`spark.sql.legacy.hive.thriftServer.allowSettingSystemProperties` to `true`.
- Since Spark 4.3, the adaptive execution rule
`org.apache.spark.sql.execution.adaptive.DynamicJoinSelection` has been renamed
to `DemoteBroadcastHashJoin`, which now only demotes broadcast hash joins
(emitting `NO_BROADCAST_HASH`). Its selection of shuffled hash join over sort
merge join has moved to a new physical rule gated by
`spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled` (default
`true`). If you previously disabled the shuffled-hash-join preference by
listing `o [...]
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 5af33da7ae8e..edb99889f3b4 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -6222,6 +6222,16 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val LEGACY_OBSERVE_METRICS_AGGREGATE_ALL_ATTEMPTS =
+ buildConf("spark.sql.legacy.observeMetricsAggregateAllAttempts")
+ .internal()
+ .doc("When true, metrics produced by `Dataset.observe` aggregate values
from all task " +
+ "attempts. When false, only values from the last successful task
attempts are used.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .booleanConf
+ .createWithDefault(false)
+
val MAX_TO_STRING_FIELDS = buildConf("spark.sql.debug.maxToStringFields")
.doc("Maximum number of fields of sequence-like entries can be converted
to strings " +
"in debug output. Any elements beyond the limit will be dropped and
replaced by a " +
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/AggregatingAccumulator.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/AggregatingAccumulator.scala
index bb785bfd18f0..530d58b57c3d 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/AggregatingAccumulator.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/AggregatingAccumulator.scala
@@ -18,7 +18,7 @@ package org.apache.spark.sql.execution
import scala.collection.mutable
-import org.apache.spark.TaskContext
+import org.apache.spark.{SparkContext, TaskContext}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Attribute,
AttributeReference, AttributeSeq, BindReferences, Expression,
InterpretedMutableProjection, InterpretedUnsafeProjection, JoinedRow,
MutableProjection, NamedExpression, Projection, SpecificInternalRow}
import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
DeclarativeAggregate, ImperativeAggregate, NoOp, TypedImperativeAggregate}
@@ -30,7 +30,7 @@ import org.apache.spark.util.AccumulatorV2
/**
* Accumulator that computes a global aggregate.
*/
-class AggregatingAccumulator private(
+class AggregatingAccumulator protected(
bufferSchema: Seq[DataType],
initialValues: Seq[Expression],
updateExpressions: Seq[Expression],
@@ -49,7 +49,7 @@ class AggregatingAccumulator private(
@transient
private var joinedRow: JoinedRow = _
- private var buffer: SpecificInternalRow = _
+ protected var buffer: SpecificInternalRow = _
private def createBuffer(): SpecificInternalRow = {
val buffer = new SpecificInternalRow(bufferSchema)
@@ -156,28 +156,7 @@ class AggregatingAccumulator private(
if (!other.isZero) {
other match {
case agg: AggregatingAccumulator =>
- val buffer = getOrCreateBuffer()
- val otherBuffer = agg.buffer
- mergeProjection.target(buffer)(joinedRow.withRight(otherBuffer))
- var i = 0
- while (i < imperatives.length) {
- imperatives(i).merge(buffer, otherBuffer)
- i += 1
- }
- i = 0
- if (isAtDriverSide) {
- while (i < typedImperatives.length) {
- // The input buffer stores serialized data
- typedImperatives(i).merge(buffer, otherBuffer)
- i += 1
- }
- } else {
- while (i < typedImperatives.length) {
- // The input buffer stores deserialized object
- typedImperatives(i).mergeBuffersObjects(buffer, otherBuffer)
- i += 1
- }
- }
+ mergeBuffer(agg.buffer)
case _ =>
throw QueryExecutionErrors.cannotMergeClassWithOtherClassError(
this.getClass.getName, other.getClass.getName)
@@ -185,6 +164,30 @@ class AggregatingAccumulator private(
}
}
+ protected def mergeBuffer(otherBuffer: InternalRow): Unit =
withSQLConf(true, ()) {
+ val buffer = getOrCreateBuffer()
+ mergeProjection.target(buffer)(joinedRow.withRight(otherBuffer))
+ var i = 0
+ while (i < imperatives.length) {
+ imperatives(i).merge(buffer, otherBuffer)
+ i += 1
+ }
+ i = 0
+ if (isAtDriverSide) {
+ while (i < typedImperatives.length) {
+ // The input buffer stores serialized data
+ typedImperatives(i).merge(buffer, otherBuffer)
+ i += 1
+ }
+ } else {
+ while (i < typedImperatives.length) {
+ // The input buffer stores a deserialized object
+ typedImperatives(i).mergeBuffersObjects(buffer, otherBuffer)
+ i += 1
+ }
+ }
+ }
+
override def value: InternalRow = withSQLConf(false, InternalRow.empty) {
// Either use the existing buffer or create a temporary one.
val input = if (!isZero) {
@@ -234,11 +237,24 @@ class AggregatingAccumulator private(
}
}
+private case class AggregatingAccumulatorArgs(
+ bufferSchema: Seq[DataType],
+ initialValues: Seq[Expression],
+ updateExpressions: Seq[Expression],
+ mergeExpressions: Seq[Expression],
+ resultExpressions: Seq[Expression],
+ imperatives: Array[ImperativeAggregate],
+ typedImperatives: Array[TypedImperativeAggregate[_]],
+ conf: SQLConf)
+
object AggregatingAccumulator {
/**
- * Create an aggregating accumulator for the given functions and input
schema.
+ * Build the bound expressions and aggregate buffers shared by every flavor
of aggregating
+ * accumulator created for the given functions and input schema.
*/
- def apply(functions: Seq[Expression], inputAttributes: Seq[Attribute]):
AggregatingAccumulator = {
+ private def buildAccumulatorArgs(
+ functions: Seq[Expression],
+ inputAttributes: Seq[Attribute]): AggregatingAccumulatorArgs = {
// There are a couple of things happening here:
// - Collect the schema's of the aggregate and input aggregate buffers.
These are needed to bind
// the expressions which will be done when we create the accumulator.
@@ -288,8 +304,7 @@ object AggregatingAccumulator {
val mergeAttrSeq: AttributeSeq = (aggBufferAttributes ++
inputAggBufferAttributes).toSeq
val aggBufferAttributesSeq: AttributeSeq = aggBufferAttributes.toSeq
- // Create the accumulator.
- new AggregatingAccumulator(
+ AggregatingAccumulatorArgs(
aggBufferAttributes.map(_.dataType).toSeq,
initialValues.toSeq,
updateExpressions.map(BindReferences.bindReference(_,
updateAttrSeq)).toSeq,
@@ -299,4 +314,69 @@ object AggregatingAccumulator {
typedImperatives.toArray,
SQLConf.get)
}
+
+ /**
+ * Create an aggregating accumulator for the given functions and input
schema.
+ */
+ def apply(functions: Seq[Expression], inputAttributes: Seq[Attribute]):
AggregatingAccumulator = {
+ val c = buildAccumulatorArgs(functions, inputAttributes)
+ new AggregatingAccumulator(
+ c.bufferSchema,
+ c.initialValues,
+ c.updateExpressions,
+ c.mergeExpressions,
+ c.resultExpressions,
+ c.imperatives,
+ c.typedImperatives,
+ c.conf)
+ }
+
+ /**
+ * Create an aggregating accumulator for the given functions and input
schema, and register it
+ * with the given SparkContext under `name`.
+ */
+ def apply(
+ sc: SparkContext,
+ name: String,
+ functions: Seq[Expression],
+ inputAttributes: Seq[Attribute]): AggregatingAccumulator = {
+ val acc = apply(functions, inputAttributes)
+ acc.register(sc, Option(name))
+ acc
+ }
+
+ /**
+ * Create a last-attempt aggregating accumulator for the given functions and
input schema.
+ */
+ def lastAttempt(
+ functions: Seq[Expression],
+ inputAttributes: Seq[Attribute]): LastAttemptAggregatingAccumulator = {
+ val c = buildAccumulatorArgs(functions, inputAttributes)
+ new LastAttemptAggregatingAccumulator(
+ c.bufferSchema,
+ c.initialValues,
+ c.updateExpressions,
+ c.mergeExpressions,
+ c.resultExpressions,
+ c.imperatives,
+ c.typedImperatives,
+ c.conf)
+ }
+
+ /**
+ * Create a last-attempt aggregating accumulator for the given functions and
input schema,
+ * register it with the given SparkContext under `name`, and initialize its
last-attempt
+ * tracking state.
+ */
+ def lastAttempt(
+ sc: SparkContext,
+ name: String,
+ functions: Seq[Expression],
+ inputAttributes: Seq[Attribute]): LastAttemptAggregatingAccumulator = {
+ val acc = lastAttempt(functions, inputAttributes)
+ acc.register(sc, Option(name))
+ acc.initializeLastAttemptAccumulator()
+ acc
+ }
+
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala
index fd89be2368af..3dd81ff2af6a 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala
@@ -25,6 +25,7 @@ import
org.apache.spark.sql.catalyst.plans.physical.Partitioning
import org.apache.spark.sql.catalyst.types.DataTypeUtils
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.StructType
/**
@@ -37,9 +38,13 @@ case class CollectMetricsExec(
extends UnaryExecNode {
private lazy val accumulator: AggregatingAccumulator = {
- val acc = AggregatingAccumulator(metricExpressions, child.output)
- acc.register(sparkContext, Option(CollectMetricsExec.ACCUMULATOR_NAME))
- acc
+ if (conf.getConf(SQLConf.LEGACY_OBSERVE_METRICS_AGGREGATE_ALL_ATTEMPTS)) {
+ AggregatingAccumulator(
+ sparkContext, CollectMetricsExec.ACCUMULATOR_NAME, metricExpressions,
child.output)
+ } else {
+ AggregatingAccumulator.lastAttempt(
+ sparkContext, CollectMetricsExec.ACCUMULATOR_NAME, metricExpressions,
child.output)
+ }
}
private[sql] def accumulatorId: Long = {
@@ -56,7 +61,13 @@ case class CollectMetricsExec(
.asInstanceOf[InternalRow => Row]
}
- def collectedMetrics: Row = toRowConverter(accumulator.value)
+ def collectedMetrics: Row = toRowConverter(collectedMetricsRow)
+
+ private def collectedMetricsRow: InternalRow = accumulator match {
+ case acc: LastAttemptAggregatingAccumulator =>
+ acc.lastAttemptValueForHighestRDDId().getOrElse(acc.value)
+ case acc => acc.value
+ }
override def output: Seq[Attribute] = child.output
@@ -66,6 +77,10 @@ case class CollectMetricsExec(
override def resetMetrics(): Unit = {
accumulator.reset()
+ accumulator match {
+ case acc: LastAttemptAggregatingAccumulator =>
acc.resetLastAttemptAccumulator()
+ case _ =>
+ }
super.resetMetrics()
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/LastAttemptAggregatingAccumulator.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/LastAttemptAggregatingAccumulator.scala
new file mode 100644
index 000000000000..f806fcf6b98c
--- /dev/null
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/LastAttemptAggregatingAccumulator.scala
@@ -0,0 +1,91 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.expressions.aggregate.ImperativeAggregate
+import
org.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggregate
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.DataType
+import org.apache.spark.util.{AccumulatorV2, LastAttemptAccumulator, Utils}
+
+class LastAttemptAggregatingAccumulator private[execution](
+ bufferSchema: Seq[DataType],
+ initialValues: Seq[Expression],
+ updateExpressions: Seq[Expression],
+ mergeExpressions: Seq[Expression],
+ @transient resultExpressions: Seq[Expression],
+ imperatives: Array[ImperativeAggregate],
+ typedImperatives: Array[TypedImperativeAggregate[_]],
+ @transient conf: SQLConf)
+ extends AggregatingAccumulator(
+ bufferSchema,
+ initialValues,
+ updateExpressions,
+ mergeExpressions,
+ resultExpressions,
+ imperatives,
+ typedImperatives,
+ conf)
+ with LastAttemptAccumulator[InternalRow, InternalRow, InternalRow] {
+
+ // Keep a serializable copy because both the base field and the
resultExpressions constructor
+ // parameter are @transient. Executor-side copyAndReset() needs it below.
+ private val serializedResultExpressions: Seq[Expression] = resultExpressions
+
+ // Keep a serializable snapshot because the base schema depends on its
transient
+ // resultExpressions field. isMergeable() needs it below after task
serialization.
+ private val serializedOutputSchema = schema
+
+ override protected def partialMergeVal: InternalRow = {
+ // The aggregate buffer is mutable, so last-attempt tracking keeps a
snapshot.
+ if (buffer == null) null else buffer.copy()
+ }
+
+ override protected def partialMerge(otherVal: InternalRow): Unit = {
+ if (otherVal != null) {
+ mergeBuffer(otherVal)
+ }
+ }
+
+ override protected def isMergeable(other: AccumulatorV2[_, _]): Boolean =
other match {
+ case o: LastAttemptAggregatingAccumulator =>
+ o.serializedOutputSchema == serializedOutputSchema
+ case _ => false
+ }
+
+ override protected def accumulatorStoresUserData: Boolean = true
+
+ override def copyAndReset(): LastAttemptAggregatingAccumulator =
+ new LastAttemptAggregatingAccumulator(
+ bufferSchema,
+ initialValues,
+ updateExpressions,
+ mergeExpressions,
+ serializedResultExpressions,
+ imperatives,
+ typedImperatives,
+ conf)
+
+ override def add(v: InternalRow): Unit = {
+ super.add(v)
+ if (isAtDriverSide && !Utils.isInRunningSparkTask) {
+ setValueIfOnDriverSide(value.copy())
+ }
+ }
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/AggregatingAccumulatorSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/AggregatingAccumulatorSuite.scala
index 42a20c7ae800..f59945f65cca 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/AggregatingAccumulatorSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/AggregatingAccumulatorSuite.scala
@@ -19,9 +19,11 @@ package org.apache.spark.sql.execution
import java.util.Properties
import org.apache.spark.{TaskContext, TaskContextImpl}
+import org.apache.spark.scheduler.{TaskInfo, TaskLocality}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.dsl.expressions._
-import org.apache.spark.sql.catalyst.expressions.{ExpressionEvalHelper, If,
SortArray, SparkPartitionID, SpecificInternalRow}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression,
ExpressionEvalHelper, If,
+ SortArray, SparkPartitionID, SpecificInternalRow}
import org.apache.spark.sql.catalyst.expressions.aggregate.CollectSet
import org.apache.spark.sql.catalyst.util.GenericArrayData
import org.apache.spark.sql.test.SharedSparkSession
@@ -40,131 +42,184 @@ class AggregatingAccumulatorSuite
private val inputAttributes = Seq(a, b, c)
private def str(s: String): UTF8String = UTF8String.fromString(s)
+ private def newAccumulator(
+ lastAttempt: Boolean,
+ functions: Seq[Expression],
+ inputAttributes: Seq[Attribute]): AggregatingAccumulator = {
+ if (lastAttempt) {
+ AggregatingAccumulator.lastAttempt(functions, inputAttributes)
+ } else {
+ AggregatingAccumulator(functions, inputAttributes)
+ }
+ }
+
+ /** Creates a last-attempt accumulator, registers it, and initializes its
tracking state. */
+ private def registeredLastAttempt(
+ functions: Seq[Expression],
+ inputAttributes: Seq[Attribute]): LastAttemptAggregatingAccumulator = {
+ AggregatingAccumulator.lastAttempt(
+ sparkContext, "test accumulator", functions, inputAttributes)
+ }
+
+ /** Builds a non-zero task-side accumulator by copying `main` and adding
`rows`. */
+ private def taskUpdate(
+ main: LastAttemptAggregatingAccumulator,
+ rows: InternalRow*): LastAttemptAggregatingAccumulator = {
+ val taskAcc = main.copyAndReset()
+ rows.foreach(taskAcc.add)
+ taskAcc
+ }
+
+ private def taskInfo(partitionId: Int, attemptNumber: Int = 0): TaskInfo = {
+ new TaskInfo(
+ taskId = partitionId.toLong,
+ index = partitionId,
+ attemptNumber = attemptNumber,
+ partitionId = partitionId,
+ launchTime = 0L,
+ executorId = "executor",
+ host = "localhost",
+ taskLocality = TaskLocality.ANY,
+ speculative = false)
+ }
+
test("empty aggregation") {
- val acc1 = AggregatingAccumulator(
- Seq(sum(a) + 1L as "sum_a", max(b) as "max_b", approxCountDistinct(c) as
"acntd_c"),
- inputAttributes)
- val expectedSchema = new StructType()
- .add("sum_a", "long")
- .add("max_b", "string")
- .add("acntd_c", "long", nullable = false)
- assert(acc1.schema === expectedSchema)
-
- val accEmpty = acc1.copy()
- val acc2 = acc1.copy()
-
- // Merge empty
- acc1.merge(accEmpty)
- assert(acc1.isZero)
-
- // No updates
- assert(acc1.isZero)
- checkResult(acc1.value, InternalRow(null, null, 0), expectedSchema, false)
- assert(acc1.isZero)
-
- // A few updates
- acc1.add(InternalRow(4L, str("foo"), 4.9d))
- acc1.add(InternalRow(98L, str("bar"), -323.9d))
- acc1.add(InternalRow(-30L, str("baz"), 4129.8d))
- assert(!acc1.isZero)
- checkResult(acc1.value, InternalRow(73L, str("baz"), 3L), expectedSchema,
false)
-
- // Idempotency of result
- checkResult(acc1.value, InternalRow(73L, str("baz"), 3L), expectedSchema,
false)
-
- // A few updates to the copied accumulator using an updater
- val updater = acc2.copyAndReset()
- updater.add(InternalRow(-2L, str("qwerty"), -6773.9d))
- updater.add(InternalRow(-35L, str("zzz-top"), -323.9d))
- assert(acc2.isZero)
- acc2.setState(updater)
- checkResult(acc2.value, InternalRow(-36L, str("zzz-top"), 2L),
expectedSchema, false)
-
- // Merge accumulators
- acc1.merge(acc2)
- acc1.merge(acc2)
- acc1.merge(accEmpty)
- acc1.merge(accEmpty)
- checkResult(acc1.value, InternalRow(1L, str("zzz-top"), 5L),
expectedSchema, false)
-
- // Reset
- acc1.reset()
- assert(acc1.isZero)
+ Seq(false, true).foreach { lastAttempt =>
+ val acc1 = newAccumulator(
+ lastAttempt,
+ Seq(sum(a) + 1L as "sum_a", max(b) as "max_b", approxCountDistinct(c)
as "acntd_c"),
+ inputAttributes)
+ val expectedSchema = new StructType()
+ .add("sum_a", "long")
+ .add("max_b", "string")
+ .add("acntd_c", "long", nullable = false)
+ assert(acc1.schema === expectedSchema)
+
+ val accEmpty = acc1.copy()
+ val acc2 = acc1.copy()
+
+ // Merge empty
+ acc1.merge(accEmpty)
+ assert(acc1.isZero)
+
+ // No updates
+ assert(acc1.isZero)
+ checkResult(acc1.value, InternalRow(null, null, 0), expectedSchema,
false)
+ assert(acc1.isZero)
+
+ // A few updates
+ acc1.add(InternalRow(4L, str("foo"), 4.9d))
+ acc1.add(InternalRow(98L, str("bar"), -323.9d))
+ acc1.add(InternalRow(-30L, str("baz"), 4129.8d))
+ assert(!acc1.isZero)
+ checkResult(acc1.value, InternalRow(73L, str("baz"), 3L),
expectedSchema, false)
+
+ // Idempotency of result
+ checkResult(acc1.value, InternalRow(73L, str("baz"), 3L),
expectedSchema, false)
+
+ // A few updates to the copied accumulator using an updater
+ val updater = acc2.copyAndReset()
+ updater.add(InternalRow(-2L, str("qwerty"), -6773.9d))
+ updater.add(InternalRow(-35L, str("zzz-top"), -323.9d))
+ assert(acc2.isZero)
+ acc2.setState(updater)
+ checkResult(acc2.value, InternalRow(-36L, str("zzz-top"), 2L),
expectedSchema, false)
+
+ // Merge accumulators
+ acc1.merge(acc2)
+ acc1.merge(acc2)
+ acc1.merge(accEmpty)
+ acc1.merge(accEmpty)
+ checkResult(acc1.value, InternalRow(1L, str("zzz-top"), 5L),
expectedSchema, false)
+
+ // Reset
+ acc1.reset()
+ assert(acc1.isZero)
+ }
}
test("non-deterministic expressions") {
- val acc_driver = AggregatingAccumulator(
- Seq(
- min(SparkPartitionID()) as "min_pid",
- max(SparkPartitionID()) as "max_pid",
- SparkPartitionID()),
- Nil)
- checkResult(acc_driver.value, InternalRow(null, null, 0),
acc_driver.schema, false)
-
- def inPartition(id: Int)(f: => Unit): Unit = {
- val ctx = new TaskContextImpl(0, 0, 1, 0, 0, 1, null, new Properties,
null)
- TaskContext.setTaskContext(ctx)
- try {
- f
- } finally {
- TaskContext.unset()
+ Seq(false, true).foreach { lastAttempt =>
+ val accDriver = newAccumulator(
+ lastAttempt,
+ Seq(
+ min(SparkPartitionID()) as "min_pid",
+ max(SparkPartitionID()) as "max_pid",
+ SparkPartitionID()),
+ Nil)
+ checkResult(accDriver.value, InternalRow(null, null, 0),
accDriver.schema, false)
+
+ def inPartition(id: Int)(f: => Unit): Unit = {
+ val ctx = new TaskContextImpl(0, 0, 1, 0, 0, 1, null, new Properties,
null)
+ TaskContext.setTaskContext(ctx)
+ try {
+ f
+ } finally {
+ TaskContext.unset()
+ }
}
- }
- val acc1 = acc_driver.copy()
- inPartition(3) {
- acc1.add(InternalRow.empty)
- }
- val acc2 = acc_driver.copy()
- inPartition(42) {
- acc2.add(InternalRow.empty)
- }
- val acc3 = acc_driver.copy()
- inPartition(96) {
- acc3.add(InternalRow.empty)
- }
+ val acc1 = accDriver.copy()
+ inPartition(3) {
+ acc1.add(InternalRow.empty)
+ }
+ val acc2 = accDriver.copy()
+ inPartition(42) {
+ acc2.add(InternalRow.empty)
+ }
+ val acc3 = accDriver.copy()
+ inPartition(96) {
+ acc3.add(InternalRow.empty)
+ }
- acc_driver.merge(acc1)
- acc_driver.merge(acc2)
- acc_driver.merge(acc3)
- assert(!acc_driver.isZero)
- checkResult(acc_driver.value, InternalRow(3, 96, 0), acc_driver.schema,
false)
+ accDriver.merge(acc1)
+ accDriver.merge(acc2)
+ accDriver.merge(acc3)
+ assert(!accDriver.isZero)
+ checkResult(accDriver.value, InternalRow(3, 96, 0), accDriver.schema,
false)
+ }
}
test("collect agg metrics on job") {
- val acc = AggregatingAccumulator(
- Seq(
- avg(a) + 1.0d as "avg_a",
- sum(a + 10L) as "sum_a",
- min(b) as "min_b",
- max(b) as "max_b",
- approxCountDistinct(b) as "acntd_b",
- SortArray(CollectSet(If(a < 1000L, a % 3L, a %
6L)).toAggregateExpression(), true)
- as "item_set",
- min(SparkPartitionID()) as "min_pid",
- max(SparkPartitionID()) as "max_pid",
- SparkPartitionID()),
- Seq(a, b))
- sparkContext.register(acc)
- def consume(ids: Iterator[Long]): Unit = {
- val row = new SpecificInternalRow(Seq(LongType, StringType))
- ids.foreach { id =>
- // Create the new row values.
- row.setLong(0, id)
- row.update(1, UTF8String.fromString(f"val_$id%06d"))
-
- // Update the accumulator
- acc.add(row)
+ Seq(false, true).foreach { lastAttempt =>
+ val acc = newAccumulator(
+ lastAttempt,
+ Seq(
+ avg(a) + 1.0d as "avg_a",
+ sum(a + 10L) as "sum_a",
+ min(b) as "min_b",
+ max(b) as "max_b",
+ approxCountDistinct(b) as "acntd_b",
+ SortArray(CollectSet(If(a < 1000L, a % 3L, a %
6L)).toAggregateExpression(), true)
+ as "item_set",
+ min(SparkPartitionID()) as "min_pid",
+ max(SparkPartitionID()) as "max_pid",
+ SparkPartitionID()),
+ Seq(a, b))
+ sparkContext.register(acc)
+ acc match {
+ case lastAttemptAcc: LastAttemptAggregatingAccumulator =>
+ lastAttemptAcc.initializeLastAttemptAccumulator()
+ case _ =>
+ }
+
+ def consume(ids: Iterator[Long]): Unit = {
+ val row = new SpecificInternalRow(Seq(LongType, StringType))
+ ids.foreach { id =>
+ // Create the new row values.
+ row.setLong(0, id)
+ row.update(1, UTF8String.fromString(f"val_$id%06d"))
+
+ // Update the accumulator
+ acc.add(row)
+ }
}
- }
- // Run job 1
- spark.sparkContext
- .range(0, 1000, 1, 8)
- .foreachPartition(consume)
- assert(checkResult(
- acc.value,
- InternalRow(
+ // Run job 1
+ spark.sparkContext
+ .range(0, 1000, 1, 8)
+ .foreachPartition(consume)
+ val expectedAfterJob1 = InternalRow(
500.5d,
509500L,
str("val_000000"),
@@ -173,17 +228,27 @@ class AggregatingAccumulatorSuite
new GenericArrayData(Seq(0L, 1L, 2L)),
0,
7,
- 0),
- acc.schema,
- false))
-
- // Run job 2
- spark.sparkContext
- .range(1000, 1200, 1, 8)
- .foreachPartition(consume)
- assert(checkResult(
- acc.value,
- InternalRow(
+ 0)
+ assert(checkResult(
+ acc.value,
+ expectedAfterJob1,
+ acc.schema,
+ false))
+ acc match {
+ case lastAttemptAcc: LastAttemptAggregatingAccumulator =>
+ assert(checkResult(
+ lastAttemptAcc.lastAttemptValueForAllRDDs().get,
+ expectedAfterJob1,
+ acc.schema,
+ false))
+ case _ =>
+ }
+
+ // Run job 2
+ spark.sparkContext
+ .range(1000, 1200, 1, 8)
+ .foreachPartition(consume)
+ val expectedAfterJob2 = InternalRow(
600.5d,
731400L,
str("val_000000"),
@@ -192,8 +257,201 @@ class AggregatingAccumulatorSuite
new GenericArrayData(Seq(0L, 1L, 2L, 3L, 4L, 5L)),
0,
7,
- 0),
+ 0)
+ assert(checkResult(
+ acc.value,
+ expectedAfterJob2,
+ acc.schema,
+ false))
+ acc match {
+ case lastAttemptAcc: LastAttemptAggregatingAccumulator =>
+ assert(checkResult(
+ lastAttemptAcc.lastAttemptValueForAllRDDs().get,
+ expectedAfterJob2,
+ acc.schema,
+ false))
+ case _ =>
+ }
+ }
+ }
+
+ //
---------------------------------------------------------------------------
+ // Tests exclusive to the last-attempt accumulator.
+ //
---------------------------------------------------------------------------
+
+ test("last attempt direct driver value is copied") {
+ val acc = registeredLastAttempt(Seq(sum(a) as "sum_a"), Seq(a))
+
+ val row = new SpecificInternalRow(Seq(LongType))
+ row.setLong(0, 1L)
+ acc.add(row)
+ val firstLastAttemptValue = acc.lastAttemptValueForHighestRDDId().get
+ checkResult(firstLastAttemptValue, InternalRow(1L), acc.schema, false)
+
+ row.setLong(0, 2L)
+ acc.add(row)
+ checkResult(acc.lastAttemptValueForHighestRDDId().get, InternalRow(3L),
acc.schema, false)
+ checkResult(firstLastAttemptValue, InternalRow(1L), acc.schema, false)
+ }
+
+ test("last attempt partial merge combines aggregate buffers") {
+ val acc = registeredLastAttempt(Seq(sum(a) as "sum_a", max(b) as "max_b"),
Seq(a, b))
+
+ val rdd = sparkContext.parallelize(Seq.empty[Int], 2)
+ val properties = new Properties
+
+ acc.mergeLastAttempt(
+ taskUpdate(acc, InternalRow(1L, str("a")), InternalRow(2L, str("b"))),
+ rdd,
+ taskInfo(partitionId = 0),
+ stageId = 0,
+ stageAttemptId = 0,
+ properties)
+ acc.mergeLastAttempt(
+ taskUpdate(acc, InternalRow(10L, str("z")), InternalRow(-5L, str("c"))),
+ rdd,
+ taskInfo(partitionId = 1),
+ stageId = 0,
+ stageAttemptId = 0,
+ properties)
+
+ checkResult(
+ acc.lastAttemptValueForRDDId(rdd.id).get,
+ InternalRow(8L, str("z")),
acc.schema,
- false))
+ false)
+ }
+
+ test("last attempt invalidates on schema mismatch") {
+ val acc = registeredLastAttempt(Seq(sum(a) as "sum_a"), Seq(a))
+
+ // A task accumulator with a different schema is not mergeable, so the
merge must bail out.
+ val mismatched = AggregatingAccumulator.lastAttempt(
+ Seq(sum(a) as "sum_a", max(b) as "max_b"),
+ Seq(a, b))
+ mismatched.add(InternalRow(1L, str("a")))
+
+ val rdd = sparkContext.parallelize(Seq.empty[Int], 1)
+ acc.mergeLastAttempt(
+ mismatched,
+ rdd,
+ taskInfo(partitionId = 0),
+ stageId = 0,
+ stageAttemptId = 0,
+ new Properties)
+
+ assert(!acc.getValid)
+ assert(acc.lastAttemptValueForAllRDDs().isEmpty)
+ assert(acc.lastAttemptValueForRDDId(rdd.id).isEmpty)
+ }
+
+ test("last attempt invalidates when driver and task updates are mixed") {
+ val acc = registeredLastAttempt(Seq(sum(a) as "sum_a"), Seq(a))
+
+ val rdd = sparkContext.parallelize(Seq.empty[Int], 1)
+ acc.mergeLastAttempt(
+ taskUpdate(acc, InternalRow(5L)),
+ rdd,
+ taskInfo(partitionId = 0),
+ stageId = 0,
+ stageAttemptId = 0,
+ new Properties)
+ checkResult(acc.lastAttemptValueForAllRDDs().get, InternalRow(5L),
acc.schema, false)
+
+ // A direct driver update on top of task updates cannot be reasoned about,
so it invalidates.
+ val row = new SpecificInternalRow(Seq(LongType))
+ row.setLong(0, 1L)
+ acc.add(row)
+
+ assert(!acc.getValid)
+ assert(acc.lastAttemptValueForAllRDDs().isEmpty)
+ }
+
+ test("last attempt value for highest RDD id selects the latest RDD") {
+ val acc = registeredLastAttempt(Seq(sum(a) as "sum_a"), Seq(a))
+
+ val rdd1 = sparkContext.parallelize(Seq.empty[Int], 1)
+ val rdd2 = sparkContext.parallelize(Seq.empty[Int], 1)
+ assert(rdd2.id > rdd1.id)
+
+ acc.mergeLastAttempt(
+ taskUpdate(acc, InternalRow(3L)),
+ rdd1,
+ taskInfo(partitionId = 0),
+ stageId = 0,
+ stageAttemptId = 0,
+ new Properties)
+ acc.mergeLastAttempt(
+ taskUpdate(acc, InternalRow(10L)),
+ rdd2,
+ taskInfo(partitionId = 0),
+ stageId = 1,
+ stageAttemptId = 0,
+ new Properties)
+
+ // Highest RDD id aggregates only the latest RDD.
+ checkResult(acc.lastAttemptValueForHighestRDDId().get, InternalRow(10L),
acc.schema, false)
+ // All RDDs aggregate both.
+ checkResult(acc.lastAttemptValueForAllRDDs().get, InternalRow(13L),
acc.schema, false)
+ // Narrowing to the earlier RDD returns only its value.
+ checkResult(acc.lastAttemptValueForRDDId(rdd1.id).get, InternalRow(3L),
acc.schema, false)
+ }
+
+ test("last attempt skips uncomputed partitions") {
+ val acc = registeredLastAttempt(Seq(sum(a) as "sum_a"), Seq(a))
+
+ // The RDD has 3 partitions but only 0 and 2 report values (e.g. early
stop in take/limit).
+ val rdd = sparkContext.parallelize(Seq.empty[Int], 3)
+ acc.mergeLastAttempt(
+ taskUpdate(acc, InternalRow(4L)),
+ rdd,
+ taskInfo(partitionId = 0),
+ stageId = 0,
+ stageAttemptId = 0,
+ new Properties)
+ acc.mergeLastAttempt(
+ taskUpdate(acc, InternalRow(6L)),
+ rdd,
+ taskInfo(partitionId = 2),
+ stageId = 0,
+ stageAttemptId = 0,
+ new Properties)
+
+ // Partition 1 never reported, so only partitions 0 and 2 contribute.
+ checkResult(acc.lastAttemptValueForRDDId(rdd.id).get, InternalRow(10L),
acc.schema, false)
+ }
+
+ test("last attempt uses the latest stage attempt for a partition") {
+ val acc = registeredLastAttempt(Seq(sum(a) as "sum_a"), Seq(a))
+
+ val rdd = sparkContext.parallelize(Seq.empty[Int], 1)
+ acc.mergeLastAttempt(
+ taskUpdate(acc, InternalRow(3L)),
+ rdd,
+ taskInfo(partitionId = 0),
+ stageId = 0,
+ stageAttemptId = 0,
+ new Properties)
+ checkResult(acc.lastAttemptValueForRDDId(rdd.id).get, InternalRow(3L),
acc.schema, false)
+
+ // A later stage attempt recomputes the same partition and replaces its
value.
+ acc.mergeLastAttempt(
+ taskUpdate(acc, InternalRow(100L)),
+ rdd,
+ taskInfo(partitionId = 0),
+ stageId = 0,
+ stageAttemptId = 1,
+ new Properties)
+ checkResult(acc.lastAttemptValueForRDDId(rdd.id).get, InternalRow(100L),
acc.schema, false)
+
+ // An older attempt arriving late is discarded.
+ acc.mergeLastAttempt(
+ taskUpdate(acc, InternalRow(7L)),
+ rdd,
+ taskInfo(partitionId = 0),
+ stageId = 0,
+ stageAttemptId = 0,
+ new Properties)
+ checkResult(acc.lastAttemptValueForRDDId(rdd.id).get, InternalRow(100L),
acc.schema, false)
}
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala
index 06b0d64944e6..33da022f7aa3 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala
@@ -16,15 +16,20 @@
*/
package org.apache.spark.sql.execution.metric
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.collection.mutable.ArrayBuffer
import scala.util.Random
import org.apache.spark.SparkException
import org.apache.spark.internal.config
-import org.apache.spark.sql.{Column, Dataset}
+import org.apache.spark.sql.{Column, Dataset, Observation, Row}
+import org.apache.spark.sql.execution.QueryExecution
import org.apache.spark.sql.execution.adaptive.{AQETestHelper,
DisableAdaptiveExecutionSuite}
-import org.apache.spark.sql.functions.udf
+import org.apache.spark.sql.functions.{approx_count_distinct, collect_list,
collect_set, count, lit, max, min, size, sort_array, sum, udf}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.util.QueryExecutionListener
class MetricsFailureInjectionSuite
extends SharedSparkSession
@@ -299,6 +304,80 @@ class MetricsFailureInjectionSuite
}
}
+ test("Observation metrics select attempts according to legacy config") {
+ withTable("test_table") {
+ setUpTestTable("test_table")
+ withSparkContextConf(
+ config.Tests.INJECT_SHUFFLE_FETCH_FAILURES.key -> "true") {
+ Seq(false, true).foreach { aggregateAllAttempts =>
+ withSQLConf(
+ SQLConf.LEGACY_OBSERVE_METRICS_AGGREGATE_ALL_ATTEMPTS.key ->
+ aggregateAllAttempts.toString) {
+ val observation = Observation("stage1")
+ val observed = spark.read.table("test_table")
+ .observe(observation, count(lit(1)).as("rows"))
+ val finalDf =
observed.groupBy("low_cardinality_col").count().as[(Int, Long)]
+ val result = finalDf.collect()
+
+ assert(result.toMap === (0 until 5).map(v => (v, 300 / 5)).toMap)
+ if (aggregateAllAttempts) {
+ assert(observation.get("rows").asInstanceOf[Long] > 300L)
+ } else {
+ assert(observation.get === Map("rows" -> 300L))
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("Observation metrics reset last-attempt state after failure injection")
{
+ val includeRows = new AtomicBoolean(true)
+ val includeRow = udf { _: Long => includeRows.get() }
+ val metricMaps = ArrayBuffer.empty[Map[String, Row]]
+ val listener = new QueryExecutionListener {
+ override def onSuccess(funcName: String, qe: QueryExecution, duration:
Long): Unit = {
+ if (qe.observedMetrics.nonEmpty) {
+ metricMaps += qe.observedMetrics
+ }
+ }
+
+ override def onFailure(funcName: String, qe: QueryExecution, exception:
Exception): Unit = {
+ // No-op
+ }
+ }
+
+ withTable("test_table") {
+ setUpTestTable("test_table")
+ withSparkContextConf(
+ config.Tests.INJECT_SHUFFLE_FETCH_FAILURES.key -> "true") {
+ val observed = spark.read.table("test_table")
+ .filter(includeRow($"id"))
+ .observe(name = "stage1", count(lit(1)).as("rows"))
+ val finalDf = observed.groupBy("low_cardinality_col").count().as[(Int,
Long)]
+
+ spark.listenerManager.register(listener)
+ try {
+ val result = finalDf.collect()
+ sparkContext.listenerBus.waitUntilEmpty()
+ assert(result.toMap === (0 until 5).map(v => (v, 300 / 5)).toMap)
+ assert(metricMaps.size === 1)
+ assert(metricMaps.head("stage1") === Row(300L))
+ metricMaps.clear()
+
+ // Reuse the same Dataset while making the second execution produce
no rows.
+ includeRows.set(false)
+ assert(finalDf.collect().isEmpty)
+ sparkContext.listenerBus.waitUntilEmpty()
+ assert(metricMaps.size === 1)
+ assert(metricMaps.head("stage1") === Row(0L))
+ } finally {
+ spark.listenerManager.unregister(listener)
+ }
+ }
+ }
+ }
+
for {
injectFailure <- BOOLEAN_DOMAIN
} test(s"Non-deterministic stage block failure injection -
injectFailure=$injectFailure") {
@@ -686,4 +765,195 @@ class MetricsFailureInjectionSuite
}
}
}
+
+ // Observed metrics are computed by the last-attempt aggregating
accumulator. This exercises
+ // every aggregate-buffer flavor when a shuffle-fetch failure forces the
observed (leaf) stage
+ // to recompute:
+ // - declarative aggregates (count / sum / min / max),
+ // - imperative, non-typed aggregate (approx_count_distinct ->
HyperLogLogPlusPlus),
+ // - typed imperative aggregates (collect_set and collect_list).
+ // Regardless of the recompute, the observed values must equal the
single-attempt values,
+ // proving the last-attempt accumulator's partial-merge and buffer
snapshotting are correct for
+ // imperative and typed-imperative buffers, not just simple long counters.
+ // collect_list is the strongest guard here: unlike collect_set /
approx_count_distinct (which
+ // are idempotent under re-merge, so merging a superseded attempt would
still yield the right
+ // set/estimate), a collect_list would double in length if a superseded
attempt were not
+ // discarded, so size(collect_list) == 300 deterministically proves
last-attempt selection for a
+ // typed-imperative buffer.
+ for {
+ injectFailure <- BOOLEAN_DOMAIN
+ } test(s"Observation metrics with all aggregate buffer types survive failure
injection - " +
+ s"injectFailure=$injectFailure") {
+ val observation = Observation("mixed_aggregate_buffers")
+
+ withTable("test_table") {
+ setUpTestTable("test_table")
+ withSparkContextConf(
+ config.Tests.INJECT_SHUFFLE_FETCH_FAILURES.key ->
injectFailure.toString) {
+ val observed = spark.read.table("test_table")
+ .observe(
+ observation,
+ count(lit(1)).as("rows"),
+ sum($"low_cardinality_col").as("sum_lcc"),
+ min($"id").as("min_id"),
+ max($"id").as("max_id"),
+
approx_count_distinct($"low_cardinality_col").as("approx_distinct_lcc"),
+ sort_array(collect_set($"low_cardinality_col")).as("distinct_lcc"),
+ size(collect_list($"low_cardinality_col")).as("list_size_lcc"))
+ val finalDf = observed.groupBy("low_cardinality_col").count().as[(Int,
Long)]
+ val result = finalDf.collect()
+
+ assert(result.toMap === (0 until 5).map(v => (v, 300 / 5)).toMap)
+
+ // low_cardinality_col = id % 5 over ids 0..299: each residue in
{0,1,2,3,4} appears 60
+ // times, so sum = 60 * (0+1+2+3+4) = 600, and the distinct set is
{0,1,2,3,4}.
+ val metrics = observation.get
+ assert(metrics("rows") === 300L)
+ assert(metrics("sum_lcc") === 600L)
+ assert(metrics("min_id") === 0L)
+ assert(metrics("max_id") === 299L)
+ assert(metrics("approx_distinct_lcc") === 5L)
+ assert(metrics("distinct_lcc") === Seq(0, 1, 2, 3, 4))
+ // 300 rows total; would be 600 if a superseded attempt's collect_list
were not discarded.
+ assert(metrics("list_size_lcc") === 300)
+ }
+ }
+ }
+
+ // Observed metrics are produced by the last-attempt aggregating
accumulator. A forced AQE
+ // cancellation re-runs the map stage feeding the shuffle, creating a new
RDD for the replanned
+ // stage - the canonical case the last-attempt tracking is designed for.
Regardless of the
+ // replan, the observed values must reflect only the last attempt, for every
aggregate-buffer
+ // flavor (declarative / imperative / typed-imperative). AQE must be enabled
for the forced
+ // cancellation to take effect (it is a no-op otherwise), so this scenario
is AQE-only.
+ test("Observation metrics AQE cancellation injection") {
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
+ val stage1Metric = SQLMetrics.createMetric(spark.sparkContext, "stage 1
counter")
+
+ def runQueryWithObservation(triggerMetrics: SQLMetric*): Unit = {
+ assert(stage1Metric.value === 0)
+ val observation = Observation("aqe_obs")
+ withTable("test_table") {
+ setUpTestTable("test_table")
+ AQETestHelper.withForcedCancellation(triggerMetrics: _*) {
+ val stage1MetricsExpr = incrementMetrics(Seq(stage1Metric))
+ val stage1 = spark.read.table("test_table")
+ .observe(
+ observation,
+ count(lit(1)).as("rows"),
+ sum($"low_cardinality_col").as("sum_lcc"),
+
approx_count_distinct($"low_cardinality_col").as("approx_distinct_lcc"),
+
sort_array(collect_set($"low_cardinality_col")).as("distinct_lcc"))
+ .filter(Column(stage1MetricsExpr))
+ val finalDf =
stage1.groupBy("low_cardinality_col").count().as[(Int, Long)]
+ val result = finalDf.collect()
+ assert(result.toMap === (0 until 5).map(v => (v, 300 / 5)).toMap)
+
+ // low_cardinality_col = id % 5 over ids 0..299: each residue
appears 60 times, so
+ // sum = 60 * (0+1+2+3+4) = 600 and the distinct set is
{0,1,2,3,4}.
+ val metrics = observation.get
+ assert(metrics("rows") === 300L)
+ assert(metrics("sum_lcc") === 600L)
+ assert(metrics("approx_distinct_lcc") === 5L)
+ assert(metrics("distinct_lcc") === Seq(0, 1, 2, 3, 4))
+ stage1Metric.reset()
+ }
+ }
+ }
+
+ // Case 1: No forced replanning (single-attempt baseline).
+ runQueryWithObservation()
+ // Case 2: Replan on stage1Metric, so the map stage feeding the shuffle
re-runs with a new
+ // RDD; the observed metrics must still reflect only the last attempt.
+ runQueryWithObservation(stage1Metric)
+ }
+ }
+
+ // The other observation failure tests all observe the leaf scan, so only a
*downstream* stage
+ // recomputes while the observed operator itself runs once. Here the
observation sits on a
+ // non-leaf operator (the output of a join that feeds a shuffle), so the
observed operator's own
+ // RDD is the one re-run when the downstream fetch failure fires. This
exercises per-RDD
+ // last-attempt selection for the aggregate buffer on a recomputed non-leaf
stage, across every
+ // aggregate-buffer flavor, rather than only on a leaf scan.
+ test("Observation metrics on a non-leaf stage survive failure injection") {
+ val observation = Observation("non_leaf_obs")
+
+ withTable("primary_table", "secondary_table") {
+ setUpTestTable("primary_table")
+ setUpTestTable("secondary_table")
+ withSparkContextConf(
+ config.Tests.INJECT_SHUFFLE_FETCH_FAILURES.key -> "true") {
+ val joined = spark.read.table("primary_table")
+ .join(
+ spark.read.table("secondary_table"),
+ usingColumn = "id",
+ joinType = "fullOuter")
+ .observe(
+ observation,
+ count(lit(1)).as("rows"),
+ sum("primary_table.low_cardinality_col").as("sum_lcc"),
+
approx_count_distinct("primary_table.low_cardinality_col").as("approx_distinct_lcc"),
+
sort_array(collect_set("primary_table.low_cardinality_col")).as("distinct_lcc"))
+ val finalDf =
joined.groupBy("primary_table.low_cardinality_col").count().as[(Int, Long)]
+ val result = finalDf.collect()
+ assert(result.toMap === (0 until 5).map(v => (v, 300 / 5)).toMap)
+
+ // A full-outer join of the two identical 300-row tables on the unique
`id` yields 300
+ // rows, so the observed values match the single-attempt leaf values:
each residue in
+ // {0,1,2,3,4} appears 60 times, so sum = 60 * (0+1+2+3+4) = 600.
+ val metrics = observation.get
+ assert(metrics("rows") === 300L)
+ assert(metrics("sum_lcc") === 600L)
+ assert(metrics("approx_distinct_lcc") === 5L)
+ assert(metrics("distinct_lcc") === Seq(0, 1, 2, 3, 4))
+ }
+ }
+ }
+
+ // The failure suite otherwise has a single observe per query, and
DataFrameCallbackSuite has
+ // multi-observe queries but no retries. This combines them: two observation
nodes at different
+ // stages (a leaf scan and a non-leaf join output), both of which recompute
under fetch-failure
+ // injection. Each observation is backed by its own CollectMetricsExec /
last-attempt
+ // accumulator, so this checks that they independently report their own
correct last-attempt
+ // value rather than leaking into one another.
+ test("Multiple observation points survive failure injection") {
+ val leafObs = Observation("leaf_obs")
+ val nonLeafObs = Observation("non_leaf_obs")
+
+ withTable("primary_table", "secondary_table") {
+ setUpTestTable("primary_table")
+ setUpTestTable("secondary_table")
+ withSparkContextConf(
+ config.Tests.INJECT_SHUFFLE_FETCH_FAILURES.key -> "true") {
+ val leaf = spark.read.table("primary_table")
+ .observe(
+ leafObs,
+ count(lit(1)).as("rows"),
+ sort_array(collect_set("low_cardinality_col")).as("distinct_lcc"))
+ val joined = leaf
+ .join(
+ spark.read.table("secondary_table"),
+ usingColumn = "id",
+ joinType = "fullOuter")
+ .observe(
+ nonLeafObs,
+ count(lit(1)).as("rows"),
+ sum("primary_table.low_cardinality_col").as("sum_lcc"))
+ val finalDf =
joined.groupBy("primary_table.low_cardinality_col").count().as[(Int, Long)]
+ val result = finalDf.collect()
+ assert(result.toMap === (0 until 5).map(v => (v, 300 / 5)).toMap)
+
+ // The leaf scan and the join output each recompute, but every
observation reports only its
+ // last successful attempt: 300 rows on both, distinct {0,1,2,3,4} on
the leaf, and
+ // sum = 60 * (0+1+2+3+4) = 600 on the join output.
+ val leafMetrics = leafObs.get
+ assert(leafMetrics("rows") === 300L)
+ assert(leafMetrics("distinct_lcc") === Seq(0, 1, 2, 3, 4))
+
+ val nonLeafMetrics = nonLeafObs.get
+ assert(nonLeafMetrics("rows") === 300L)
+ assert(nonLeafMetrics("sum_lcc") === 600L)
+ }
+ }
+ }
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/util/DataFrameCallbackSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/util/DataFrameCallbackSuite.scala
index cccd0a87d4cc..2912917a7be3 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/util/DataFrameCallbackSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/util/DataFrameCallbackSuite.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.util
import java.lang.{Long => JLong}
import java.util.concurrent.{CopyOnWriteArrayList, CountDownLatch, TimeUnit}
+import java.util.concurrent.atomic.AtomicBoolean
import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
@@ -401,6 +402,42 @@ class DataFrameCallbackSuite extends SharedSparkSession
}
}
+ test("observed metrics reset last-attempt state between reused executions") {
+ val includeRows = new AtomicBoolean(true)
+ val includeRow = udf { _: Long => includeRows.get() }
+ val df = spark.range(0, 10, 1, 2)
+ .filter(includeRow($"id"))
+ .observe(name = "my_event", count(lit(1)).as("rows"))
+
+ val metricMaps = ArrayBuffer.empty[Map[String, Row]]
+ val listener = new QueryExecutionListener {
+ override def onSuccess(funcName: String, qe: QueryExecution, duration:
Long): Unit = {
+ metricMaps += qe.observedMetrics
+ }
+
+ override def onFailure(funcName: String, qe: QueryExecution, exception:
Exception): Unit = {
+ // No-op
+ }
+ }
+ spark.listenerManager.register(listener)
+ try {
+ assert(df.collect().length === 10)
+ sparkContext.listenerBus.waitUntilEmpty()
+ assert(metricMaps.size === 1)
+ assert(metricMaps.head("my_event") === Row(10L))
+ metricMaps.clear()
+
+ // Reuse the same Dataset while making the second execution produce no
rows.
+ includeRows.set(false)
+ assert(df.collect().isEmpty)
+ sparkContext.listenerBus.waitUntilEmpty()
+ assert(metricMaps.size === 1)
+ assert(metricMaps.head("my_event") === Row(0L))
+ } finally {
+ spark.listenerManager.unregister(listener)
+ }
+ }
+
test("SPARK-50581: support observe with udaf") {
withUserDefinedFunction(("someUdaf", true)) {
spark.udf.register("someUdaf", functions.udaf(new Aggregator[JLong,
JLong, JLong] {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
