add fix for sorting also

prakharjain09 · prakharjain09 · commit d5a0fbe779a6 · 2020-11-12T16:22:07.000+05:30
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala
@@ -31,8 +31,10 @@ trait AliasAwareOutputExpression extends UnaryExecNode {
 
   protected def hasAlias: Boolean = aliasMap.nonEmpty
 
-  protected def replaceAlias(attr: AttributeReference): Option[Attribute] = {
-    aliasMap.get(attr)
+  protected def normalizeExpression(exp: Expression): Expression = {
+    exp.transform {
+      case attr: AttributeReference => aliasMap.get(attr).getOrElse(attr)
+    }
   }
 }
 
@@ -45,10 +47,7 @@ trait AliasAwareOutputPartitioning extends AliasAwareOutputExpression {
     if (hasAlias) {
       child.outputPartitioning match {
         case e: Expression =>
-          val normalizedExpr = e.transform {
-            case attr: AttributeReference => replaceAlias(attr).getOrElse(attr)
-          }
-          normalizedExpr.asInstanceOf[Partitioning]
+          normalizeExpression(e).asInstanceOf[Partitioning]
         case other => other
       }
     } else {
@@ -66,11 +65,8 @@ trait AliasAwareOutputOrdering extends AliasAwareOutputExpression {
 
   final override def outputOrdering: Seq[SortOrder] = {
     if (hasAlias) {
-      orderingExpressions.map { s =>
-        s.child match {
-          case a: AttributeReference => s.copy(child = replaceAlias(a).getOrElse(a))
-          case _ => s
-        }
+      orderingExpressions.map { sortOrder =>
+        normalizeExpression(sortOrder).asInstanceOf[SortOrder]
       }
     } else {
       orderingExpressions
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
@@ -953,6 +953,31 @@ class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("SPARK-33399: aliased should be handled properly " +
+    "for partitioning and sortorder involving complex expressions") {
+    withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+      withTempView("t1", "t2", "t3") {
+        spark.range(10).createTempView("t1")
+        spark.range(20).createTempView("t2")
+        spark.range(30).createTempView("t3")
+        val planned = sql(
+          """
+            |SELECT t3.id as t3id
+            |FROM (
+            |    SELECT t1.id as t1id
+            |    FROM t1, t2
+            |    WHERE t1.id % 10 = t2.id % 10
+            |) t12, t3
+            |WHERE t1id % 10 = t3.id % 10
+          """.stripMargin).queryExecution.executedPlan
+        val sortNodes = planned.collect { case s: SortExec => s }
+        assert(sortNodes.size == 3)
+        val exchangeNodes = planned.collect { case e: ShuffleExchangeExec => e }
+        assert(exchangeNodes.size == 3)
+      }
+    }
+  }
+
   test("SPARK-33399: alias handling should happen properly for SinglePartition") {
     withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
       val df = spark.range(1, 100, 1, 1)
