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

python/README.md

@@ -146,6 +146,20 @@ 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
+# 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
+
+# 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

@@ -110,6 +110,16 @@ 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"
+
+# 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):
     """Raised when not connected to a Ballista cluster."""
@@ -221,6 +231,51 @@ 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`` 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
+        non-positive explicit ``--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_FLAG:
+                no_display = True
+            elif token == LIMIT_FLAG:
+                # 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
+
+        return var_name, no_display, limit
+
     @line_cell_magic
     def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
         """
@@ -243,28 +298,46 @@ 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]    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.
         """
         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)
+            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
 
-            # Store in user namespace if variable name provided
             if var_name and self.shell is not None:
                 self.shell.user_ns[var_name] = result
-            return result
+
+            # 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."""
@@ -487,8 +560,8 @@ def _show_help(self) -> Optional[str]:
 
     %%sql [options] [var]     - Execute multi-line SQL query
         Options:
-            --no-display      - Don't display results
-            --limit N         - Limit displayed rows (default: 50)
+            --no-display      - Run the query and store the result without displaying it
+            --limit [N]       - Add LIMIT to the query (N, or default 50 if bare)
         var                   - Store result in variable
 
 History:

python/python/tests/test_jupyter.py

@@ -222,9 +222,10 @@ def test_sql_magic_line_returns_dataframe(self, connected_magics):
         assert result is not None
 
     def test_sql_magic_cell_returns_dataframe(self, connected_magics):
-        """Test %%sql cell magic returns a DistributedDataFrame."""
+        """%%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_stores_in_shell_namespace(self, connected_magics):
         """Test %%sql stores result in shell namespace when var name given."""
@@ -241,6 +242,56 @@ 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_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
+
+        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()) == 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 mock_shell.user_ns["other_var"] is not None
+
+    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
+
     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 +326,34 @@ 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):
+        from ballista.jupyter import DEFAULT_LIMIT
+
+        # (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)),
+            # 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
+
+        # 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)
+
+
 class TestIPythonExtension:
     """Tests for IPython extension loading."""