peter-toth commented on code in PR #59149:
URL: https://github.com/apache/spark/pull/59149#discussion_r4177785729
##########
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,
+ Literal(null, pivotColumn.dataType)),
+ Some(namedPivotCol match {
+ case ne: NamedExpression => ne.name
+ case _ => "__pivot_col"
+ }))
+ } else {
+ namedPivotCol
+ }
+ val extendedGroupingExpressions =
+ groupByExpressions :+ effectivePivotCol
Review Comment:
**Finding 2.** With this key, `firstAgg` can need a shuffle that base didn't
need.
Base groups `firstAgg` by the pivot attribute itself. A child that is
already hash-partitioned by that column satisfies
`ClusteredDistribution(groupBy :+ course)`, so the final `firstAgg` runs
without an exchange. Here the key is a new attribute (`_groupingexpression`
after `PullOutGroupingExpressions`). No child partitioning covers it, so
`EnsureRequirements` adds an exchange.
Shuffle counts for `groupBy("year").pivot("course", Seq("dotNET",
"Java")).agg(sum("earnings"))`:
| child | base | this PR |
|---|---|---|
| table bucketed by `course` (AQE on or off) | 1 | 2 |
| `repartition($"course")` (AQE off) | 2 | 3 |
| join on `course` (AQE off) | 3 | 4 |
A pivot over a subquery grouped by the pivot column has the same shape. When
every row matches a pivot value there is nothing to collapse, so the extra
stage is pure cost. In that case a struct pivot column (the `FOR (a, b) IN`
form) was also about 10% slower in each of three runs (30M rows, local[2]).
Long and string pivot columns were within noise.
The win in the target case is large (11.7 s vs 0.4 s, see the summary), so
this is a trade-off. Could you put it behind an internal `SQLConf`? Users with
a pre-partitioned child could then turn it off. Adding these numbers to the
description would also help decide the default.
##########
sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala:
##########
@@ -654,4 +654,107 @@ class DataFramePivotSuite extends SharedSparkSession {
checkAnswer(nanDf, Row(null, null, 20L))
}
}
+
+ test("SPARK-55569: pivot pushes a filter on the pivot column values") {
+ val df = Seq(
+ ("dotNET", 2012, 10000),
+ ("Java", 2012, 20000),
+ ("dotNET", 2013, 48000),
+ ("Java", 2013, 30000),
+ ("C#", 2013, 15000)
+ ).toDF("course", "year", "earnings")
+
+ val pivoted = df
+ .groupBy("year")
+ .pivot("course", Seq("dotNET", "Java"))
+ .agg(sum("earnings"))
+
+ val analyzed = pivoted.queryExecution.analyzed
+ // The fast-path group-by key should contain an If that collapses
+ // non-matching pivot column values to null.
+ assert(analyzed.exists { node =>
+ node.expressions.exists(_.exists {
+ case catalyst.expressions.If(_, _,
+ catalyst.expressions.Literal(null, _)) => true
+ case _ => false
+ })
+ }, s"Expected If(..., col, null):\n${analyzed.treeString}")
+
+ // Results are identical to the un-optimized plan.
+ checkAnswer(pivoted,
+ Seq(Row(2012, 10000, 20000), Row(2013, 48000, 30000)))
+ }
+
+ test("SPARK-55569: NULL pivot value skips the collapse optimization") {
+ // When a NULL pivot value is present, the collapse optimization is
+ // skipped because non-matching rows mapped to null would merge with
+ // legitimate null-key rows. Results must still be correct.
+ val df = Seq(
+ (Some("a"), 1), (Some("b"), 2), (None, 3)
+ ).toDF("key", "value")
+
+ val pivoted = df.groupBy().pivot("key", Seq("a", null))
+ .agg(sum("value"))
+
+ // a=1 (key='a'), null=3 (key IS NULL)
+ checkAnswer(pivoted, Row(1, 3))
+ }
+
+ test("SPARK-55569: pivot filter is not pushed for aggregate pivot column") {
Review Comment:
**Finding 6.** This test still passes with the aggregate guard removed. It
only checks the error condition, and the condition is `GROUP_BY_AGGREGATE`
either way. The guard is already pinned by the existing `SPARK-24722: aggregate
as the pivot column` test. That test checks `sqlExpr`, which turns into the
`IF(...)` expression without the guard. I'd drop this test, along with the
matching line in the description.
##########
sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala:
##########
@@ -654,4 +654,107 @@ class DataFramePivotSuite extends SharedSparkSession {
checkAnswer(nanDf, Row(null, null, 20L))
}
}
+
+ test("SPARK-55569: pivot pushes a filter on the pivot column values") {
Review Comment:
**Finding 7.** This name, and the ones at lines 703 and 719, say that the PR
pushes a filter. It doesn't, so someone looking for a pushdown test would be
misled. Something like `collapses non-matching pivot values in the first
aggregate` fits better.
##########
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,
+ Literal(null, pivotColumn.dataType)),
+ Some(namedPivotCol match {
+ case ne: NamedExpression => ne.name
+ case _ => "__pivot_col"
+ }))
Review Comment:
**Finding 5.** A few branches here can't fire.
- `canCollapse` already excludes null pivot values, so the `v != null`
filter keeps everything.
- `namedPivotCol` is a `NamedExpression`, so `case _ => "__pivot_col"` is
unreachable.
- `In(pivotColumn, Nil)` evaluates to false (or null under the legacy flag),
so the `FalseLiteral` branch isn't needed for an empty value list either.
```suggestion
val matchesPivotValue = In(pivotColumn, pivotValues.map { value =>
Cast(value, pivotColumn.dataType, Some(conf.sessionLocalTimeZone))
})
newAlias(
If(matchesPivotValue, pivotColumn, Literal(null,
pivotColumn.dataType)),
Some(namedPivotCol.name))
```
With this change `DataFramePivotSuite` and the pivot golden files pass
unchanged, and an empty value list still works.
--
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]
