[SPARK-58093][SQL] Push down LIMIT through UNION ALL to Data Source V2 scan

[j1wonpark]

What changes were proposed in this pull request?

V2ScanRelationPushDown.pushDownLimit had no case for Union, so a LIMIT on top of a UNION ALL was never pushed to the underlying Data Source V2 scans through SupportsPushDownLimit; it fell through to case other => (other, false).

Note the plan shape: LimitPushDown already pushes a LocalLimit into each Union branch (operator level), but the source-level limit hint (pushDownLimit) stops at the Union.

This PR adds two cases to pushDownLimit:

• case u: Union — recurse into each branch, pushing the limit as a partial limit so the limit operators above the union are kept.
• case l @ LocalLimit(IntegerLiteral(_), _) — look through the per-branch LocalLimit that LimitPushDown inserts, so the limit reaches the scan.

Results are unchanged. UNION ALL does not de-duplicate rows, so the union of branches that each keep at most limit rows still contains the first limit rows of the whole union — the same reason the existing Union case in LimitPushDown is safe. The union may return more than limit rows, so the push is only partial and the limit operators stay in the plan.

Two things worth calling out for review:

• A branch may fully push a top N and drop its Sort (the Sort case already returns a plan without it). This is fine because Union does not attach meaning to the order in which it emits its children, and the limit above the union still caps the result. LIMIT over an ORDER BY that sits above the union is not pushed at all, which a test covers.
• Because the Union case always reports the limit as not removable, plans that combine LIMIT with OFFSET can call pushLimit on the same builder twice with the same value. This already happens today for partially pushed scans and is idempotent for the sources in the repository.

Why are the changes needed?

SELECT ... FROM a UNION ALL SELECT ... FROM b LIMIT n never delivers the LIMIT to the scans, so each source does more work than necessary. Any Data Source V2 connector implementing SupportsPushDownLimit benefits from receiving the limit per branch — e.g. for JDBC, each branch's pushed-down SQL then carries LIMIT n.

Does this PR introduce any user-facing change?

No. Query results are identical. The only observable difference is that the pushed limit now appears on each UNION ALL branch's scan (e.g. PushedLimit: LIMIT n in the plan, or a LIMIT clause in the JDBC SQL), reducing work at the source.

How was this patch tested?

Added five tests to JDBCV2Suite, each asserting the limit pushed to every branch scan:

• a two-branch union, filtered to a single matching employee so the result is deterministic and can be checked with checkAnswer;
• a three-branch union;
• a union with a sorted branch, where the branch limit is pushed as a top N and the branch Sort is dropped;
• a union combined with OFFSET, covering both LIMIT n OFFSET m (branches keep n + m rows) and limit(n).offset(m) (branches keep n rows);
• a negative test: a LIMIT over an ORDER BY on top of the union is not pushed to the scans.

Ran the full JDBCV2Suite (86 tests) — no regression.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus)

[yadavay-amzn]

Looks correct, a couple of non-blocking notes

I traced the per-branch limit values, both OFFSET plan shapes, the ORDER-BY-above-union case, and UNION DISTINCT against master, and I think the correctness story holds. The key guard, where the Union case always reports the push as not removable so the outer limit stays above the union, is right.

Left two non-blocking comments inline, neither is a correctness issue. Thanks for the clear write-up in the description, the plan-shape explanation made this much easier to follow.

What I verified:

• Each branch scan receives the full union-level limit, and branches with a smaller or larger explicit LIMIT still don't under-produce, since the intermediate limits are preserved and GlobalLimit-wrapped branches aren't recursed into.
• LIMIT 4 OFFSET 2 lands on the OffsetAndLimit pattern and pushes n+m=6, while .limit(4).offset(2) lands on LimitAndOffset and pushes 4. Both match the tests.
• ORDER BY above the union isn't pushed, since the Sort's child is the Union rather than a scan.
• UNION DISTINCT is Distinct(Union(...)), which stops the recursion, so nothing is pushed.

[yadavay-amzn]

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.

[j1wonpark]

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.

[yadavay-amzn]

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.

[j1wonpark]

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.