[CARMEL-6057] Fix limit push down through window (#1016)

* Fix limit push down through window

* fix

* fix

* fix

* flaky test

Author: wangyum

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/LimitPushDownThroughWindow.scala

@@ -61,7 +61,7 @@ object LimitPushDownThroughWindow extends Rule[LogicalPlan] with AliasHelper {
       attr: Attribute, order: Seq[SortOrder], window: Window): Boolean = {
     window.windowExpressions.size == 1 &&
       getAliasMap(window.windowExpressions).get(attr).exists(a => pushableWindowExprs(a :: Nil)) &&
-      (order.isEmpty || (window.partitionSpec.nonEmpty &&
+      (order.isEmpty || (window.partitionSpec.nonEmpty && window.partitionSpec.size >= order.size &&
         window.partitionSpec.zip(order).forall {
           case (e: Expression, s: SortOrder) => s.child.semanticEquals(e)
         }))
@@ -76,7 +76,8 @@ object LimitPushDownThroughWindow extends Rule[LogicalPlan] with AliasHelper {
   private def pushedOrder(order: Seq[SortOrder], window: Window): Seq[SortOrder] = {
     val windowSize = window.partitionSpec.size
     order.take(windowSize) ++
-      sortWindowSpec(window.partitionSpec.takeRight(windowSize - order.size), window.orderSpec)
+      sortWindowSpec(window.partitionSpec.takeRight(windowSize - order.size), window.orderSpec) ++
+      order.takeRight(order.size - windowSize).filter(_.references.subsetOf(window.child.outputSet))
   }
 
   def apply(plan: LogicalPlan): LogicalPlan = plan.transform {
@@ -105,13 +106,6 @@ object LimitPushDownThroughWindow extends Rule[LogicalPlan] with AliasHelper {
         limit < conf.topKSortFallbackThreshold && supportsPushDown(attr, Nil, window) =>
       val newWindowChild = Limit(limitExpr, Sort(pushedOrder(Nil, window), true, window.child))
       f.copy(child = window.copy(child = newWindowChild))
-
-    case Limit(limitExpr @ IntegerLiteral(limit), s @ Sort(order, _,
-    f @ Filter(ExtractFilterCondition(attr, filterVal), window: Window)))
-      if limit < filterVal && window.maxRows.forall(_ > limit) &&
-        limit < conf.topKSortFallbackThreshold && supportsPushDown(attr, order, window) =>
-      val newWindowChild = Limit(limitExpr, Sort(pushedOrder(order, window), true, window.child))
-      s.copy(child = f.copy(child = window.copy(child = newWindowChild)))
   }
 
   private object ExtractFilterCondition {

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/LimitPushdownThroughWindowSuite.scala

@@ -228,8 +228,8 @@ class LimitPushdownThroughWindowSuite extends PlanTest {
   }
 
   test("SPARK-39600: Limit should LessThan/LessThanOrEqual filter value") {
-    Seq(3, 4, 5, 6).foreach { limitVal =>
-      Seq(3, 4, 5, 6).foreach { filterVal =>
+    Seq(40, 80, 120).foreach { limitVal =>
+      Seq(40, 80, 120).foreach { filterVal =>
         val originalQuery = testRelation
           .select(a, b, c,
             windowExpr(RowNumber(), windowSpec(a :: Nil, c.desc :: Nil, windowFrame)).as("rn"))
@@ -242,35 +242,7 @@ class LimitPushdownThroughWindowSuite extends PlanTest {
             windowExpr(RowNumber(), windowSpec(a :: Nil, c.desc :: Nil, windowFrame)).as("rn"))
           .where('rn < Literal(filterVal))
 
-        if (limitVal < filterVal) {
-          comparePlans(
-            Optimize.execute(originalQuery.analyze),
-            WithoutOptimize.execute(correctAnswer.analyze))
-        } else {
-          comparePlans(
-            Optimize.execute(originalQuery.analyze),
-            WithoutOptimize.execute(originalQuery.analyze))
-        }
-      }
-    }
-
-    Seq(3, 4, 5, 6).foreach { limitVal =>
-      Seq(3, 4, 5, 6).foreach { filterVal =>
-        val originalQuery = testRelation
-          .select(a, b, c,
-            windowExpr(RowNumber(), windowSpec(a :: Nil, c.desc :: Nil, windowFrame)).as("rn"))
-          .where('rn <= Literal(filterVal))
-          .sortBy(a.asc, b.asc)
-          .limit(limitVal)
-        val correctAnswer = testRelation
-          .orderBy(a.asc, c.desc)
-          .limit(limitVal)
-          .select(a, b, c,
-            windowExpr(RowNumber(), windowSpec(a :: Nil, c.desc :: Nil, windowFrame)).as("rn"))
-          .where('rn <= Literal(filterVal))
-          .sortBy(a.asc, b.asc)
-
-        if (limitVal <= filterVal) {
+        if (limitVal < filterVal && limitVal < 100) {
           comparePlans(
             Optimize.execute(originalQuery.analyze),
             WithoutOptimize.execute(correctAnswer.analyze))
@@ -283,26 +255,6 @@ class LimitPushdownThroughWindowSuite extends PlanTest {
     }
   }
 
-  test("SPARK-39600: Push down if partitionSpec is the prefix of sort exprs") {
-    val originalQuery = testRelation
-      .select(a, b, c,
-        windowExpr(RowNumber(), windowSpec(b :: Nil, c.desc :: Nil, windowFrame)).as("rn"))
-      .where('rn <= Literal(5))
-      .orderBy(b.asc)
-      .limit(5)
-    val correctAnswer = testRelation
-      .orderBy(b.asc, c.desc)
-      .limit(5)
-      .select(a, b, c,
-        windowExpr(RowNumber(), windowSpec(b :: Nil, c.desc :: Nil, windowFrame)).as("rn"))
-      .where('rn <= Literal(5))
-      .orderBy(b.asc)
-
-    comparePlans(
-      Optimize.execute(originalQuery.analyze),
-      WithoutOptimize.execute(correctAnswer.analyze))
-  }
-
   test("SPARK-39600: Push down if sort exprs is not exist") {
     val originalQuery = testRelation
       .select(a, b, c,
@@ -321,46 +273,6 @@ class LimitPushdownThroughWindowSuite extends PlanTest {
       WithoutOptimize.execute(correctAnswer.analyze))
   }
 
-  test("SPARK-39600: Push down should keep the origin order") {
-    val originalQuery = testRelation
-      .select(a, b, c,
-        windowExpr(RowNumber(), windowSpec(b :: Nil, c.desc :: Nil, windowFrame)).as("rn"))
-      .where('rn <= Literal(5))
-      .orderBy(b.desc, c.asc, a.desc)
-      .limit(5)
-    val correctAnswer = testRelation
-      .orderBy(b.desc, c.desc)
-      .limit(5)
-      .select(a, b, c,
-        windowExpr(RowNumber(), windowSpec(b :: Nil, c.desc :: Nil, windowFrame)).as("rn"))
-      .where('rn <= Literal(5))
-      .orderBy(b.desc, c.asc, a.desc)
-
-    comparePlans(
-      Optimize.execute(originalQuery.analyze),
-      WithoutOptimize.execute(correctAnswer.analyze))
-  }
-
-  test("SPARK-39600: Order exprs is subset of partitionSpec") {
-    val originalQuery = testRelation
-      .select(a, b, c,
-        windowExpr(RowNumber(), windowSpec(a :: b :: Nil, c.desc :: Nil, windowFrame)).as("rn"))
-      .where('rn <= Literal(5))
-      .orderBy(a.desc)
-      .limit(5)
-    val correctAnswer = testRelation
-      .orderBy(a.desc, b.asc, c.desc)
-      .limit(5)
-      .select(a, b, c,
-        windowExpr(RowNumber(), windowSpec(a :: b :: Nil, c.desc :: Nil, windowFrame)).as("rn"))
-      .where('rn <= Literal(5))
-      .orderBy(a.desc)
-
-    comparePlans(
-      Optimize.execute(originalQuery.analyze),
-      WithoutOptimize.execute(correctAnswer.analyze))
-  }
-
   test("SPARK-39600: Unsupported case: do not push down if partitionSpec is empty") {
     val originalQuery = testRelation
       .select(a, b, c,

sql/core/src/test/scala/org/apache/spark/sql/sources/DisableUnnecessaryBucketedScanSuite.scala

@@ -377,7 +377,7 @@ abstract class DisableUnnecessaryBucketedScanSuite extends QueryTest
           """
          SELECT * FROM
          (SELECT j, row_number() OVER (partition by j order by i) rn, i
-         FROM t1 order by j) t
+         FROM t1 order by j, rn) t
          WHERE rn <= 3
          """.stripMargin, 0, 1),