HyukjinKwon commented on code in PR #59149:
URL: https://github.com/apache/spark/pull/59149#discussion_r4173656666
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/PivotTransformer.scala:
##########
@@ -120,12 +121,50 @@ object PivotTransformer extends AliasHelper with
SQLConfHelper {
case _ =>
newAlias(pivotColumn, Some("__pivot_col"))
}
- val extendedGroupingExpressions = groupByExpressions :+ namedPivotCol
+ // SPARK-55569: Collapse non-matching pivot column values to null so the
firstAgg produces
+ // at most one group per (groupBy-key, null) instead of one group per
distinct non-matching
+ // value. PivotFirst ignores the null group, but the groupBy key is
preserved in secondAgg,
+ // keeping the semantics identical to the un-optimized plan. This
reduces the number of
+ // groups when the table has many distinct values outside the pivot IN
list.
+ // Skip when the pivot column contains an aggregate (SPARK-24722) or
when a NULL pivot
+ // value is present (non-matching rows would merge with legitimate null
rows).
+ val hasNullPivotValue =
+ evalPivotValues.contains(null)
+ val canCollapse =
+ !pivotColumn.exists(_.isInstanceOf[AggregateFunction]) &&
+ !hasNullPivotValue
+ val effectivePivotCol = if (canCollapse) {
+ val castPivotExprs = pivotValues.map { value =>
+ Cast(value, pivotColumn.dataType,
+ Some(conf.sessionLocalTimeZone))
+ }
+ val nonNullExprs = castPivotExprs
+ .zip(evalPivotValues)
+ .collect { case (expr, v) if v != null => expr }
+ val matchesPivotValue = if (nonNullExprs.nonEmpty) {
+ In(pivotColumn, nonNullExprs)
+ } else {
+ Literal.FalseLiteral
+ }
+ newAlias(
+ If(matchesPivotValue, pivotColumn,
Review Comment:
**Non-blocking (P2):** The collapse maps every non-matching pivot value
within a group-by key to a single `null` group, but `firstAgg` still evaluates
the pivot aggregates (`sum`, `avg`, ...) over that merged group; `PivotFirst`
only discards the null group's result afterwards in the second aggregation. For
a `LongType` or `DecimalType` measure, the merged non-matching rows can sum to
a value that overflows the accumulator even when no individual non-matching
value's group would. Under ANSI mode (enabled by default), that overflow raises
`SparkArithmeticException` during `firstAgg`, so a pivot that previously
returned results — including the high-cardinality case this optimization
targets — can now fail at runtime, even though the overflowing aggregate is
immediately discarded.
**Recommended change:** In the collapse branch, mask the non-matching rows
out of the firstAgg aggregates (not only out of the grouping key) so the
discarded null group never aggregates real measure values. Reuse the
already-constructed matchesPivotValue (In) predicate to wrap each pivot
aggregate's input, mirroring the standard path's createIfExpression masking but
with a single shared predicate rather than one per pivot value. Add a
DataFramePivotSuite regression test exercising an overflow-magnitude
non-matching measure under ANSI, and regenerate the affected pivot
analyzer-result golden files.
**Why this works:** Within a firstAgg group the collapsed grouping key is
uniform: a group is either all-matching (In=true, key=col) or all-non-matching
(In=false/null, key=null). Wrapping each aggregate's measure input in
If(matchesPivotValue, input, null) therefore leaves every matching group's
aggregate byte-identical (all its rows satisfy the predicate) while forcing the
merged null group's inputs to null, so sum/avg over it yields null instead of
overflowing. PivotFirst still discards the null group and preserves group-by
keys, so observable results equal the un-optimized plan under both ANSI and
non-ANSI modes.
**Scope:** Prevent the collapsed non-matching null group from computing an
overflow-prone aggregate in the fast-path firstAgg, keeping pivot results
identical to the un-optimized plan under default ANSI mode while preserving the
group-reduction optimization.
**Compatibility:** All currently-succeeding pivot queries return identical
results, and group-by keys whose rows are entirely non-matching are still
emitted with all-null pivot columns.
**Risks:** Masking must reuse the existing matchesPivotValue predicate
rather than re-deriving per-value conditions, or it would reintroduce the
per-row cost the fast path avoids. The regenerated golden files must be
reviewed to confirm only the expected If-masking of firstAgg aggregates
changed, with pivot output projections unchanged.
**Constraints:** Must preserve group-by keys whose rows are all non-matching
(do not filter rows); the null group row must still be produced so secondAgg
emits the all-null pivot output.
**Success:** A pivot query whose merged non-matching measure sum overflows
the accumulator while per-value sums do not returns the same results as the
un-optimized plan under default ANSI mode, with no SparkArithmeticException.
Matching-value pivot outputs, multiple-aggregate results, and preserved
all-null group-by keys are unchanged (existing DataFramePivotSuite cases pass
after golden-file regeneration). Non-ANSI behavior is unchanged (the null group
was and remains discarded).
**Verification:**
- **Regression:** With ANSI enabled, a pivot over a high-cardinality
non-matching set with a Long or Decimal measure whose merged non-matching sum
overflows returns correct results instead of raising SparkArithmeticException;
this case fails before the fix.
- **Behavior:** Existing result-equality pivot cases (single/multiple
aggregates, preserved all-null group-by key, null-pivot skip) continue to
return identical rows.
- **Compatibility:** Non-ANSI pivot results remain unchanged because the
collapsed null group is still discarded.
- **Inspection:** Regenerated pivot analyzer-result golden files show the
firstAgg pivot aggregates masked by the shared matchesPivotValue predicate and
otherwise-unchanged pivot output projection.
--
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]
