Skip to content
Snippets Groups Projects
Commit f1fca81b authored by Andrew Ray's avatar Andrew Ray Committed by Herman van Hovell
Browse files

[SPARK-17760][SQL] AnalysisException with dataframe pivot when groupBy column is not attribute

## What changes were proposed in this pull request?

Fixes AnalysisException for pivot queries that have group by columns that are expressions and not attributes by substituting the expressions output attribute in the second aggregation and final projection.

## How was this patch tested?

existing and additional unit tests

Author: Andrew Ray <ray.andrew@gmail.com>

Closes #16177 from aray/SPARK-17760.
parent c496d03b
No related branches found
No related tags found
No related merge requests found
......@@ -463,14 +463,15 @@ class Analyzer(
.toAggregateExpression()
, "__pivot_" + a.sql)()
}
val secondAgg = Aggregate(groupByExprs, groupByExprs ++ pivotAggs, firstAgg)
val groupByExprsAttr = groupByExprs.map(_.toAttribute)
val secondAgg = Aggregate(groupByExprsAttr, groupByExprsAttr ++ pivotAggs, firstAgg)
val pivotAggAttribute = pivotAggs.map(_.toAttribute)
val pivotOutputs = pivotValues.zipWithIndex.flatMap { case (value, i) =>
aggregates.zip(pivotAggAttribute).map { case (aggregate, pivotAtt) =>
Alias(ExtractValue(pivotAtt, Literal(i), resolver), outputName(value, aggregate))()
}
}
Project(groupByExprs ++ pivotOutputs, secondAgg)
Project(groupByExprsAttr ++ pivotOutputs, secondAgg)
} else {
val pivotAggregates: Seq[NamedExpression] = pivotValues.flatMap { value =>
def ifExpr(expr: Expression) = {
......
......@@ -208,4 +208,12 @@ class DataFramePivotSuite extends QueryTest with SharedSQLContext{
)
}
test("pivot with column definition in groupby") {
checkAnswer(
courseSales.groupBy(substring(col("course"), 0, 1).as("foo"))
.pivot("year", Seq(2012, 2013))
.sum("earnings"),
Row("d", 15000.0, 48000.0) :: Row("J", 20000.0, 30000.0) :: Nil
)
}
}
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment