fix(spark): apply limit via DataFrame.limit before toPandas - #2574
fix(spark): apply limit via DataFrame.limit before toPandas#2574Bartok9 wants to merge 8 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
Walkthrough
ChangesSpark SQL limit handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant SparkConnector
participant SparkSQLSession
participant SparkDataFrame
Client->>SparkConnector: Submit SQL and optional limit
SparkConnector->>SparkConnector: Strip semicolon and validate limit
SparkConnector->>SparkSQLSession: Execute wrapped or unchanged SQL
SparkSQLSession-->>SparkConnector: Return DataFrame
SparkConnector->>SparkDataFrame: Apply DataFrame limit when required
SparkConnector-->>Client: Return PyArrow table
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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
📒 Files selected for processing (2)
core/wren/src/wren/connector/spark.pycore/wren/tests/unit/test_spark_semicolon.py
|
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). |
Code reviewThe 1. WrenAI/core/wren/src/wren/connector/spark.py Lines 44 to 49 in e1474b4
If the goal is only to keep a trailing 2. The comment describes a change that isn't happening. It says the wrapper is preferred "so EXPLAIN is unnecessary", but 3. Please confirm the wrap is safe on Spark for duplicate output columns. This codebase has already rejected subquery-wrapping once for exactly that reason: WrenAI/core/wren/src/wren/connector/mysql.py Lines 36 to 41 in 5d2637c Spark is more permissive than MySQL here, so it may well be fine — but given the documented history, a note in the description confirming a join that projects two same-named columns still works under the wrap would be reassuring. Rebase before merge — 5 commits behind |
|
To clean my review queue, I changed the PR status to draft. After addressing the review comment, you can request me again. |
e1474b4 to
b9ee3ee
Compare
|
Thanks @goldmedal — all three points addressed in the latest push.
Test updated to assert the DataFrame validation path. Rebased onto current |
goldmedal
left a comment
goldmedal
left a comment
There was a problem hiding this comment.
Summary
The direction is right and matches the established pushdown pattern in this codebase (duckdb / snowflake / databricks all subquery-wrap). Two things I'd like addressed before merge, both stemming from the fact that the wrap path is far hotter than it looks: mcp_server.py::_query_with_limit_probe never passes limit=None (it falls back to DEFAULT_ROW_LIMIT), so every Spark query issued through MCP run_sql now goes through the new subquery wrap.
1. Trailing line comment breaks the wrap (spark.py)
The single-line f-string means user SQL ending in a line comment produces invalid SQL:
-- input: SELECT 1 -- note
SELECT * FROM (SELECT 1 -- note) AS _q LIMIT 10
-- ^ swallows ") AS _q LIMIT 10"snowflake.py hit exactly this and fixed it by putting the user SQL on its own line, with the rationale in a comment:
# Place the user SQL on its own line so a trailing line comment
# (`-- ...`) cannot swallow the closing paren, alias, or LIMIT.
executed = (
"SELECT * FROM (\n"
f"{strip_trailing_semicolon(sql)}\n"
f") AS _wren_sub LIMIT {int(limit)}"
)Worth mirroring here. On main this SQL works (the limit was applied client-side), so as written this is a behavioural regression for LLM/agent-generated SQL, which routinely carries comments.
2. query() now breaks the statements dry_run deliberately protects
The new dry_run comment states the rule explicitly:
Validate via the DataFrame API so statements that are not legal as a subquery (SHOW TABLES, DESCRIBE, ...) are still accepted, exactly as before.
But query() wraps unconditionally whenever limit is not None, and via MCP that is always. So SHOW TABLES passes dry_run and fails query — and it succeeded on main. Either the same carve-out should apply to query() (skip the wrap / fall back to a client-side slice for non-subqueryable statements), or the trade-off should be stated as a deliberate accepted regression rather than contradicted one method apart.
3. _coerce_limit already exists (minor)
mysql.py:44 has this exact validate-and-coerce, down to the f"limit must be non-negative, got {coerced}" message. This is now a second copy. Promoting it to connector/base.py and calling it from spark/mysql would also give duckdb and snowflake the negative-limit check they currently lack (both interpolate a bare int(limit)).
4. PR description is stale (minor)
The body still says "dry_run uses SELECT * FROM (...) LIMIT 0 instead of client .limit(0)", but commit b9ee3ee reverted that — the dry_run diff is now a no-op (comment + local variable, identical behaviour). Worth correcting, since the body becomes the squash-merge message.
5. Test coverage (minor)
The tests are mock-only string assertions, which is reasonable for shape, but two behaviours added in this PR are untested:
- negative
limitraisingValueError(added in5e45530) limit=0producingLIMIT 0rather than being treated as falsy
A trailing-line-comment case would also lock in the fix for (1). Note there is no Spark entry in tests/connectors/, so string construction is the only evidence we have that the generated SQL is valid — nothing exercises a real Spark session.
Verdict
(1) and (2) are blocking; (3)–(5) are non-blocking.
|
Thanks @goldmedal — addressed the two blocking items (plus the easy test/description nits) in the latest push.
|
Address goldmedal review on Canner#2574: - Mirror snowflake multiline wrap so trailing `--` comments cannot swallow the closing paren/alias/LIMIT - When limit is set for SHOW/DESCRIBE/etc., keep DataFrame client slice (MCP always passes DEFAULT_ROW_LIMIT) - Cover comment wrap, limit=0, negative limit, and SHOW TABLES path
There was a problem hiding this comment.
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 11-15: Update the classification around _NON_SUBQUERYABLE to skip
leading SQL comments, including line comments followed by newlines, before
testing the command keyword. Ensure standalone commands such as SHOW TABLES
remain classified as non-subqueryable when prefixed by comments, while
preserving the existing keyword matching behavior.
🪄 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: 246f278c-979d-42a1-81ef-f9c21aee1fd4
📒 Files selected for processing (2)
core/wren/src/wren/connector/spark.pycore/wren/tests/unit/test_spark_semicolon.py
|
Addressed CodeRabbit’s leading-comment classification note. `_is_non_subqueryable` now strips leading whitespace and `--` / `/* */` comments before the keyword match, so e.g. `-- metadata\nSHOW TABLES` with a limit stays on the DataFrame path instead of being subquery-wrapped. Added a unit test; `pytest tests/unit/test_spark_semicolon.py -v` — 10 passed. |
goldmedal
left a comment
goldmedal
left a comment
There was a problem hiding this comment.
Code review
Re-checked all three points from my last round. They're genuinely addressed, and your duplicate-column claim holds up — I verified it. But in fixing them the patch grew a 35-line SQL-string classifier, and I'm asking for it to come out before merge.
🔴 Blocking — the subquery wrap is unnecessary; DataFrame.limit(n) already is the pushdown
I measured both paths on a local PySpark 4.1.1 session:
spark.sql("SELECT * FROM t").limit(3).explain()
== Physical Plan ==
CollectLimit 3
+- *(1) Range (0, 1000, step=1, splits=1)
spark.sql("SELECT * FROM (\nSELECT * FROM t\n) AS _q LIMIT 3").explain()
== Physical Plan ==
CollectLimit 3
+- *(1) Range (0, 1000, step=1, splits=1)
Identical. DataFrame.limit() is a Limit node in the logical plan — it is server-side, not a client slice.
The bug on main was never "no SQL LIMIT"; it was that the limit was applied to the Arrow table after toPandas():
arrow_table = pa.Table.from_pandas(df) # full result already materialized
if limit is not None:
arrow_table = arrow_table.slice(0, limit)So the else branch this PR already added is the entire fix:
frame = self.connection.sql(strip_trailing_semicolon(sql))
if coerced is not None:
frame = frame.limit(coerced)
df = frame.toPandas()That works for every statement class — SELECT, SHOW TABLES, DESCRIBE alike (verified: SHOW TABLES + .limit(5) → OK; the same statement wrapped → AnalysisException [TABLE_OR_VIEW_NOT_FOUND]: table or view 'SHOW').
Please drop the wrap branch. It deletes _NON_SUBQUERYABLE, _LEADING_SQL_NOISE, _strip_leading_sql_comments and _is_non_subqueryable outright, and the trailing-line-comment fix and the non-subqueryable carve-out both stop existing rather than being patched — which is the difference between a fix and a fix plus a permanent maintenance surface.
The reason Spark differs from snowflake/duckdb/postgres is that those hold a raw DBAPI cursor with no plan-level limit, so string surgery is their only option. Spark has a first-class DataFrame API — and databricks.py already skips the wrap in query() for the same reason (cursor.fetchmany_arrow(limit)). "Matching the other connectors" is the right instinct applied to the wrong axis: the shared contract is limit server-side, not limit via string interpolation.
🔴 Blocking (alternative) — if the wrap stays, the classifier has to change shape
I'd rather not merge the denylist in its current form either way, so if you disagree with the above, this is what needs fixing instead.
_NON_SUBQUERYABLE is wrong in both directions. Measured false positive:
_is_non_subqueryable(
"WITH t AS (SELECT CASE WHEN a=1 THEN 'INSERT' ELSE 'UPDATE' END AS op FROM db.log) SELECT * FROM t"
) # -> TrueWITH\s+.*\bINSERT\b under re.DOTALL scans the whole statement, so any INSERT token anywhere in a CTE query matches. That silently drops the pushdown this PR exists to add — and it matters more than it looks, because the connector never sees raw user SQL: engine.query() runs dry_plan() first, so what reaches SparkConnector.query is CTE-injected planned SQL that begins with WITH for essentially every model-backed query.
The other direction: any non-subqueryable statement not enumerated gets wrapped and fails, and the list can't be completed — it's a denylist over an open set.
An allowlist inverts both failure modes, and sqlglot is already a dependency used for exactly this in sql_classify.py — parse_one(sql, dialect="spark") plus an isinstance(..., exp.Select | exp.Union | ...) check is shorter and doesn't need a hand-rolled comment stripper.
Minor (non-blocking)
_strip_leading_sql_comments'swhileloop is dead code._LEADING_SQL_NOISEends in+, so a singlesub(..., count=1)already consumes all leading whitespace/--//* */runs. Verified: single-pass output == loop output on mixed inputs. (Moot if the wrap goes.)_coerce_limitis now the second copy (mysql.py:44), and #2635 is adding a third inredshift.py. Three copies is the point at whichconnector/base.pywins — and it would give duckdb/snowflake/athena/datafusion the negative-limit check they currently lack.dry_runis a behavioural no-op in this diff (comment + local variable). Correct outcome, but it's churn in a method that didn't change.- Still 3 commits behind
main.
Verified / closing my earlier points
- Duplicate output columns under the wrap: your claim is correct.
SELECT * FROM (SELECT a.id, b.id FROM t a JOIN t b ON a.id=b.id) AS _q LIMIT 2returns['id', 'id']and both columns round-trip. The MySQLER_DUP_FIELDNAMEhistory doesn't apply to Spark. Point 3 closed. - Inner
ORDER BYsurvives the outerLIMIT—[999, 998, 997]from both paths. pytest tests/unit/test_spark_semicolon.py→ 10 passed;ruff check/ruff format --checkclean. (Fulltests/unithas 4 unrelatedtest_memory.pyModuleNotFoundErrorfailures from the uninstalledmemoryextra.)
Caveat on my measurements: local mode, not Spark Connect. There's still no tests/connectors/test_spark.py, so nothing in CI exercises a real session either way — which is another argument for the path that requires no generated SQL to be valid.
Verdict
Request changes. The direction is right and the behaviour is now correct, so this is close. Blocking is the implementation shape, not the outcome: keep frame.limit(n), delete the wrap branch and the classifier with it — same plan, all statement classes, ~40 fewer lines. If you want to keep the wrap, then the denylist→allowlist change is required in its place.
Avoid full result materialization before Arrow slice; wrap dry_run in LIMIT 0 subquery after stripping trailing semicolons.
Address goldmedal review on Canner#2574: - Mirror snowflake multiline wrap so trailing `--` comments cannot swallow the closing paren/alias/LIMIT - When limit is set for SHOW/DESCRIBE/etc., keep DataFrame client slice (MCP always passes DEFAULT_ROW_LIMIT) - Cover comment wrap, limit=0, negative limit, and SHOW TABLES path
CI ruff format --check wants the subquery wrap as a single f-string.
Classify SHOW/DESCRIBE/... after stripping leading -- and /* */ comments so MCP default limits still use the DataFrame path for commented meta SQL.
goldmedal measured DataFrame.limit() as server-side CollectLimit — identical plan to subquery LIMIT wrap. Drop wrap + classifier; fix is limit before toPandas, not string surgery. Keep non-negative coerce.
8c19d73 to
3c63da0
Compare
Bartok9
changed the title
|
Thanks @goldmedal — agreed on the measurements and the shape. Dropped the subquery wrap and the entire classifier ( coerced = _coerce_limit(limit)
frame = self.connection.sql(strip_trailing_semicolon(sql))
if coerced is not None:
frame = frame.limit(coerced)
df = frame.toPandas()That matches the path you identified as the real fix (server-side |
2 participants
Summary
limitwith SparkDataFrame.limit(n)beforetoPandas(), so the engine gets a server-sideCollectLimitinstead of materializing the full result and slicing the Arrow table on the clientSHOW/DESCRIBE/SELECT alike)ValueErrorbefore executiondry_rununchanged vsmain(trailing;strip +.limit(0).count())License
Apache-2.0 (
core/**).Motivation
On
main, Spark appliedlimitonly aftertoPandas()via Arrowslice, so bounded asks still fully materialized. Maintainer review confirmedDataFrame.limit()is already plan-level pushdown (not a client slice); string-wrapping LIMIT is the right tool for DBAPI connectors, not Spark.Verification
cd core/wren && .venv/bin/python -m pytest tests/unit/test_spark_semicolon.py -v— 8 passedDuplicate check