yadavay-amzn commented on code in PR #57200:
URL: https://github.com/apache/spark/pull/57200#discussion_r3584746578
##########
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:
Non-blocking suggestion: it might be worth a `UNION DISTINCT` case here
asserting the limit isn't pushed (`checkPushedLimits(df, None, None)`). The
behavior already looks correct to me since `UNION DISTINCT` shows up as
`Distinct(Union(...))` and the recursion stops at the `Distinct`, but there's
nothing pinning that today, so a future change to the matched cases could
regress it quietly. A mixed case (one branch on a source that supports limit
pushdown, one that doesn't) would be a nice one to have too.
##########
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:
Non-blocking, and I don't think it's a correctness problem: this arm isn't
scoped to the `Union` case, so it fires for any `LocalLimit` the recursion
reaches. I tried to build a wrong-results case out of it and couldn't, since
the `LocalLimit` node is preserved by `withNewChildren` and the value that gets
pushed is always the tighter outer cap, so functionally it looks fine to me.
The one thing I'd mention is the comment says the per-branch limit has "the
same value as the limit above the union," which I don't think is guaranteed if
a branch carries its own smaller `LIMIT`. Might be worth either doing the
look-through inside the `Union` case so the intent stays local, or softening
the comment to say the pushed value is the outer cap and the branch's own
`LocalLimit` still caps the result regardless. Happy to be wrong here if you've
already convinced yourself the arm can only be reached under a union.
--
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]
