cloud-fan commented on code in PR #59529:
URL: https://github.com/apache/spark/pull/59529#discussion_r4346075552
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushProjectionThroughLimitAndOffsetSuite.scala:
##########
@@ -159,4 +159,103 @@ class PushProjectionThroughLimitAndOffsetSuite extends
PlanTest {
.analyze
comparePlans(optimized6, expected6)
}
+
+ test("SPARK-59438: do not push projection with scalar subquery through
limit") {
+ // A scalar subquery is deterministic, but later rewrites transform it
into a join
+ // that can destroy the sort order Limit depends on.
+ val testRelation = LocalRelation.fromExternalRows(
+ Seq("a".attr.int, "b".attr.int, "c".attr.int),
+ 1.to(20).map(_ => Row(1, 2, 3)))
+ val subqueryRelation = LocalRelation.fromExternalRows(
+ Seq("x".attr.int, "y".attr.int),
+ Seq(Row(1, 10)))
+ val scalarSubquery = ScalarSubquery(
+ subqueryRelation.groupBy($"x")(sum($"y").as("total")).analyze)
+
+ // Project with scalar subquery over GlobalLimit + LocalLimit: should NOT
be pushed down
+ val query1 = testRelation
+ .sortBy($"a".asc)
+ .limit(3)
+ .select($"a", scalarSubquery.as("sub"))
+ .analyze
+ val optimized1 = Optimize.execute(query1)
+ // The Project should remain above the Limit, not be pushed below it
+ comparePlans(optimized1, query1)
+
+ // Project with scalar subquery over LocalLimit only: should NOT be pushed
down
+ val query2 = testRelation
+ .limit(10)
+ .select($"a", scalarSubquery.as("sub"))
+ .limit(5).analyze
+ val optimized2 = Optimize.execute(query2)
+ // After EliminateLimits merges the limits to 5, the subquery Project
should stay on top
+ val expected2 = testRelation
Review Comment:
**Nit (P3):** This expectation can't match the optimizer output, with or
without the fix. `query2` is `GlobalLimit(5, LocalLimit(5, Project(a, sub,
GlobalLimit(10, LocalLimit(10, rel)))))`. With the guard, the subquery Project
stays between the two limits. `EliminateLimits` only merges directly nested
limits, or drops a limit whose child's `maxRows` fits (`Project.maxRows` is 10
here, which is more than 5). `LimitPushDown` has no plain-Project case. So
`optimized2 == query2`. Without the guard, the `Limit(le, Limit(ne, _))` case
merges the limits into `GlobalLimit(5, LocalLimit(5, Project(a, sub, rel)))`,
which doesn't match either.
The comments are off too. `testRelation.limit(10)` builds
`GlobalLimit(LocalLimit(...))`, not a bare LocalLimit, so this scenario repeats
`query1`'s arm, and the limits are never merged. The intended expectation is
`query2` itself. Covering the `Project`-over-`LocalLimit` arm needs
`localLimit(...)`, as the existing offset tests use.
See **Shared repair plan 1** in the review body.
##########
sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala:
##########
@@ -3227,4 +3227,142 @@ class SubquerySuite extends SharedSparkSession
}
}
}
+
+ test("SPARK-59438: scalar subquery in SELECT over sorted LIMIT returns
correct rows") {
+ // Verifies that PushProjectionThroughLimitAndOffset does NOT push a
Project
+ // containing a scalar subquery below a Limit. If pushed, the
subquery-to-join
+ // rewrite inserts a join below the Limit, destroying the sort order and
+ // returning wrong rows.
+ withTempView("p", "k") {
+ // Table p: 10 rows with id 1..10, name, age
+ sql(
+ """CREATE TEMPORARY VIEW p AS
+ |SELECT * FROM VALUES
+ | (1, 'Alice', 30), (2, 'Bob', 25), (3, 'Carol', 35),
+ | (4, 'Dave', 28), (5, 'Eve', 22), (6, 'Frank', 40),
+ | (7, 'Grace', 33), (8, 'Heidi', 27), (9, 'Ivan', 31),
+ | (10, 'Judy', 29)
+ |AS t(id, name, age)
+ """.stripMargin)
+
+ // Table k: counts per id
+ sql(
+ """CREATE TEMPORARY VIEW k AS
+ |SELECT * FROM VALUES (1, 10), (2, 20), (3, 30), (5, 50), (7, 70)
+ |AS t(a, b)
+ """.stripMargin)
+
+ // CTE: top 3 by age DESC => Frank(40), Carol(35), Grace(33)
+ // Then for each, compute a scalar subquery count from k.
+ // The bug would return wrong 3 rows because the join destroys sort
order.
+ val result = sql(
+ """WITH s AS (SELECT id, name FROM p ORDER BY age DESC LIMIT 3)
+ |SELECT name, (SELECT count(*) FROM k WHERE k.a = s.id) AS cnt
+ |FROM s ORDER BY name
+ """.stripMargin)
+
+ // Frank(id=6, no k match -> 0), Carol(id=3, one k match -> 1),
Grace(id=7, one k match -> 1)
+ checkAnswer(result, Seq(
+ Row("Carol", 1),
+ Row("Frank", 0),
+ Row("Grace", 1)
+ ))
+ }
+ }
+
+ test("SPARK-59438: scalar subquery in SELECT over OFFSET returns correct
rows") {
+ withTempView("p2", "k2") {
+ sql(
+ """CREATE TEMPORARY VIEW p2 AS
+ |SELECT * FROM VALUES
+ | (1, 'A', 10), (2, 'B', 20), (3, 'C', 30), (4, 'D', 40), (5, 'E',
50)
+ |AS t(id, name, score)
+ """.stripMargin)
+
+ sql(
+ """CREATE TEMPORARY VIEW k2 AS
+ |SELECT * FROM VALUES (3, 100), (4, 200), (5, 300)
+ |AS t(a, b)
+ """.stripMargin)
+
+ // ORDER BY score ASC LIMIT 5 OFFSET 2 => skip first 2 (A=10, B=20),
take C(30), D(40), E(50)
+ val result = sql(
+ """WITH s AS (SELECT id, name FROM p2 ORDER BY score ASC LIMIT 5
OFFSET 2)
+ |SELECT name, (SELECT sum(b) FROM k2 WHERE k2.a = s.id) AS total
+ |FROM s ORDER BY name
+ """.stripMargin)
+
+ checkAnswer(result, Seq(
+ Row("C", 100),
+ Row("D", 200),
+ Row("E", 300)
+ ))
+ }
+ }
+
+ test("SPARK-59438: EXISTS in SELECT over sorted LIMIT returns correct rows")
{
+ // EXISTS is also a subquery expression that gets rewritten to a join.
+ withTempView("p3", "k3") {
+ sql(
+ """CREATE TEMPORARY VIEW p3 AS
+ |SELECT * FROM VALUES
+ | (1, 'A', 10), (2, 'B', 20), (3, 'C', 30), (4, 'D', 40), (5, 'E',
50)
+ |AS t(id, name, score)
+ """.stripMargin)
+
+ sql(
+ """CREATE TEMPORARY VIEW k3 AS
+ |SELECT * FROM VALUES (1), (3), (5) AS t(a)
+ """.stripMargin)
+
+ // Top 3 by score DESC => E(50), D(40), C(30). Check EXISTS against k3.
+ val result = sql(
+ """WITH s AS (SELECT id, name FROM p3 ORDER BY score DESC LIMIT 3)
+ |SELECT name, EXISTS(SELECT 1 FROM k3 WHERE k3.a = s.id) AS has_match
+ |FROM s ORDER BY name
+ """.stripMargin)
+
+ // C(id=3, match), D(id=4, no match), E(id=5, match)
+ checkAnswer(result, Seq(
+ Row("C", true),
+ Row("D", false),
+ Row("E", true)
+ ))
+ }
+ }
+
+ test("SPARK-59438: forced SortMergeJoin reproduces the bug without the fix")
{
+ // When broadcast is disabled, the subquery-to-join rewrite uses
SortMergeJoin
+ // which repartitions inputs and destroys sort order. This test ensures
the fix
+ // works even with SMJ.
+ withTempView("p4", "k4") {
+ sql(
+ """CREATE TEMPORARY VIEW p4 AS
+ |SELECT * FROM VALUES
+ | (1, 'A', 100), (2, 'B', 90), (3, 'C', 80), (4, 'D', 70), (5, 'E',
60)
Review Comment:
**Nit (P3):** Of the new end-to-end tests, only the first one (`ORDER BY age
DESC LIMIT 3` over `p`) fails if `canPushThrough` is reverted to
`projectList.forall(_.deterministic)`. I traced the others; I didn't run them:
- Forced SMJ (this fixture): `p4`'s ids, which are the join key, ascend in
the same order as `score DESC`. Without the fix, the Sort under the pushed
`LocalLimit(3)` is eliminated, and the sort-merge join re-sorts the surviving
rows by id. With AQE coalescing the tiny shuffle, the outer limit still keeps
A, B, C, the expected rows. So despite its name, this test does not reproduce
the bug. (`SharedSparkSession` excludes `ConvertToLocalRelation`, so the reason
is the join-key order, not LocalRelation folding.)
- OFFSET: `p2` is inserted in `score ASC` order, and the broadcast join
streams it in that order, so skipping two rows still yields C, D, E.
- EXISTS: `RewritePredicateSubquery` runs after the Eliminate Sorts batch,
so the logical Sort survives, and the default broadcast existence join
preserves the streamed sorted order. Order loss only shows up with a
non-order-preserving join such as SMJ.
- Catalyst: no case puts the subquery Project directly on `localLimit(...)`,
so the bare `LocalLimit` arm is never exercised. Every negative case also uses
only an uncorrelated scalar subquery.
As a result, narrowing the guard to `ScalarSubquery`, or dropping it from
the bare `LocalLimit` arm, would make no additional new test fail. Suggest
using fixtures whose insertion order, and for the SMJ case also id order,
differ from the requested sort order. Run the EXISTS case (plus an IN case)
with `spark.sql.autoBroadcastJoinThreshold=-1`, and add a `localLimit(...)`
plan case.
See **Shared repair plan 1** in the review body.
--
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]
