diff --git a/core/wren/src/wren/connector/spark.py b/core/wren/src/wren/connector/spark.py index 6612d807ec..9ef294c8f7 100644 --- a/core/wren/src/wren/connector/spark.py +++ b/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() diff --git a/core/wren/tests/unit/test_spark_semicolon.py b/core/wren/tests/unit/test_spark_semicolon.py index 4d1f5e9af0..bb2149fee5 100644 --- a/core/wren/tests/unit/test_spark_semicolon.py +++ b/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: