fix(spark): push LIMIT into SQL instead of client slice

[Bartok9]

Summary

• Push LIMIT n into Spark SQL via subquery wrap (after semicolon strip)
• dry_run uses SELECT * FROM (...) LIMIT 0 instead of client .limit(0)

License

Apache-2.0 (core/**).

Motivation

Spark connector applied limit only after toPandas(), so engine still produced full results for bounded asks. Matches DuckDB/Postgres pushdown style.

Verification

cd core/wren && .venv/bin/python -m pytest tests/unit/test_spark_semicolon.py -v — 5 passed

Duplicate check

• No open PR for Spark LIMIT pushdown

Summary by CodeRabbit

• Bug Fixes

  • Improved Spark query handling when a result limit is specified by pushing LIMIT into Spark SQL (instead of client-side slicing), reducing unnecessary data processing.
  • Improved dry-run behavior by wrapping the cleaned query with LIMIT 0 before counting, and continuing to handle trailing semicolons safely.
  • Added validation to ensure provided limits are non-negative integers before executing.

• Tests

  • Added/updated unit tests to verify SQL limit pushdown and trailing-semicolon stripping, including dry-run LIMIT 0 wrapping behavior.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 251c3f1a-e200-4c6e-847f-b139e4744796

📥 Commits

Reviewing files that changed from the base of the PR and between 6875986 and e1474b4.

📒 Files selected for processing (1)

• core/wren/src/wren/connector/spark.py

🚧 Files skipped from review as they are similar to previous changes (1)

• core/wren/src/wren/connector/spark.py

---

Walkthrough

SparkConnector now strips trailing semicolons and pushes query limits into Spark SQL subquery wrappers. Dry runs use cleaned SQL with LIMIT 0, and unit tests validate both generated SQL paths.

Changes

Spark SQL limit handling

Layer / File(s)	Summary
Query LIMIT pushdown core/wren/src/wren/connector/spark.py, core/wren/tests/unit/test_spark_semicolon.py	query cleans SQL, validates the optional limit, wraps limited queries, and converts results directly to PyArrow; tests cover semicolon stripping and LIMIT generation.
Dry-run LIMIT zero core/wren/src/wren/connector/spark.py, core/wren/tests/unit/test_spark_semicolon.py	dry_run executes cleaned SQL through a LIMIT 0 subquery and calls count(), replacing the prior DataFrame .limit(0) behavior.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• Canner/WrenAI#2430: Applies analogous semicolon stripping and subquery-based limit handling.
• Canner/WrenAI#2457: Updates connector query logic with matching SQL wrapping and limit pushdown.
• Canner/WrenAI#2464: Modifies Spark SQL semicolon handling and related tests.

Suggested reviewers: goldmedal

Poem

I’m a bunny hopping through SQL,
Stripping semicolons with a twirl.
Limits now land before rows grow,
Dry runs count what they need to know.
Spark paths neatly wrap and flow!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: pushing Spark LIMIT handling into SQL instead of client-side slicing.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@core/wren/src/wren/connector/spark.py`:
- Around line 29-30: Validate limit before the SQL construction in the
limit-handling branch, rejecting negative values with a clear connector-level
validation error. Preserve the existing int conversion and LIMIT generation for
non-negative limits, including zero.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 0efd5983-ebbc-421f-83dd-a5fd8d924243

📥 Commits

Reviewing files that changed from the base of the PR and between d472877 and 6875986.

📒 Files selected for processing (2)

• core/wren/src/wren/connector/spark.py
• core/wren/tests/unit/test_spark_semicolon.py

[Bartok9]

Addressed in e1474b4 — negative limits now raise a clear connector-level ValueError before SQL construction, matching the MySQL connector's _coerce_limit pattern (zero still allowed).