support order by index and group by index

refactor

style didn't catched in local

style

forgot to enbale style checking for test in local

override resolved

refactor

a version

refactor

Author: zhichao-li

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

@@ -71,6 +71,7 @@ class Analyzer(
     Batch("Resolution", fixedPoint,
       ResolveRelations ::
       ResolveReferences ::
+      ResolveIndexReferences ::
       ResolveGroupingAnalytics ::
       ResolvePivot ::
       ResolveUpCast ::
@@ -321,6 +322,62 @@ class Analyzer(
     }
   }
 
+  /**
+    * Replace the index with the related attribute for order by and group by.
+    */
+  object ResolveIndexReferences extends Rule[LogicalPlan] {
+
+    def indexToColumn(
+        sortOrder: SortOrder,
+        child: LogicalPlan,
+        index: Int,
+        direction: SortDirection): SortOrder = {
+      val orderNodes = child.output
+      if (index > 0 && index <= orderNodes.size) {
+        SortOrder(orderNodes(index - 1), direction)
+      } else {
+        throw new UnresolvedException(sortOrder,
+          s"""Order by position: $index does not exist \n
+              |The Select List is indexed from 1 to ${orderNodes.size}""".stripMargin)
+      }
+    }
+
+    def indexToColumn(
+        agg: Aggregate,
+        child: LogicalPlan,
+        index: Int): Expression = {
+      val attributes = child.output
+      if (index > 0 && index <= attributes.size) {
+        attributes(index - 1)
+      } else {
+        throw new UnresolvedException(agg,
+          s"""Order by position: $index does not exist \n
+              |The Select List is indexed from 1 to ${attributes.size}""".stripMargin)
+      }
+    }
+
+    def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
+      case s @ Sort(orders, global, child) if child.resolved =>
+        val newOrders = orders map {
+          case s @ SortOrder(IntegerLiteral(index), direction) =>
+            indexToColumn(s, child, index, direction)
+          case s @ SortOrder(UnaryMinus(IntegerLiteral(index)), direction) =>
+            indexToColumn(s, child, -index, direction)
+          case other => other
+        }
+        Sort(newOrders, global, child)
+      case a @ Aggregate(groups, aggs, child) =>
+        val newGroups = groups map {
+          case IntegerLiteral(index) =>
+            indexToColumn(a, child, index)
+          case UnaryMinus(IntegerLiteral(index)) =>
+            indexToColumn(a, child, -index)
+          case other => other
+        }
+        Aggregate(newGroups, aggs, child)
+    }
+  }
+
   /**
    * Replaces [[UnresolvedAttribute]]s with concrete [[AttributeReference]]s from
    * a logical plan node's children.
@@ -446,6 +503,10 @@ class Analyzer(
         val newOrdering = resolveSortOrders(ordering, child, throws = false)
         Sort(newOrdering, global, child)
 
+      case s @ Sort(ordering, global, child) if child.resolved && !s.resolved =>
+        val newOrdering = resolveSortOrders(ordering, child, throws = false)
+        Sort(newOrdering, global, child)
+
       // A special case for Generate, because the output of Generate should not be resolved by
       // ResolveReferences. Attributes in the output will be resolved by ResolveGenerate.
       case g @ Generate(generator, join, outer, qualifier, output, child)

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

@@ -21,10 +21,9 @@ import java.math.MathContext
 import java.sql.Timestamp
 
 import org.apache.spark.AccumulatorSuite
-import org.apache.spark.sql.catalyst.CatalystQl
-import org.apache.spark.sql.catalyst.analysis.FunctionRegistry
-import org.apache.spark.sql.catalyst.parser.ParserConf
-import org.apache.spark.sql.execution.{aggregate, SparkQl}
+import org.apache.spark.sql.catalyst.analysis.{FunctionRegistry, UnresolvedException}
+import org.apache.spark.sql.catalyst.expressions.SortOrder
+import org.apache.spark.sql.execution.aggregate
 import org.apache.spark.sql.execution.joins.{CartesianProduct, SortMergeJoin}
 import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
 import org.apache.spark.sql.functions._
@@ -456,21 +455,36 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
 
   test("literal in agg grouping expressions") {
     checkAnswer(
-      sql("SELECT a, count(1) FROM testData2 GROUP BY a, 1"),
-      Seq(Row(1, 2), Row(2, 2), Row(3, 2)))
+      sql("SELECT SUM(a) FROM testData2 GROUP BY 2"),
+      Seq(Row(6), Row(6)))
+
     checkAnswer(
-      sql("SELECT a, count(2) FROM testData2 GROUP BY a, 2"),
+      sql("SELECT a, SUM(b) FROM testData2 GROUP BY 1"),
+      Seq(Row(1, 3), Row(2, 3), Row(3, 3)))
+
+    checkAnswer(
+      sql("SELECT a, SUM(b) FROM testData2 GROUP BY 1"),
+      Seq(Row(1, 3), Row(2, 3), Row(3, 3)))
+
+    checkAnswer(
+      sql("SELECT a, count(1) FROM testData2 GROUP BY a, 1"),
       Seq(Row(1, 2), Row(2, 2), Row(3, 2)))
 
     checkAnswer(
       sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a, 1"),
       sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a"))
+
     checkAnswer(
-      sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a, 1 + 2"),
-      sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a"))
+      sql("SELECT a, count(2) FROM testData2 GROUP BY a, 2"),
+      Seq(Row(1, 1), Row(1, 1), Row(2, 1), Row(2, 1), Row(3, 1), Row(3, 1)))
+
     checkAnswer(
       sql("SELECT 1, 2, sum(b) FROM testData2 GROUP BY 1, 2"),
-      sql("SELECT 1, 2, sum(b) FROM testData2"))
+      sql("SELECT 1, 2, sum(b) FROM testData2 GROUP BY a, b"))
+
+    checkAnswer(
+      sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a, 1 + 2"),
+      sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a"))
   }
 
   test("aggregates with nulls") {
@@ -493,6 +507,21 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
       Row("1"))
   }
 
+  test("order by number index")  {
+    intercept[UnresolvedException[SortOrder]] {
+      sql("SELECT * FROM testData2 ORDER BY -1 DESC, b ASC").collect()
+    }
+    checkAnswer(
+      sql("SELECT * FROM testData2 ORDER BY 1 DESC"),
+      Seq(Row(3, 1), Row(3, 2), Row(2, 1), Row(2, 2), Row(1, 1), Row(1, 2)))
+    checkAnswer(
+      sql("SELECT * FROM testData2 ORDER BY 1 DESC, b ASC"),
+      sql("SELECT * FROM testData2 ORDER BY a DESC, b ASC"))
+    checkAnswer(
+      sql("SELECT * FROM testData2 ORDER BY 1 ASC, b ASC"),
+      Seq(Row(1, 1), Row(1, 2), Row(2, 1), Row(2, 2), Row(3, 1), Row(3, 2)))
+  }
+
   def sortTest(): Unit = {
     checkAnswer(
       sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC"),