Address remaining comments

- Use Literal's for the pivot column values instead of strings.
- Change seperator when using multiple aggregates to `_` instead of space.
- Some additional unit testing

Author: aray

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -260,7 +260,7 @@ class Analyzer(
         val singleAgg = aggregates.size == 1
         val pivotAggregates: Seq[NamedExpression] = pivotValues.flatMap { value =>
           def ifExpr(expr: Expression) = {
-            If(EqualTo(pivotColumn, Literal(value)), expr, Literal(null))
+            If(EqualTo(pivotColumn, value), expr, Literal(null))
           }
           aggregates.map { aggregate =>
             val filteredAggregate = aggregate.transformDown {
@@ -278,7 +278,7 @@ class Analyzer(
               throw new AnalysisException(
                 s"Aggregate expression required for pivot, found '$aggregate'")
             }
-            val name = if (singleAgg) value else value + " " + aggregate.prettyString
+            val name = if (singleAgg) value.toString else value + "_" + aggregate.prettyString
             Alias(filteredAggregate, name)()
           }
         }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala

@@ -388,13 +388,13 @@ case class Rollup(
 case class Pivot(
     groupByExprs: Seq[NamedExpression],
     pivotColumn: Expression,
-    pivotValues: Seq[String],
+    pivotValues: Seq[Literal],
     aggregates: Seq[Expression],
     child: LogicalPlan) extends UnaryNode {
   override def output: Seq[Attribute] = groupByExprs.map(_.toAttribute) ++ aggregates match {
-    case aggregate :: Nil => pivotValues.map(AttributeReference(_, aggregate.dataType)())
+    case agg :: Nil => pivotValues.map(value => AttributeReference(value.toString, agg.dataType)())
     case _ => pivotValues.flatMap{ value =>
-      aggregates.map(agg => AttributeReference(value + " " + agg.prettyString, agg.dataType)())
+      aggregates.map(agg => AttributeReference(value + "_" + agg.prettyString, agg.dataType)())
     }
   }
 }

sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala

@@ -297,18 +297,25 @@ class GroupedData protected[sql](
     * @since 1.6.0
     */
   @scala.annotation.varargs
-  def pivot(pivotColumn: Column, values: String*): GroupedData = groupType match {
+  def pivot(pivotColumn: Column, values: Column*): GroupedData = groupType match {
     case _: GroupedData.PivotType =>
       throw new UnsupportedOperationException("repeated pivots are not supported")
     case GroupedData.GroupByType =>
       val pivotValues = if (values.nonEmpty) {
-        values
+        values.map {
+          case Column(literal: Literal) => literal
+          case other =>
+            throw new UnsupportedOperationException(
+              s"The values of a pivot must be literals, found $other")
+        }
       } else {
         // Get the distinct values of the column and sort them so its consistent
-        df.select(pivotColumn.cast(StringType))
+        df.select(pivotColumn)
           .distinct()
-          .map(_.getString(0))
-          .collect().sorted.toSeq
+          .sort(pivotColumn)
+          .map(_.get(0))
+          .collect()
+          .map(Literal(_)).toSeq
       }
       new GroupedData(df, groupingExprs, GroupedData.PivotType(pivotColumn.expr, pivotValues))
     case _ =>
@@ -330,9 +337,9 @@ class GroupedData protected[sql](
     * @since 1.6.0
     */
   @scala.annotation.varargs
-  def pivot(pivotColumn: String, values: String*): GroupedData = {
+  def pivot(pivotColumn: String, values: Any*): GroupedData = {
     val resolvedPivotColumn = Column(df.resolve(pivotColumn))
-    pivot(resolvedPivotColumn, values: _*)
+    pivot(resolvedPivotColumn, values.map(functions.lit): _*)
   }
 }
 
@@ -372,5 +379,5 @@ private[sql] object GroupedData {
   /**
     * To indicate it's the PIVOT
     */
-  private[sql] case class PivotType(pivotCol: Expression, values: Seq[String]) extends GroupType
+  private[sql] case class PivotType(pivotCol: Expression, values: Seq[Literal]) extends GroupType
 }

sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala

@@ -23,45 +23,53 @@ import org.apache.spark.sql.test.SharedSQLContext
 class DataFramePivotSuite extends QueryTest with SharedSQLContext{
   import testImplicits._
 
-  test("pivot courses groupBy") {
+  test("pivot courses with literals") {
     checkAnswer(
-      courseSales.groupBy($"year").pivot($"course", "dotNET", "Java").agg(sum($"earnings")),
+      courseSales.groupBy($"year").pivot($"course", lit("dotNET"), lit("Java"))
+        .agg(sum($"earnings")),
       Row(2012, 15000.0, 20000.0) :: Row(2013, 48000.0, 30000.0) :: Nil
     )
   }
 
-  test("pivot year groupBy") {
+  test("pivot year with literals") {
     checkAnswer(
-      courseSales.groupBy($"course").pivot($"year", "2012", "2013").agg(sum($"earnings")),
+      courseSales.groupBy($"course").pivot($"year", lit(2012), lit(2013)).agg(sum($"earnings")),
       Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil
     )
   }
 
-  test("pivot courses groupBy multiple") {
+  test("pivot courses with literals and multiple aggregations") {
     checkAnswer(
-      courseSales.groupBy($"year").pivot($"course", "dotNET", "Java")
+      courseSales.groupBy($"year").pivot($"course", lit("dotNET"), lit("Java"))
         .agg(sum($"earnings"), avg($"earnings")),
       Row(2012, 15000.0, 7500.0, 20000.0, 20000.0) ::
         Row(2013, 48000.0, 48000.0, 30000.0, 30000.0) :: Nil
     )
   }
 
-  test("pivot year groupBy with strings") {
+  test("pivot year with string values (cast)") {
     checkAnswer(
       courseSales.groupBy("course").pivot("year", "2012", "2013").sum("earnings"),
       Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil
     )
   }
 
-  test("pivot courses groupBy with no values") {
+  test("pivot year with int values") {
+    checkAnswer(
+      courseSales.groupBy("course").pivot("year", 2012, 2013).sum("earnings"),
+      Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil
+    )
+  }
+
+  test("pivot courses with no values") {
     // Note Java comes before dotNet in sorted order
     checkAnswer(
       courseSales.groupBy($"year").pivot($"course").agg(sum($"earnings")),
       Row(2012, 20000.0, 15000.0) :: Row(2013, 30000.0, 48000.0) :: Nil
     )
   }
 
-  test("pivot year groupBy with no values") {
+  test("pivot year with no values") {
     checkAnswer(
       courseSales.groupBy($"course").pivot($"year").agg(sum($"earnings")),
       Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil