j1wonpark commented on code in PR #57200:
URL: https://github.com/apache/spark/pull/57200#discussion_r3609540552
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -1047,6 +1047,17 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
case p: Project =>
val (newChild, isPartiallyPushed) = pushDownLimit(p.child, limit)
(p.withNewChildren(Seq(newChild)), isPartiallyPushed)
+ case u: Union =>
+ // The union of branches that each keep at most `limit` rows still
contains the first `limit`
+ // rows of the whole union, as UNION ALL does not de-duplicate. The
union may return more than
+ // `limit` rows, so this is only a partial push.
+ val newChildren = u.children.map(child => pushDownLimit(child, limit)._1)
+ (u.withNewChildren(newChildren), false)
+ case l @ LocalLimit(IntegerLiteral(_), _) =>
Review Comment:
Good catch, thanks. You're right that the arm wasn't scoped to the union and
the comment overstated the invariant. I moved the look-through inside the
`Union` case so it only applies to union branches, and it now pushes `min(outer
limit, branch limit)` so the pushed value no longer relies on the two limits
being equal. Also added a test where a branch carries its own tighter `LIMIT`.
##########
sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala:
##########
@@ -402,6 +402,93 @@ class JDBCV2Suite extends SharedSparkSession with
ExplainSuiteHelper {
}
}
+ private def checkPushedLimits(df: DataFrame, limits: Option[Int]*): Unit = {
+ val pushedLimits = df.queryExecution.optimizedPlan.collect {
+ case relation: DataSourceV2ScanRelation => relation.scan match {
+ case v1: V1ScanWrapper => v1.pushedDownOperators.limit
+ case other => fail(s"Expected a V1ScanWrapper but got $other")
+ }
+ }
+ assert(pushedLimits === limits)
+ }
+
+ test("SPARK-58093: push down LIMIT through UNION ALL to each data source
scan") {
+ // dept = 6 matches a single employee, so the union of the two branches is
deterministic
+ val df = sql(
+ """
+ |SELECT name FROM h2.test.employee WHERE dept = 6
+ |UNION ALL
+ |SELECT name FROM h2.test.employee WHERE dept = 6
+ |LIMIT 1
+ |""".stripMargin)
+ checkPushedLimits(df, Some(1), Some(1))
+ checkPushedInfo(df, "PushedFilters: [DEPT IS NOT NULL, DEPT = 6]",
"PushedLimit: LIMIT 1")
+ checkLimitRemoved(df, removed = false)
+ checkAnswer(df, Seq(Row("jen")))
+ }
+
+ test("SPARK-58093: push down LIMIT through UNION ALL with more than two
branches") {
+ val df = sql(
+ """
+ |SELECT name FROM h2.test.employee WHERE dept = 6
+ |UNION ALL
+ |SELECT name FROM h2.test.employee WHERE dept = 6
+ |UNION ALL
+ |SELECT name FROM h2.test.employee WHERE dept = 6
+ |LIMIT 2
+ |""".stripMargin)
+ checkPushedLimits(df, Some(2), Some(2), Some(2))
+ checkAnswer(df, Seq(Row("jen"), Row("jen")))
+ }
+
+ test("SPARK-58093: push down LIMIT through UNION ALL over a sorted branch") {
+ // the sort survives above the scan of the first branch, so the limit
reaches it as a top N
+ val df =
spark.read.table("h2.test.employee").select($"name").orderBy($"salary")
+ .union(spark.read.table("h2.test.employee").select($"name"))
+ .limit(1)
+ checkPushedLimits(df, Some(1), Some(1))
+ checkSortRemoved(df)
+ checkLimitRemoved(df, removed = false)
+ assert(df.collect().length == 1)
+ }
+
+ test("SPARK-58093: push down LIMIT through UNION ALL together with OFFSET") {
+ // `LIMIT n OFFSET m` skips before it takes, so each branch has to keep `n
+ m` rows
+ val df1 = sql(
+ """
+ |SELECT name FROM h2.test.employee
+ |UNION ALL
+ |SELECT name FROM h2.test.employee
+ |LIMIT 4 OFFSET 2
+ |""".stripMargin)
+ checkPushedLimits(df1, Some(6), Some(6))
+ checkOffsetRemoved(df1, removed = false)
+ assert(df1.collect().length == 4)
+
+ // `limit(n).offset(m)` takes before it skips, so `n` rows per branch are
enough
+ val df2 = spark.read.table("h2.test.employee").select($"name")
+ .union(spark.read.table("h2.test.employee").select($"name"))
+ .limit(4)
+ .offset(2)
+ checkPushedLimits(df2, Some(4), Some(4))
+ checkOffsetRemoved(df2, removed = false)
+ assert(df2.collect().length == 2)
+ }
+
+ test("SPARK-58093: LIMIT with ORDER BY on top of UNION ALL is not pushed
down") {
+ // the limit must not be pushed below the sort: a branch capped before the
global sort could
+ // drop rows that belong to the overall top N
+ val df = sql(
+ """
+ |SELECT name FROM h2.test.employee
+ |UNION ALL
+ |SELECT name FROM h2.test.employee
+ |ORDER BY name LIMIT 2
+ |""".stripMargin)
+ checkPushedLimits(df, None, None)
+ checkAnswer(df, Seq(Row("alex"), Row("alex")))
+ }
Review Comment:
Added both: a `UNION DISTINCT` test asserting nothing is pushed, and a mixed
case using a second H2 catalog with `pushDownLimit=false` so only one branch
pushes the limit. One note on the `UNION DISTINCT` test: with two identical
branches `RemoveNoopUnion` collapses the union entirely, so the test uses two
different filters to keep the `Distinct(Union(...))` shape.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
