[SQL][SPARK-51635] Merge PushProjectionThroughLimit and PushProjectionThroughOffset

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

@@ -100,8 +100,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
       Seq(
         // Operator push down
         PushProjectionThroughUnion,
-        PushProjectionThroughLimit,
-        PushProjectionThroughOffset,
+        PushProjectionThroughLimitAndOffset,
         ReorderJoin,
         EliminateOuterJoin,
         PushDownPredicates,

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushProjectionThroughLimit.scala → sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushProjectionThroughLimitAndOffset.scala

@@ -17,16 +17,16 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
-import org.apache.spark.sql.catalyst.plans.logical.{GlobalLimit, LocalLimit, LogicalPlan, Project}
+import org.apache.spark.sql.catalyst.plans.logical.{GlobalLimit, LocalLimit, LogicalPlan, Offset, Project}
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.catalyst.trees.TreePattern.{LIMIT, PROJECT}
+import org.apache.spark.sql.catalyst.trees.TreePattern.{LIMIT, OFFSET, PROJECT}
 
 /**
  * Pushes Project operator through Limit operator.
  */
-object PushProjectionThroughLimit extends Rule[LogicalPlan] {
-  def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning(
-    _.containsAllPatterns(PROJECT, LIMIT)) {
+object PushProjectionThroughLimitAndOffset extends Rule[LogicalPlan] {
+  def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning(treeBits =>
+    treeBits.containsPattern(PROJECT) && treeBits.containsAnyPattern(OFFSET, LIMIT)) {
 
     case p @ Project(projectList, limit @ LocalLimit(_, child))
         if projectList.forall(_.deterministic) =>
@@ -35,5 +35,9 @@ object PushProjectionThroughLimit extends Rule[LogicalPlan] {
     case p @ Project(projectList, g @ GlobalLimit(_, limit @ LocalLimit(_, child)))
         if projectList.forall(_.deterministic) =>
       g.copy(child = limit.copy(child = p.copy(projectList, child)))
+
+    case p @ Project(projectList, offset @ Offset(_, child))
+      if projectList.forall(_.deterministic) =>
+      offset.copy(child = p.copy(projectList, child))
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushProjectionThroughOffsetSuite.scala → sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushProjectionThroughLimitAndOffsetSuite.scala

@@ -25,17 +25,71 @@ import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
 
-class PushProjectionThroughOffsetSuite extends PlanTest {
-
+class PushProjectionThroughLimitAndOffsetSuite extends PlanTest {
   object Optimize extends RuleExecutor[LogicalPlan] {
     val batches = Batch("Optimizer Batch",
       FixedPoint(100),
-      PushProjectionThroughLimit,
-      PushProjectionThroughOffset,
+      PushProjectionThroughLimitAndOffset,
       EliminateLimits,
       LimitPushDown) :: Nil
   }
 
+  test("SPARK-40501: push projection through limit") {
+    val testRelation = LocalRelation.fromExternalRows(
+      Seq("a".attr.int, "b".attr.int, "c".attr.int),
+      1.to(20).map(_ => Row(1, 2, 3)))
+
+    val query1 = testRelation
+      .limit(10)
+      .select(Symbol("a"), Symbol("b"), 'c')
+      .limit(15).analyze
+    val optimized1 = Optimize.execute(query1)
+    val expected1 = testRelation
+      .select(Symbol("a"), Symbol("b"), 'c')
+      .limit(10).analyze
+    comparePlans(optimized1, expected1)
+
+    val query2 = testRelation
+      .sortBy($"a".asc)
+      .limit(10)
+      .select(Symbol("a"), Symbol("b"), 'c')
+      .limit(15).analyze
+    val optimized2 = Optimize.execute(query2)
+    val expected2 = testRelation
+      .sortBy($"a".asc)
+      .select(Symbol("a"), Symbol("b"), 'c')
+      .limit(10).analyze
+    comparePlans(optimized2, expected2)
+
+    val query3 = testRelation
+      .limit(10)
+      .select(Symbol("a"), Symbol("b"), 'c')
+      .limit(20)
+      .select(Symbol("a"))
+      .limit(15).analyze
+    val optimized3 = Optimize.execute(query3)
+    val expected3 = testRelation
+      .select(Symbol("a"), Symbol("b"), 'c')
+      .select(Symbol("a"))
+      .limit(10).analyze
+    comparePlans(optimized3, expected3)
+
+    val query4 = testRelation
+      .sortBy($"a".asc)
+      .limit(10)
+      .select(Symbol("a"), Symbol("b"), 'c')
+      .limit(20)
+      .select(Symbol("a"))
+      .limit(15).analyze
+    val optimized4 = Optimize.execute(query4)
+    val expected4 = testRelation
+      .sortBy($"a".asc)
+      .select(Symbol("a"), Symbol("b"), 'c')
+      .select(Symbol("a"))
+      .limit(10).analyze
+    comparePlans(optimized4, expected4)
+  }
+
   test("push projection through offset") {
     val testRelation = LocalRelation.fromExternalRows(
       Seq("a".attr.int, "b".attr.int, "c".attr.int),

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala

@@ -87,7 +87,7 @@ class SparkOptimizer(
       ColumnPruning,
       LimitPushDown,
       PushPredicateThroughNonJoin,
-      PushProjectionThroughLimit,
+      PushProjectionThroughLimitAndOffset,
       RemoveNoopOperators),
     Batch("Infer window group limit", Once,
       InferWindowGroupLimit,

sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFsSuite.scala

@@ -198,8 +198,8 @@ class ExtractPythonUDFsSuite extends SparkPlanTest with SharedSparkSession {
   test("Infers LocalLimit for Python evaluator") {
     val df = Seq(("Hello", 4), ("World", 8)).toDF("a", "b")
 
-    // Check that PushProjectionThroughLimit brings GlobalLimit - LocalLimit to the top (for
-    // CollectLimit) and that LimitPushDown keeps LocalLimit under UDF.
+    // Check that PushProjectionThroughLimitAndOffset brings GlobalLimit - LocalLimit to the top
+    // (for CollectLimit) and that LimitPushDown keeps LocalLimit under UDF.
     val df2 = df.limit(1).select(batchedPythonUDF(col("b")))
     assert(df2.queryExecution.optimizedPlan match {
       case Limit(_, _) => true