diff --git a/python/README.md b/python/README.md index c357cccc16..c4d15d7ae5 100644 --- a/python/README.md +++ b/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: diff --git a/python/python/ballista/jupyter.py b/python/python/ballista/jupyter.py index d1d8c38ed9..66302e8123 100644 --- a/python/python/ballista/jupyter.py +++ b/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: diff --git a/python/python/tests/test_jupyter.py b/python/python/tests/test_jupyter.py index 73c69ff4fb..4fa72f763b 100644 --- a/python/python/tests/test_jupyter.py +++ b/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."""