From 3f8b09ec6e4a721031824248f00deb928f5c1c65 Mon Sep 17 00:00:00 2001 From: Jogesh Anand Date: Sun, 21 Jun 2026 20:29:59 -0700 Subject: [PATCH 1/4] feat: add %%sql --limit and --no-display support --- python/README.md | 13 +++++ python/python/ballista/jupyter.py | 78 +++++++++++++++++++++++++---- python/python/tests/test_jupyter.py | 55 ++++++++++++++++++++ 3 files changed, 137 insertions(+), 9 deletions(-) diff --git a/python/README.md b/python/README.md index c357cccc16..d900b72262 100644 --- a/python/README.md +++ b/python/README.md @@ -146,6 +146,19 @@ You can also store results in a variable: SELECT * FROM orders WHERE status = 'pending' ``` +The cell magic accepts two optional flags before the variable name: + +```python +# Render at most 10 rows in the cell (display only; the stored result keeps +# every row). Without --limit the display defaults to 50 rows. +%%sql --limit 10 my_result +SELECT * FROM orders + +# Run the query and store the result without displaying it. +%%sql --no-display my_result +SELECT * FROM orders +``` + ### Execution Plan Visualization Visualize query execution plans directly in notebooks: diff --git a/python/python/ballista/jupyter.py b/python/python/ballista/jupyter.py index d1d8c38ed9..dff19772ed 100644 --- a/python/python/ballista/jupyter.py +++ b/python/python/ballista/jupyter.py @@ -108,8 +108,15 @@ def decorator(func): return decorator +from datafusion import configure_formatter + from .extension import BallistaSessionContext, DistributedDataFrame +# Default number of rows rendered for a ``%%sql`` cell when ``--limit`` is not +# given. This caps only the display (via datafusion's HTML formatter); the +# underlying result keeps all of its rows. +DEFAULT_DISPLAY_LIMIT = 50 + class BallistaConnectionError(Exception): """Raised when not connected to a Ballista cluster.""" @@ -221,6 +228,45 @@ def register(self, line: str) -> Optional[str]: "Currently not supporting the inserted file format" ) + @staticmethod + def _parse_cell_magic_args(line: str): + """Parse the argument line of a ``%%sql`` cell magic. + + Recognises the ``--no-display`` and ``--limit N`` flags (space form, + e.g. ``--limit 5``, consistent with the other magics in this module). + The first non-flag token, if any, is treated as the variable name to + store the result in. + + Returns a ``(var_name, no_display, limit)`` tuple where ``limit`` is + ``None`` when ``--limit`` was not supplied. Raises ``ValueError`` for a + missing or invalid ``--limit`` value. + """ + tokens = line.strip().split() + var_name = None + no_display = False + limit = None + + i = 0 + while i < len(tokens): + token = tokens[i] + if token == "--no-display": + no_display = True + elif token == "--limit": + i += 1 + if i >= len(tokens): + raise ValueError("--limit requires a number, e.g. --limit 5") + try: + limit = int(tokens[i]) + except ValueError: + raise ValueError(f"--limit expects an integer, got '{tokens[i]}'") + if limit < 1: + raise ValueError("--limit must be a positive integer") + elif not token.startswith("--") and var_name is None: + var_name = token + i += 1 + + return var_name, no_display, limit + @line_cell_magic def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]: """ @@ -243,27 +289,41 @@ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]: LIMIT 5 `my_result` will store the result of the SQL-query + + The cell magic accepts two optional flags before the variable name: + --limit N Render at most N rows in the cell output (default + 50). This caps the display only; the stored result + keeps every row. + --no-display Run the query and store the result without + displaying it. """ if not cell: return self._execute_sql(line.strip()) if line.strip() else None else: - var_name = None query = cell.strip() if not query: return None - args = line.strip().split() - i = 0 - while i < len(args): - if not args[i].startswith("--"): - var_name = args[i] - i += 1 + try: + var_name, no_display, limit = self._parse_cell_magic_args(line) + except ValueError as e: + return str(e) result = self._execute_sql(query) - # Store in user namespace if variable name provided + # The stored variable always holds the full, untruncated result. if var_name and self.shell is not None: self.shell.user_ns[var_name] = result + + if no_display: + return None + + # Display-only cap: limits the rows rendered in the cell, never the + # underlying data, so an in-query LIMIT always takes effect. Both + # min_rows and max_rows are set because the formatter requires + # min_rows <= max_rows (datafusion itself defaults them equal). + rows = limit if limit is not None else DEFAULT_DISPLAY_LIMIT + configure_formatter(max_rows=rows, min_rows=rows) return result def _connect(self, address: str) -> Optional[str]: @@ -487,7 +547,7 @@ def _show_help(self) -> Optional[str]: %%sql [options] [var] - Execute multi-line SQL query Options: - --no-display - Don't display results + --no-display - Run the query and store the result without displaying it --limit N - Limit displayed rows (default: 50) var - Store result in variable diff --git a/python/python/tests/test_jupyter.py b/python/python/tests/test_jupyter.py index 73c69ff4fb..20a8853502 100644 --- a/python/python/tests/test_jupyter.py +++ b/python/python/tests/test_jupyter.py @@ -241,6 +241,40 @@ def test_sql_magic_empty_line_returns_none(self, connected_magics): result = connected_magics.sql("") assert result is None + def test_sql_cell_magic_limit_and_no_display(self, connected_magics): + """%%sql honors --limit (display-only) and --no-display.""" + from datafusion.dataframe_formatter import get_formatter + from ballista.jupyter import DEFAULT_DISPLAY_LIMIT + + mock_shell = MagicMock() + mock_shell.user_ns = {} + connected_magics.shell = mock_shell + + # --limit caps the display formatter but never truncates the stored data. + connected_magics.sql( + "--limit 2 my_var", + cell="SELECT * FROM (VALUES (1),(2),(3),(4),(5)) AS t(x)", + ) + stored = mock_shell.user_ns["my_var"] + assert sum(batch.num_rows for batch in stored.collect()) == 5 + assert get_formatter().max_rows == 2 + + # Without --limit the formatter falls back to the default cap. + connected_magics.sql("", cell="SELECT 1 as value") + assert get_formatter().max_rows == DEFAULT_DISPLAY_LIMIT + + # --no-display stores the result but renders nothing. + result = connected_magics.sql( + "--no-display other_var", cell="SELECT 1 as value" + ) + assert result is None + assert mock_shell.user_ns["other_var"] is not None + + # An invalid --limit returns the error string instead of raising. + result = connected_magics.sql("--limit abc", cell="SELECT 1") + assert isinstance(result, str) + assert "--limit" in result + def test_schema_missing_table_name_returns_usage(self, connected_magics): """Test _schema with no table name returns usage string.""" result = connected_magics._schema("") @@ -275,6 +309,27 @@ def test_register_missing_file_path_returns_message(self, connected_magics): assert result is not None +class TestSqlCellMagicArgParsing: + """Cluster-free tests for %%sql argument parsing.""" + + def test_parse_cell_magic_args(self, magics): + # (line, expected (var_name, no_display, limit)) + valid_cases = [ + ("", (None, False, None)), + ("my_var", ("my_var", False, None)), + ("--no-display", (None, True, None)), + ("--limit 10 my_var", ("my_var", False, 10)), + ("my_var --no-display --limit 3", ("my_var", True, 3)), + ] + for line, expected in valid_cases: + assert magics._parse_cell_magic_args(line) == expected, line + + # Missing, non-integer, and non-positive --limit values are rejected. + for line in ("--limit", "--limit abc", "--limit 0", "--limit -1"): + with pytest.raises(ValueError): + magics._parse_cell_magic_args(line) + + class TestIPythonExtension: """Tests for IPython extension loading.""" From bf30f2901eb5c0408bfe3c1d622d708e4f7b50d4 Mon Sep 17 00:00:00 2001 From: Jogesh A Dinavahi Date: Mon, 13 Jul 2026 19:23:39 -0700 Subject: [PATCH 2/4] rebase + --limit option to jupyter cell --- python/python/ballista/jupyter.py | 34 ++++++++--- python/python/tests/test_jupyter.py | 88 +++++++++++++++++++++++------ 2 files changed, 99 insertions(+), 23 deletions(-) diff --git a/python/python/ballista/jupyter.py b/python/python/ballista/jupyter.py index dff19772ed..b1a81df2de 100644 --- a/python/python/ballista/jupyter.py +++ b/python/python/ballista/jupyter.py @@ -108,7 +108,11 @@ def decorator(func): return decorator -from datafusion import configure_formatter +from datafusion.dataframe_formatter import ( + configure_formatter, + get_formatter, + set_formatter, +) from .extension import BallistaSessionContext, DistributedDataFrame @@ -318,13 +322,29 @@ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]: if no_display: return None - # Display-only cap: limits the rows rendered in the cell, never the - # underlying data, so an in-query LIMIT always takes effect. Both - # min_rows and max_rows are set because the formatter requires - # min_rows <= max_rows (datafusion itself defaults them equal). + # Outside IPython there is no auto-render to cap or suppress, so + # just return the result rather than swallowing it. + if not IPYTHON_AVAILABLE: + return result + + # Display-only cap: limits the rows rendered for THIS cell, never + # the underlying data, so an in-query LIMIT always takes effect. + # datafusion's formatter is a process-global singleton, so we + # render eagerly with the cap applied and then restore the previous + # formatter, keeping the effect scoped to this cell instead of + # leaking into the rest of the session. Both min_rows and max_rows + # are set because the formatter requires min_rows <= max_rows. rows = limit if limit is not None else DEFAULT_DISPLAY_LIMIT - configure_formatter(max_rows=rows, min_rows=rows) - return result + previous_formatter = get_formatter() + try: + configure_formatter(max_rows=rows, min_rows=rows) + display(result) + finally: + set_formatter(previous_formatter) + + # Returning None avoids a second, un-capped auto-render by IPython; + # the result is still available via the variable and _last_result. + return None def _connect(self, address: str) -> Optional[str]: """Connect to a Ballista cluster.""" diff --git a/python/python/tests/test_jupyter.py b/python/python/tests/test_jupyter.py index 20a8853502..833f489a8d 100644 --- a/python/python/tests/test_jupyter.py +++ b/python/python/tests/test_jupyter.py @@ -221,10 +221,33 @@ def test_sql_magic_line_returns_dataframe(self, connected_magics): result = connected_magics.sql("SELECT 1 as value") assert result is not None - def test_sql_magic_cell_returns_dataframe(self, connected_magics): - """Test %%sql cell magic returns a DistributedDataFrame.""" - result = connected_magics.sql("", cell="SELECT 1 as value") + def test_sql_magic_cell_returns_dataframe_without_ipython(self, connected_magics): + """Outside IPython, %%sql returns the DataFrame instead of swallowing it. + + The ``return None`` path exists only to suppress IPython's second, + un-capped auto-render; with no IPython there is nothing to display or + suppress, so the result must be returned so it is not lost. + """ + with patch("ballista.jupyter.IPYTHON_AVAILABLE", False): + result = connected_magics.sql("", cell="SELECT 1 as value") assert result is not None + assert connected_magics._last_result is not None + + def test_sql_magic_cell_renders_and_returns_none_in_ipython(self, connected_magics): + """Inside IPython, %%sql renders (capped, cell-locally) and returns None. + + It displays the result itself and returns None to avoid a second + un-capped auto-render; the DataFrame remains available via + ``_last_result``. + """ + with ( + patch("ballista.jupyter.IPYTHON_AVAILABLE", True), + patch("ballista.jupyter.display", create=True) as mock_display, + ): + result = connected_magics.sql("", cell="SELECT 1 as value") + assert result is None + mock_display.assert_called_once() + assert connected_magics._last_result is not None def test_sql_magic_cell_stores_in_shell_namespace(self, connected_magics): """Test %%sql stores result in shell namespace when var name given.""" @@ -250,24 +273,57 @@ def test_sql_cell_magic_limit_and_no_display(self, connected_magics): mock_shell.user_ns = {} connected_magics.shell = mock_shell - # --limit caps the display formatter but never truncates the stored data. - connected_magics.sql( - "--limit 2 my_var", - cell="SELECT * FROM (VALUES (1),(2),(3),(4),(5)) AS t(x)", - ) + # The row cap is applied only while the cell renders and is reverted + # afterwards, so it never leaks into the rest of the session. + formatter_max_rows_before = get_formatter().max_rows + + def _capture_max_rows_at_display(_obj): + captured.append(get_formatter().max_rows) + + # --limit caps the rows rendered for this cell (checked at display + # time) but never truncates the stored data, and is scoped to the cell. + captured = [] + with ( + patch("ballista.jupyter.IPYTHON_AVAILABLE", True), + patch( + "ballista.jupyter.display", _capture_max_rows_at_display, create=True + ), + ): + result = connected_magics.sql( + "--limit 2 my_var", + cell="SELECT * FROM (VALUES (1),(2),(3),(4),(5)) AS t(x)", + ) stored = mock_shell.user_ns["my_var"] assert sum(batch.num_rows for batch in stored.collect()) == 5 - assert get_formatter().max_rows == 2 - - # Without --limit the formatter falls back to the default cap. - connected_magics.sql("", cell="SELECT 1 as value") - assert get_formatter().max_rows == DEFAULT_DISPLAY_LIMIT + assert captured == [2] # cap active while rendering + assert result is None # no second, un-capped auto-render + assert get_formatter().max_rows == formatter_max_rows_before # restored + + # Without --limit the default cap is used at render time, then restored. + captured = [] + with ( + patch("ballista.jupyter.IPYTHON_AVAILABLE", True), + patch( + "ballista.jupyter.display", _capture_max_rows_at_display, create=True + ), + ): + connected_magics.sql("", cell="SELECT 1 as value") + assert captured == [DEFAULT_DISPLAY_LIMIT] + assert get_formatter().max_rows == formatter_max_rows_before # restored # --no-display stores the result but renders nothing. - result = connected_magics.sql( - "--no-display other_var", cell="SELECT 1 as value" - ) + captured = [] + with ( + patch("ballista.jupyter.IPYTHON_AVAILABLE", True), + patch( + "ballista.jupyter.display", _capture_max_rows_at_display, create=True + ), + ): + result = connected_magics.sql( + "--no-display other_var", cell="SELECT 1 as value" + ) assert result is None + assert captured == [] # display was never called assert mock_shell.user_ns["other_var"] is not None # An invalid --limit returns the error string instead of raising. From 45eeeea1d658a7373590febf8a01719622dd45ff Mon Sep 17 00:00:00 2001 From: Jogesh A Dinavahi Date: Tue, 14 Jul 2026 23:12:22 -0700 Subject: [PATCH 3/4] store ipython constant flags separately --- python/python/ballista/jupyter.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/python/python/ballista/jupyter.py b/python/python/ballista/jupyter.py index b1a81df2de..6cc0c25f38 100644 --- a/python/python/ballista/jupyter.py +++ b/python/python/ballista/jupyter.py @@ -116,6 +116,11 @@ def decorator(func): from .extension import BallistaSessionContext, DistributedDataFrame +# Flags accepted by the ``%%sql`` cell magic, kept as constants so the parser +# and its error messages share a single source of truth. +LIMIT_FLAG = "--limit" +NO_DISPLAY_FLAG = "--no-display" + # Default number of rows rendered for a ``%%sql`` cell when ``--limit`` is not # given. This caps only the display (via datafusion's HTML formatter); the # underlying result keeps all of its rows. @@ -253,18 +258,22 @@ def _parse_cell_magic_args(line: str): i = 0 while i < len(tokens): token = tokens[i] - if token == "--no-display": + if token == NO_DISPLAY_FLAG: no_display = True - elif token == "--limit": + elif token == LIMIT_FLAG: i += 1 if i >= len(tokens): - raise ValueError("--limit requires a number, e.g. --limit 5") + raise ValueError( + f"{LIMIT_FLAG} requires a number, e.g. {LIMIT_FLAG} 5" + ) try: limit = int(tokens[i]) except ValueError: - raise ValueError(f"--limit expects an integer, got '{tokens[i]}'") + raise ValueError( + f"{LIMIT_FLAG} expects an integer, got '{tokens[i]}'" + ) if limit < 1: - raise ValueError("--limit must be a positive integer") + raise ValueError(f"{LIMIT_FLAG} must be a positive integer") elif not token.startswith("--") and var_name is None: var_name = token i += 1 From 16c46d4859c7d4e985bc7ccf1ec939a0f51a2a52 Mon Sep 17 00:00:00 2001 From: Jogesh A Dinavahi Date: Sun, 19 Jul 2026 19:14:39 -0700 Subject: [PATCH 4/4] address review comments: push limit clause to execution layer and not just display side --- python/README.md | 5 +- python/python/ballista/jupyter.py | 102 +++++++++------------ python/python/tests/test_jupyter.py | 134 +++++++++++----------------- 3 files changed, 97 insertions(+), 144 deletions(-) diff --git a/python/README.md b/python/README.md index d900b72262..c4d15d7ae5 100644 --- a/python/README.md +++ b/python/README.md @@ -149,8 +149,9 @@ SELECT * FROM orders WHERE status = 'pending' The cell magic accepts two optional flags before the variable name: ```python -# Render at most 10 rows in the cell (display only; the stored result keeps -# every row). Without --limit the display defaults to 50 rows. +# Add LIMIT 10 to the query: only 10 rows are computed, collected, and stored +# in my_result. A bare `--limit` uses the default (50); omit --limit entirely +# to return every row (or the query's own LIMIT). %%sql --limit 10 my_result SELECT * FROM orders diff --git a/python/python/ballista/jupyter.py b/python/python/ballista/jupyter.py index 6cc0c25f38..66302e8123 100644 --- a/python/python/ballista/jupyter.py +++ b/python/python/ballista/jupyter.py @@ -108,12 +108,6 @@ def decorator(func): return decorator -from datafusion.dataframe_formatter import ( - configure_formatter, - get_formatter, - set_formatter, -) - from .extension import BallistaSessionContext, DistributedDataFrame # Flags accepted by the ``%%sql`` cell magic, kept as constants so the parser @@ -121,10 +115,10 @@ def decorator(func): LIMIT_FLAG = "--limit" NO_DISPLAY_FLAG = "--no-display" -# Default number of rows rendered for a ``%%sql`` cell when ``--limit`` is not -# given. This caps only the display (via datafusion's HTML formatter); the -# underlying result keeps all of its rows. -DEFAULT_DISPLAY_LIMIT = 50 +# LIMIT applied when ``--limit`` is given without an explicit number +# (e.g. ``%%sql --limit my_var``). ``--limit N`` overrides it; omitting +# ``--limit`` entirely returns every row. +DEFAULT_LIMIT = 50 class BallistaConnectionError(Exception): @@ -241,14 +235,14 @@ def register(self, line: str) -> Optional[str]: def _parse_cell_magic_args(line: str): """Parse the argument line of a ``%%sql`` cell magic. - Recognises the ``--no-display`` and ``--limit N`` flags (space form, - e.g. ``--limit 5``, consistent with the other magics in this module). - The first non-flag token, if any, is treated as the variable name to - store the result in. + Recognises the ``--no-display`` and ``--limit`` flags. ``--limit`` may + be given bare (uses ``DEFAULT_LIMIT``) or with an explicit count + (``--limit 5``). The first non-flag token, if any, is the variable name + to store the result in. Returns a ``(var_name, no_display, limit)`` tuple where ``limit`` is ``None`` when ``--limit`` was not supplied. Raises ``ValueError`` for a - missing or invalid ``--limit`` value. + non-positive explicit ``--limit`` value. """ tokens = line.strip().split() var_name = None @@ -261,19 +255,21 @@ def _parse_cell_magic_args(line: str): if token == NO_DISPLAY_FLAG: no_display = True elif token == LIMIT_FLAG: - i += 1 - if i >= len(tokens): - raise ValueError( - f"{LIMIT_FLAG} requires a number, e.g. {LIMIT_FLAG} 5" - ) - try: - limit = int(tokens[i]) - except ValueError: - raise ValueError( - f"{LIMIT_FLAG} expects an integer, got '{tokens[i]}'" - ) - if limit < 1: - raise ValueError(f"{LIMIT_FLAG} must be a positive integer") + # Bare --limit uses the default. A following integer overrides + # it; a following non-integer (e.g. the variable name) is left + # for the var-name branch, so `--limit my_var` is a default + # LIMIT stored in `my_var`. + limit = DEFAULT_LIMIT + if i + 1 < len(tokens): + try: + explicit = int(tokens[i + 1]) + except ValueError: + explicit = None + if explicit is not None: + if explicit < 1: + raise ValueError(f"{LIMIT_FLAG} must be a positive integer") + limit = explicit + i += 1 # consume the number token elif not token.startswith("--") and var_name is None: var_name = token i += 1 @@ -304,9 +300,11 @@ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]: `my_result` will store the result of the SQL-query The cell magic accepts two optional flags before the variable name: - --limit N Render at most N rows in the cell output (default - 50). This caps the display only; the stored result - keeps every row. + --limit [N] Add ``LIMIT`` to the query, bounding the rows that + are computed and stored. ``--limit N`` uses N; + a bare ``--limit`` uses the default (50). Omit it + entirely to return every row (or whatever ``LIMIT`` + the query itself specifies). --no-display Run the query and store the result without displaying it. """ @@ -323,37 +321,23 @@ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]: return str(e) result = self._execute_sql(query) + if result is None: + return None + + # --limit pushes a LIMIT into the query plan (not a display cap), so + # only N rows are computed, collected, and stored. + if limit is not None: + result = result.limit(limit) + self._last_result = result + if self.shell is not None and hasattr(self.shell, "user_ns"): + self.shell.user_ns["_last_result"] = result - # The stored variable always holds the full, untruncated result. if var_name and self.shell is not None: self.shell.user_ns[var_name] = result - if no_display: - return None - - # Outside IPython there is no auto-render to cap or suppress, so - # just return the result rather than swallowing it. - if not IPYTHON_AVAILABLE: - return result - - # Display-only cap: limits the rows rendered for THIS cell, never - # the underlying data, so an in-query LIMIT always takes effect. - # datafusion's formatter is a process-global singleton, so we - # render eagerly with the cap applied and then restore the previous - # formatter, keeping the effect scoped to this cell instead of - # leaking into the rest of the session. Both min_rows and max_rows - # are set because the formatter requires min_rows <= max_rows. - rows = limit if limit is not None else DEFAULT_DISPLAY_LIMIT - previous_formatter = get_formatter() - try: - configure_formatter(max_rows=rows, min_rows=rows) - display(result) - finally: - set_formatter(previous_formatter) - - # Returning None avoids a second, un-capped auto-render by IPython; - # the result is still available via the variable and _last_result. - return None + # Returning the DataFrame lets IPython render it (datafusion's normal + # preview); --no-display stores it but suppresses that render. + return None if no_display else result def _connect(self, address: str) -> Optional[str]: """Connect to a Ballista cluster.""" @@ -577,7 +561,7 @@ def _show_help(self) -> Optional[str]: %%sql [options] [var] - Execute multi-line SQL query Options: --no-display - Run the query and store the result without displaying it - --limit N - Limit displayed rows (default: 50) + --limit [N] - Add LIMIT to the query (N, or default 50 if bare) var - Store result in variable History: diff --git a/python/python/tests/test_jupyter.py b/python/python/tests/test_jupyter.py index 833f489a8d..4fa72f763b 100644 --- a/python/python/tests/test_jupyter.py +++ b/python/python/tests/test_jupyter.py @@ -221,34 +221,12 @@ def test_sql_magic_line_returns_dataframe(self, connected_magics): result = connected_magics.sql("SELECT 1 as value") assert result is not None - def test_sql_magic_cell_returns_dataframe_without_ipython(self, connected_magics): - """Outside IPython, %%sql returns the DataFrame instead of swallowing it. - - The ``return None`` path exists only to suppress IPython's second, - un-capped auto-render; with no IPython there is nothing to display or - suppress, so the result must be returned so it is not lost. - """ - with patch("ballista.jupyter.IPYTHON_AVAILABLE", False): - result = connected_magics.sql("", cell="SELECT 1 as value") + def test_sql_magic_cell_returns_dataframe(self, connected_magics): + """%%sql cell magic returns the DataFrame for IPython to render.""" + result = connected_magics.sql("", cell="SELECT 1 as value") assert result is not None assert connected_magics._last_result is not None - def test_sql_magic_cell_renders_and_returns_none_in_ipython(self, connected_magics): - """Inside IPython, %%sql renders (capped, cell-locally) and returns None. - - It displays the result itself and returns None to avoid a second - un-capped auto-render; the DataFrame remains available via - ``_last_result``. - """ - with ( - patch("ballista.jupyter.IPYTHON_AVAILABLE", True), - patch("ballista.jupyter.display", create=True) as mock_display, - ): - result = connected_magics.sql("", cell="SELECT 1 as value") - assert result is None - mock_display.assert_called_once() - assert connected_magics._last_result is not None - def test_sql_magic_cell_stores_in_shell_namespace(self, connected_magics): """Test %%sql stores result in shell namespace when var name given.""" mock_shell = MagicMock() @@ -264,70 +242,53 @@ def test_sql_magic_empty_line_returns_none(self, connected_magics): result = connected_magics.sql("") assert result is None - def test_sql_cell_magic_limit_and_no_display(self, connected_magics): - """%%sql honors --limit (display-only) and --no-display.""" - from datafusion.dataframe_formatter import get_formatter - from ballista.jupyter import DEFAULT_DISPLAY_LIMIT - + def test_sql_cell_magic_limit_bounds_the_query(self, connected_magics): + """--limit N pushes LIMIT N into the query, bounding the stored result.""" mock_shell = MagicMock() mock_shell.user_ns = {} connected_magics.shell = mock_shell - # The row cap is applied only while the cell renders and is reverted - # afterwards, so it never leaks into the rest of the session. - formatter_max_rows_before = get_formatter().max_rows - - def _capture_max_rows_at_display(_obj): - captured.append(get_formatter().max_rows) - - # --limit caps the rows rendered for this cell (checked at display - # time) but never truncates the stored data, and is scoped to the cell. - captured = [] - with ( - patch("ballista.jupyter.IPYTHON_AVAILABLE", True), - patch( - "ballista.jupyter.display", _capture_max_rows_at_display, create=True - ), - ): - result = connected_magics.sql( - "--limit 2 my_var", - cell="SELECT * FROM (VALUES (1),(2),(3),(4),(5)) AS t(x)", - ) + five_rows = "SELECT * FROM (VALUES (1),(2),(3),(4),(5)) AS t(x)" + + # --limit 2 on a 5-row source: only 2 rows are computed and stored. + connected_magics.sql("--limit 2 my_var", cell=five_rows) stored = mock_shell.user_ns["my_var"] - assert sum(batch.num_rows for batch in stored.collect()) == 5 - assert captured == [2] # cap active while rendering - assert result is None # no second, un-capped auto-render - assert get_formatter().max_rows == formatter_max_rows_before # restored - - # Without --limit the default cap is used at render time, then restored. - captured = [] - with ( - patch("ballista.jupyter.IPYTHON_AVAILABLE", True), - patch( - "ballista.jupyter.display", _capture_max_rows_at_display, create=True - ), - ): - connected_magics.sql("", cell="SELECT 1 as value") - assert captured == [DEFAULT_DISPLAY_LIMIT] - assert get_formatter().max_rows == formatter_max_rows_before # restored - - # --no-display stores the result but renders nothing. - captured = [] - with ( - patch("ballista.jupyter.IPYTHON_AVAILABLE", True), - patch( - "ballista.jupyter.display", _capture_max_rows_at_display, create=True - ), - ): - result = connected_magics.sql( - "--no-display other_var", cell="SELECT 1 as value" - ) + assert sum(batch.num_rows for batch in stored.collect()) == 2 + + # Without --limit the full result is stored (all 5 rows). + connected_magics.sql("full_var", cell=five_rows) + full = mock_shell.user_ns["full_var"] + assert sum(batch.num_rows for batch in full.collect()) == 5 + + def test_sql_cell_magic_bare_limit_uses_default(self): + """A bare --limit uses DEFAULT_LIMIT; a non-integer next token is the var.""" + from ballista.jupyter import DEFAULT_LIMIT + + parse = BallistaMagics._parse_cell_magic_args + # bare --limit, next token is the variable name + assert parse("--limit my_var") == ("my_var", False, DEFAULT_LIMIT) + # bare --limit, no variable + assert parse("--limit") == (None, False, DEFAULT_LIMIT) + # explicit number overrides the default + assert parse("--limit 3 v") == ("v", False, 3) + # no --limit at all -> no limit (all rows) + assert parse("v") == ("v", False, None) + + def test_sql_cell_magic_no_display(self, connected_magics): + """--no-display stores the result but returns None (renders nothing).""" + mock_shell = MagicMock() + mock_shell.user_ns = {} + connected_magics.shell = mock_shell + + result = connected_magics.sql( + "--no-display other_var", cell="SELECT 1 as value" + ) assert result is None - assert captured == [] # display was never called assert mock_shell.user_ns["other_var"] is not None - # An invalid --limit returns the error string instead of raising. - result = connected_magics.sql("--limit abc", cell="SELECT 1") + def test_sql_cell_magic_rejects_non_positive_limit(self, connected_magics): + """An explicit non-positive --limit returns the error string.""" + result = connected_magics.sql("--limit 0", cell="SELECT 1") assert isinstance(result, str) assert "--limit" in result @@ -369,6 +330,8 @@ class TestSqlCellMagicArgParsing: """Cluster-free tests for %%sql argument parsing.""" def test_parse_cell_magic_args(self, magics): + from ballista.jupyter import DEFAULT_LIMIT + # (line, expected (var_name, no_display, limit)) valid_cases = [ ("", (None, False, None)), @@ -376,12 +339,17 @@ def test_parse_cell_magic_args(self, magics): ("--no-display", (None, True, None)), ("--limit 10 my_var", ("my_var", False, 10)), ("my_var --no-display --limit 3", ("my_var", True, 3)), + # bare --limit uses the default... + ("--limit", (None, False, DEFAULT_LIMIT)), + # ...and a non-integer next token is the variable name, not the count + ("--limit my_var", ("my_var", False, DEFAULT_LIMIT)), + ("--limit abc", ("abc", False, DEFAULT_LIMIT)), ] for line, expected in valid_cases: assert magics._parse_cell_magic_args(line) == expected, line - # Missing, non-integer, and non-positive --limit values are rejected. - for line in ("--limit", "--limit abc", "--limit 0", "--limit -1"): + # Only an explicit non-positive --limit value is rejected. + for line in ("--limit 0", "--limit -1"): with pytest.raises(ValueError): magics._parse_cell_magic_args(line)