Add Sort() case

Author: Derek Sabry

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

@@ -70,7 +70,7 @@ class Analyzer(
     Batch("Resolution", fixedPoint,
       ResolveRelations ::
       ResolveReferences ::
-      ResolveGroupByReferences ::
+      ResolveNumericReferences ::
       ResolveGroupingAnalytics ::
       ResolvePivot ::
       ResolveSortReferences ::
@@ -180,23 +180,42 @@ class Analyzer(
   }
 
   /**
-   * Replaces queries of the form "SELECT expr FROM A GROUP BY 1"
-   * with a query of the form "SELECT expr FROM A GROUP BY expr"
+   * Replaces queries of the form "SELECT expr FROM A GROUP BY 1 ORDER BY 1"
+   * with a query of the form "SELECT expr FROM A GROUP BY expr ORDER BY expr"
    */
-  object ResolveGroupByReferences extends Rule[LogicalPlan] {
+  object ResolveNumericReferences extends Rule[LogicalPlan] {
 
     def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
       case Aggregate(groups, aggs, child) =>
-        Aggregate(groups.map(group => group match {
-          case g if g.prettyString forall Character.isDigit =>
-            aggs(g.prettyString.toInt - 1) match {
-              case u : UnresolvedAlias =>
-                u.child
-              case a : Alias =>
-                a.child
+        val newGroups = groups.map {
+          case group if group.isInstanceOf[Literal] && group.dataType.isInstanceOf[IntegralType] =>
+            aggs(group.toString.toInt - 1) match {
+              case u: UnresolvedAlias =>
+                u.child match {
+                  case UnresolvedStar(_) => // Can't replace literal with column yet
+                    group
+                  case _ => u.child
+                }
+              case a: Alias => a.child
+              case a: AttributeReference => a
             }
-          case _ => group
-        }), aggs, child)
+          case group => group
+        }
+        Aggregate(newGroups, aggs, child)
+      case Sort(ordering, global, child) =>
+        val newOrdering = ordering.map {
+          case o if o.child.isInstanceOf[Literal] && o.dataType.isInstanceOf[IntegralType] =>
+            val newExpr = child.asInstanceOf[Project].projectList(o.child.toString.toInt - 1)
+                match {
+                  case u: UnresolvedAlias =>
+                    u.child
+                  case a: Alias =>
+                    a.child
+                }
+            SortOrder(newExpr, o.direction)
+          case other => other
+        }
+       Sort(newOrdering, global, child)
     }
   }
 

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

@@ -2028,10 +2028,17 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
       Row(false) :: Row(true) :: Nil)
   }
 
-  test("SPARK-12063: Group by Column Number identifier") {
+  test("SPARK-12063: Group by Columns Number") {
     checkAnswer(
       sql("SELECT a, SUM(b) FROM testData2 GROUP BY 1"),
       Seq(Row(1, 3), Row(2, 3), Row(3, 3)))
   }
+  
+  test("SPARK-12063: Order by Column Number") {
+    Seq(("one", 1), ("two", 2), ("three", 3), ("one", 5)).toDF("k", "v").registerTempTable("ord")
+    checkAnswer(
+      sql("SELECT v from ord order by 1 desc"),
+      Row(5) :: Row(3) :: Row(2) :: Row(1) :: Nil)
+  }
 
 }