fix(spark): apply limit via DataFrame.limit before toPandas

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

@@ -4,6 +4,16 @@
 from wren.model import SparkConnectionInfo
 
 
+def _coerce_limit(limit: int | None) -> int | None:
+    """Validate and coerce a user-supplied ``limit`` to a non-negative ``int``."""
+    if limit is None:
+        return None
+    coerced = int(limit)
+    if coerced < 0:
+        raise ValueError(f"limit must be non-negative, got {coerced}")
+    return coerced
+
+
 class SparkConnector(ConnectorABC):
     def __init__(self, connection_info: SparkConnectionInfo):
         self.connection_info = connection_info
@@ -22,17 +32,22 @@ def _create_session(self):
         )
 
     def query(self, sql: str, limit: int | None = None) -> pa.Table:
-        df = self.connection.sql(strip_trailing_semicolon(sql)).toPandas()
+        # Apply limit via DataFrame.limit before toPandas so Spark pushes a
+        # CollectLimit into the plan (server-side). Avoid post-Arrow slice and
+        # SQL subquery wraps — both unnecessary on the DataFrame API and the
+        # latter breaks SHOW/DESCRIBE-style statements.
+        coerced = _coerce_limit(limit)
+        frame = self.connection.sql(strip_trailing_semicolon(sql))
+        if coerced is not None:
+            frame = frame.limit(coerced)
+        df = frame.toPandas()
         if hasattr(df, "attrs") and df.attrs:
             df.attrs = {
                 k: v
                 for k, v in df.attrs.items()
                 if k not in ("metrics", "observed_metrics")
             }
-        arrow_table = pa.Table.from_pandas(df)
-        if limit is not None:
-            arrow_table = arrow_table.slice(0, limit)
-        return arrow_table
+        return pa.Table.from_pandas(df)
 
     def dry_run(self, sql: str) -> None:
         self.connection.sql(strip_trailing_semicolon(sql)).limit(0).count()

core/wren/tests/unit/test_spark_semicolon.py

@@ -1,7 +1,10 @@
-"""Trailing-semicolon stripping for the Spark connector (mocked session)."""
+"""Trailing-semicolon stripping + DataFrame limit for the Spark connector."""
 
 from unittest.mock import MagicMock
 
+import pandas as pd
+import pytest
+
 from wren.connector.base import strip_trailing_semicolon
 from wren.connector.spark import SparkConnector
 
@@ -16,19 +19,54 @@ def _make_mock_connector() -> tuple[SparkConnector, MagicMock]:
 
 def test_query_strips_trailing_semicolon_before_sql() -> None:
     connector, session = _make_mock_connector()
-    # pandas DF mock for pa.Table.from_pandas
-    import pandas as pd
-
     session.sql.return_value.toPandas.return_value = pd.DataFrame({"x": [1, 2, 3]})
-    connector.query("SELECT 1;", limit=2)
+    connector.query("SELECT 1;")
     session.sql.assert_called_once_with("SELECT 1")
 
 
-def test_dry_run_strips_trailing_semicolon() -> None:
+def test_query_limit_uses_dataframe_limit_before_to_pandas() -> None:
+    connector, session = _make_mock_connector()
+    frame = session.sql.return_value
+    frame.limit.return_value.toPandas.return_value = pd.DataFrame({"x": [1, 2]})
+    connector.query("SELECT 1 AS x;", limit=2)
+    session.sql.assert_called_once_with("SELECT 1 AS x")
+    frame.limit.assert_called_once_with(2)
+    frame.limit.return_value.toPandas.assert_called_once_with()
+    # Must not call toPandas on the unlimited frame.
+    frame.toPandas.assert_not_called()
+
+
+def test_query_limit_zero_uses_dataframe_limit_zero() -> None:
+    connector, session = _make_mock_connector()
+    frame = session.sql.return_value
+    frame.limit.return_value.toPandas.return_value = pd.DataFrame({"x": []})
+    connector.query("SELECT 1 AS x", limit=0)
+    session.sql.assert_called_once_with("SELECT 1 AS x")
+    frame.limit.assert_called_once_with(0)
+
+
+def test_query_negative_limit_raises() -> None:
+    connector, session = _make_mock_connector()
+    with pytest.raises(ValueError, match="non-negative"):
+        connector.query("SELECT 1", limit=-1)
+    session.sql.assert_not_called()
+
+
+def test_query_show_tables_with_limit_uses_dataframe_limit() -> None:
+    connector, session = _make_mock_connector()
+    frame = session.sql.return_value
+    frame.limit.return_value.toPandas.return_value = pd.DataFrame({"tableName": ["t"]})
+    connector.query("SHOW TABLES", limit=500)
+    session.sql.assert_called_once_with("SHOW TABLES")
+    frame.limit.assert_called_once_with(500)
+
+
+def test_dry_run_validates_via_dataframe_after_strip() -> None:
     connector, session = _make_mock_connector()
     connector.dry_run("SELECT 1;  \n")
     session.sql.assert_called_once_with("SELECT 1")
     session.sql.return_value.limit.assert_called_once_with(0)
+    session.sql.return_value.limit.return_value.count.assert_called_once_with()
 
 
 def test_helper_preserves_semicolon_inside_string_literal() -> None: