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]

Reply via email to