Skip to content

Commit

Permalink
[SPARK-17760][SQL] AnalysisException with dataframe pivot when groupB…
Browse files Browse the repository at this point in the history
…y 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 <[email protected]>

Closes apache#16177 from aray/SPARK-17760.
  • Loading branch information
aray authored and hvanhovell committed Dec 7, 2016
1 parent c496d03 commit f1fca81
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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) = {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 comments on commit f1fca81

Please sign in to comment.