vboo123 opened a new pull request, #57347:
URL: https://github.com/apache/spark/pull/57347
### What changes were proposed in this pull request?
This PR fixes `LogicalQueryStage.maxRows` to only use runtime row counts
when the underlying query stage is materialized. Before materialization,
`computeStats()` returns `None` and the code falls back to `logicalPlan.stats`,
which is a cost estimate, not a hard upper bound. That estimate was being
surfaced as `maxRows`, causing `EliminateLimits` to incorrectly remove a global
LIMIT.
The fix adds a materialization check: if the stage is materialized, use the
runtime `rowCount` (an exact upper bound). Otherwise, fall back to
`logicalPlan.maxRows`, which is always a sound upper bound derived from plan
structure rather than cardinality estimates.
### Why are the changes needed?
With AQE enabled, a root global LIMIT can be dropped during adaptive
re-optimization, returning too many rows. For example:
```sql
WITH top AS (SELECT id % 1 AS k FROM range(0, 2, 1, 1) OFFSET 1),
b AS (SELECT * FROM VALUES (0), (0) AS b(k))
SELECT 1 AS c FROM top LEFT JOIN b USING (k) LIMIT 1
```
This query returns 2 rows instead of 1.
The root cause: `EliminateLimits` removes a `GlobalLimit` when
`child.maxRows <= limit`. The child here is a `Project` over a `LeftOuter Join`
whose left input is a `LogicalQueryStage` wrapping `Offset 1 +- Range`. Before
that stage materializes, `computeStats()` falls back to `logicalPlan.stats` and
yields `rowCount = 0`. `Join.maxRows` for `LeftOuter` is `max(left * right,
left)`, which correctly propagates the 0 upward. `EliminateLimits` then sees
`maxRows = 0 <= 1` and drops the LIMIT.
The failure is timing-dependent because it depends on whether the stage has
materialized when re-optimization runs. This was introduced by SPARK-36424 (AQE
EliminateLimits support).
### Does this PR introduce _any_ user-facing change?
No API or behavior change. This fixes a correctness bug where queries with a
LIMIT over certain AQE plans could return too many rows.
### How was this patch tested?
Added a regression test in `AdaptiveQueryExecSuite`: "SPARK-57956: AQE must
not eliminate a global LIMIT over a row-increasing outer join". Because the
trigger is timing-dependent, the test runs the query 20 times and asserts each
result has exactly 1 row.
Verified:
- Without the fix, the test fails (returns 2 rows).
- With the fix, the test passes.
- The existing test "SPARK-36424: Support eliminate limits in AQE Optimizer"
still passes, confirming that legitimate limit elimination over materialized
aggregate stages is preserved.
Ran `AdaptiveQueryExecSuite` locally. Full test matrix relies on GitHub
Actions.
### Was this patch authored or co-authored using generative AI tooling?
Generated using Kiro (Claude Opus 4.8)
--
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]
