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

.cursor/rules.md

@@ -0,0 +1,5 @@
+When making a code review never link to any issue or a pull request
+in any GitHub repository!
+
+Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md!
+

.gemini/rules.md

@@ -0,0 +1,5 @@
+When making a code review never link to any issue or a pull request
+in any GitHub repository!
+
+Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md!
+

AGENTS.md

@@ -0,0 +1,5 @@
+When making a code review never link to any issue or a pull request
+in any GitHub repository!
+
+Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md!
+

CLAUDE.md

@@ -0,0 +1,5 @@
+When making a code review never link to any issue or a pull request
+in any GitHub repository!
+
+Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md!
+

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,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
 

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."""