shrirangmhalgi opened a new pull request, #59149: URL: https://github.com/apache/spark/pull/59149
### What changes were proposed in this pull request? In the PivotFirst fast path, `firstAgg` groups by `(groupByKeys, pivotColumn)`. When the table has many distinct pivot column values outside the explicit `IN` list, `firstAgg` creates one group per distinct non-matching value - all of which `PivotFirst` ignores in `secondAgg`. This PR wraps the pivot column group-by key in `IF(pivotColumn IN (pivotValues), pivotColumn, null)`, collapsing all non-matching values into a single null group. `PivotFirst` skips the null group as before, and the group-by keys are preserved in `secondAgg`, keeping query results identical to the un-optimized plan. The optimization is skipped when: - The pivot column contains an aggregate expression (SPARK-24722) - A NULL pivot value is present (non-matching rows mapped to null would merge with legitimate null-key rows) ### Why are the changes needed? For a table with 10,000 distinct course values pivoting on 2: - Before: `firstAgg` creates 10,000 groups per group-by key (99.98% wasted) - After: `firstAgg` creates 3 groups per group-by key (2 matching + 1 null) This reduces memory (fewer HashAggregate entries), CPU (fewer hash/compare/aggregate operations), and spill risk. ### Does this PR introduce *any* user-facing change? No. Query results are identical. This only affects the internal execution plan for the PivotFirst fast path. ### How was this patch tested? Five new tests in `DataFramePivotSuite`: - Plan-level assertion that `If(In(...), col, null)` appears in the analyzed plan - NULL pivot value correctly skips the optimization - Aggregate pivot column (SPARK-24722) correctly skips the optimization - Multiple aggregates (sum + avg) with the optimization - Group-by key with only non-pivot values is preserved as `Row(year, null, null)` Existing golden file tests regenerated (`pivot.sql`, `udf-pivot.sql`, `pipe-operators.sql`). 49/49 `DataFramePivotSuite` tests pass. ### Was this patch authored or co-authored using generative AI tooling? Yes. Co-authored using Claude Opus 4.8. -- 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]
