[SPARK-45369][SQL] Push down limit through generate

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

@@ -690,7 +690,7 @@ object RemoveNoopUnion extends Rule[LogicalPlan] {
 }
 
 /**
- * Pushes down [[LocalLimit]] beneath UNION ALL, OFFSET and joins.
+ * Pushes down [[LocalLimit]] beneath UNION ALL, OFFSET, joins, Generate and Python UDFs.
  */
 object LimitPushDown extends Rule[LogicalPlan] {
 
@@ -773,6 +773,10 @@ object LimitPushDown extends Rule[LogicalPlan] {
     // Push down local limit 1 if join type is LeftSemiOrAnti and join condition is empty.
     case j @ Join(_, right, LeftSemiOrAnti(_), None, _) if !right.maxRows.exists(_ <= 1) =>
       j.copy(right = maybePushLocalLimit(Literal(1, IntegerType), right))
+    case LocalLimit(exp, g: Generate) if g.outer =>
+      LocalLimit(exp, g.copy(child = maybePushLocalLimit(exp, g.child)))
+    case LocalLimit(exp, p @ Project(_, g: Generate)) if g.outer =>
+      LocalLimit(exp, p.copy(child = g.copy(child = maybePushLocalLimit(exp, g.child))))
     // Push down limits through Python UDFs.
     case LocalLimit(le, udf: BatchEvalPython) =>
       LocalLimit(le, udf.copy(child = maybePushLocalLimit(le, udf.child)))

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

@@ -21,10 +21,11 @@ import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.analysis.EliminateSubqueryAliases
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.expressions.Add
+import org.apache.spark.sql.catalyst.expressions.{Add, Explode}
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules._
+import org.apache.spark.sql.types.IntegerType
 
 class LimitPushdownSuite extends PlanTest {
 
@@ -45,6 +46,9 @@ class LimitPushdownSuite extends PlanTest {
   private val testRelation2 = LocalRelation.fromExternalRows(
     Seq("d".attr.int, "e".attr.int, "f".attr.int),
     1.to(6).map(_ => Row(1, 2, 3)))
+  private val arrayRelation = LocalRelation.fromExternalRows(
+    Seq("a".attr.int, "b".attr.array(IntegerType)),
+    1.to(6).map(i => Row(i, Array(1, i))))
   private val x = testRelation.subquery("x")
   private val y = testRelation.subquery("y")
 
@@ -352,4 +356,31 @@ class LimitPushdownSuite extends PlanTest {
       comparePlans(Optimize.execute(originalQuery2), originalQuery2)
     }
   }
+
+  test("Push down limit through generate") {
+    comparePlans(
+      Optimize.execute(arrayRelation.generate(Explode($"b"), outer = true).limit(2).analyze),
+      LocalLimit(2, arrayRelation).generate(Explode($"b"), outer = true).limit(2).analyze)
+
+    // Do not push down if outer = false
+    val originQuery = arrayRelation.generate(Explode($"b"), outer = false).limit(2).analyze
+    comparePlans(
+      Optimize.execute(originQuery),
+      originQuery)
+  }
+
+  test("Push down limit through generate and join") {
+    val condition = Some("x.a".attr === "y.a".attr)
+    val originQuery = arrayRelation.as("x").join(y, LeftOuter, condition)
+      .generate(Explode($"x.b"), outer = true)
+      .limit(2)
+    val correctAnswer =
+      LocalLimit(2, LocalLimit(2, arrayRelation).as("x").join(y, LeftOuter, condition))
+      .generate(Explode($"x.b"), outer = true)
+      .limit(2)
+
+    comparePlans(
+      Optimize.execute(originQuery.analyze),
+      correctAnswer.analyze)
+  }
 }