feat: add %%sql --limit and --no-display support

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:

python/python/ballista/jupyter.py

@@ -108,8 +108,19 @@ def decorator(func):
             return decorator
 
 
+from datafusion.dataframe_formatter import (
+    configure_formatter,
+    get_formatter,
+    set_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 +232,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,28 +293,58 @@ 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
-            return 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
 
     def _connect(self, address: str) -> Optional[str]:
         """Connect to a Ballista cluster."""
@@ -487,7 +567,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
 

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."""
@@ -241,6 +264,73 @@ 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
+
+        # 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 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 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")
+        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 +365,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."""