ulysses-you commented on code in PR #57742:
URL: https://github.com/apache/spark/pull/57742#discussion_r3755857808
##########
sql/core/benchmarks/AdaptivePartialAggregationBenchmark-results.txt:
##########
@@ -0,0 +1,56 @@
+================================================================================================
+high-cardinality input, no-spill pass-through (Tier 1)
Review Comment:
The sentinel nit is fixed too: both spill scenarios now use `0` (the
documented sentinel that disables the periodic check) instead of
`Long.MaxValue`.
##########
docs/sql-performance-tuning.md:
##########
@@ -181,6 +181,58 @@ Missing or inaccurate statistics will hinder Spark's
ability to select an optima
- **Query plan estimates**: You can inspect Spark's cost estimates in the
optimized query plan via [`EXPLAIN COST`](sql-ref-syntax-qry-explain.html) or
`DataFrame.explain(mode="cost")`.
- **Runtime statistics**: You can inspect these statistics in the [SQL
UI](web-ui.html#sql-tab) under the "Details" section as a query is running.
Look for `Statistics(..., isRuntime=true)` in the plan.
+## Optimizing the Aggregate
+
+### Adaptive Partial Aggregation
+
+A grouping aggregation normally runs in two phases: a partial aggregation
before the shuffle and a
+final aggregation after it. The partial aggregation is only worthwhile when it
actually reduces the
+number of rows; when the grouping keys are close to unique it maintains -- and
possibly spills -- an
+aggregation map roughly as large as its input while emitting almost as many
rows as it consumed.
+
+When adaptive partial aggregation is enabled, hash aggregation measures the
compaction ratio (the
+number of processed rows divided by the number of keys held in its aggregation
maps) at runtime
+and, if the partial aggregation is not collapsing enough rows to be
worthwhile, stops populating
+the aggregation map and passes the remaining rows through as single-row
partial aggregation buffers
+for the final aggregation to merge. Query results are unchanged. The ratio is
evaluated
Review Comment:
The docs now implement the wording we agreed: the periodic evaluation uses
the cumulative rows and keys of the current map epoch, so "a favorable prefix
can mask a later distinct-heavy tail, keeping the aggregation on until a spill
restarts the accounting". A late turn is caught only when a spill starts a new
epoch, and the tuning notes point at `minCompaction` and `minRows`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala:
##########
@@ -354,6 +410,47 @@ class TungstenAggregationIterator(
}
}
+ ///////////////////////////////////////////////////////////////////////////
+ // Part 5b: Methods and fields used by adaptive partial aggregation
pass-through.
+ ///////////////////////////////////////////////////////////////////////////
+
+ // Indicates that partial aggregation has been bypassed and the remaining
input rows should be
+ // passed through as single-row partial buffers. Set in `processInputs` by
either check point.
+ // It may coexist with earlier spills, so the output order is: sort-based
(or map) output first,
+ // then the pass-through rows.
+ private[this] var passThrough: Boolean = false
+
+ // The row that could not be inserted at the spill check. It is stashed here
(as
+ // a copy) so it becomes the first pass-through row rather than being lost.
+ private[this] var pendingPassThroughRow: InternalRow = null
+
+ // A reused aggregation buffer for building single-row partial buffers
during pass-through. It is
+ // re-initialized from `initialAggregationBuffer` for every passed-through
row.
+ private[this] lazy val passThroughAggregationBuffer: UnsafeRow =
createNewAggregationBuffer()
+
+ // Whether there are remaining pass-through rows to emit.
+ private def passThroughHasNext: Boolean =
+ passThrough && (pendingPassThroughRow != null || inputIter.hasNext)
+
+ // Emits the next input row as a single-row partial aggregation buffer, i.e.
a group of size one.
+ // The output (grouping key ++ buffer) is a valid partial buffer that the
downstream Final
+ // aggregation merges, so the result is identical to running partial
aggregation on this row.
+ private def nextPassThroughOutput(): UnsafeRow = {
Review Comment:
Fixed in `dd0f15465e9`: the interpreted path now times each
`nextPassThroughOutput()` drain into `aggTime` (threaded alongside
`numBypassingRows`), so the SQL UI covers the pass-through work the iterator
does after the constructor, which is exactly the workload the feature targets.
The codegen side already accounted for it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -141,6 +148,40 @@ case class HashAggregateExec(
.map(_.asInstanceOf[DeclarativeAggregate])
private val bufferSchema =
DataTypeUtils.fromAttributes(aggregateBufferAttributes)
+ /**
+ * Whether adaptive partial aggregation applies to this operator. When it
does, the aggregation
+ * may bypass partial aggregation at runtime and pass the remaining input
rows through as
+ * single-row partial buffers (see
[[SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED]]). It only
+ * applies to a pre-shuffle partial aggregation with grouping keys:
+ * - `Partial` mode only: the downstream `Final` aggregation merges the
passed-through
+ * single-row buffers, so the output contract is unchanged.
`Final`/`Complete` produce the
+ * result themselves and have no such downstream. `PartialMerge` does
have one and could be
+ * supported by passing its incoming buffer through unchanged, but that
is left for later.
+ * - grouping keys present: a global aggregation produces a single output
row, so partial
+ * aggregation achieves the maximum reduction and must never be bypassed.
+ * - DISTINCT aggregate functions are allowed: the intermediate
`PartialMerge` phase of the
+ * multi-phase distinct plan is not `Partial` mode (and requires a
distribution), so it always
+ * aggregates and de-duplicates, and the passed-through rows from the
distinct `Partial` phase
+ * therefore carry exactly one distinct value each.
+ */
+ private val adaptivePartialAggEnabled: Boolean = {
+ conf.adaptivePartialAggregationEnabled &&
+ groupingExpressions.nonEmpty &&
+ // Only the pre-shuffle partial aggregation has a downstream `Final` to
merge passed-through
+ // single-row buffers. `requiredChildDistributionExpressions` is `None`
exactly for that
+ // pre-shuffle phase and `Some` for the `Final`/`Complete` phase. This
check is what keeps a
+ // group-by-only aggregate (no aggregate functions, so an empty
`aggregateExpressions`) from
+ // being admitted vacuously: `aggregateExpressions.forall(_.mode ==
Partial)` alone is true
+ // for the empty list, which would wrongly make the `Final` phase
eligible as well.
+ requiredChildDistributionExpressions.isEmpty &&
+ aggregateExpressions.forall(a => a.mode == Partial)
Review Comment:
Fixed. The feature is gated on `!isStreaming`, and a batch `session_window`
grouping is also kept out by checking `SessionWindow.marker` in the grouping
expressions, matching the static sibling's session-window opt-out. A regression
test (`no pass-through for a session_window grouping key`) runs fully distinct
sessions with a small `minRows` and asserts no bypass; the full suite passes 49
tests. Streaming and `session_window` support are left as a TODO.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala:
##########
@@ -191,29 +209,67 @@ class TungstenAggregationIterator(
}
} else {
var i = 0
- while (inputIter.hasNext) {
+ var processedRows = 0L
+ val minRows = adaptiveMinRows
+ // The processed-row count at which the compaction ratio is evaluated
next. It advances by
+ // `minRows` after every check, and restarts after a spill so the new
in-memory map epoch is
+ // judged on its own rows. `minRows = 0` disables the periodic check:
the count is only ever
+ // compared after being incremented past 0, so it never matches and only
the spill check
+ // below remains.
+ var nextCheckRow = minRows
+ // The partial aggregation is ineffective when it does not collapse
`minCompaction` rows into
+ // one key. There is no fast map on this path, so the map's keys are all
the operator holds.
+ def ineffective(): Boolean =
+ processedRows < hashMap.getNumKeys().toDouble * adaptiveMinCompaction
+ while (inputIter.hasNext && !passThrough) {
Review Comment:
Not adopted, for the reasons you listed as the honest cost: probe-retention
keeps the per-row probe on the aggregation-is-worth-keeping path (roughly the
map-work part of the win the skip avoids today), holds the map to the end of
the task instead of freeing it early, changes the meaning of
`numBypassingRows`, and needs the find-only accessor in all three map
implementations. And as you noted on finding 2, it only rescues a tail whose
keys the prefix already inserted. The order property it guarantees is real, and
your finding 1 repro is exactly the case it would settle; we fixed that repro
directly instead and documented the residual first-row collision in the order
test.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -4156,6 +4156,46 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val ADAPTIVE_PARTIAL_AGGREGATION_ENABLED =
+
buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation.enabled")
+ .doc("When true, hash aggregation adaptively bypasses the pre-shuffle
partial aggregation " +
+ "at runtime when it observes that the partial aggregation is not
reducing the number of " +
+ "rows enough to be worthwhile. Once bypassed, the remaining input rows
are passed " +
+ "through as single-row partial aggregation buffers for the final
aggregation to merge, " +
+ "which avoids the cost of maintaining and spilling a large aggregation
map with little " +
+ "reduction benefit. This applies only to hash aggregation with
grouping keys.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(true)
Review Comment:
Agreed it is a small real cost, and the benchmark comment no longer claims
the runs "must match". The low-cardinality no-spill scenario now says the
adaptive path pays a small per-row overhead (the stop check, the bypass counter
and the check-point compare) even when nothing bypasses, and the measured
difference is a few percent rather than a regression.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -535,19 +642,27 @@ case class HashAggregateExec(
// `addNewFunction` spills this helper into a nested class (as can happen
// once the outer class passes the code-size threshold), the bare field
// reference fails with `IllegalAccessError`.
- val doAggFuncName = ctx.addNewFunction(doAgg,
- s"""
- |private void $doAgg(int partitionIndex) throws java.io.IOException {
- | ${child.asInstanceOf[CodegenSupport].produce(ctx, this)}
- | $finishHashMap
- |}
- """.stripMargin)
- // generate code for output
+ // Generate code for output. This must happen before the `doAgg` helper
below, because with
+ // adaptive partial aggregation enabled, `doConsumeWithKeys` (invoked from
the child's produce
+ // inside `doAgg`) emits pass-through rows by calling this output function
directly.
val keyTerm = ctx.freshName("aggKey")
val bufferTerm = ctx.freshName("aggBuffer")
- val outputFunc = generateResultFunction(ctx)
-
+ outputFunc = generateResultFunction(ctx)
Review Comment:
Fixed. The output function is generated before `doAgg` only when the feature
is enabled; for a non-adaptive plan it is generated after `doAgg` registers
(still before `adaptiveFinalOutput` reads it), so `generateResultFunction`'s
`consume(...)`, fresh names and mutable-state slots reach plans the feature
never touches in the upstream order.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]