feat: store %%sparksql results in _df by default

docs/sparksql.md

@@ -34,9 +34,9 @@ To cap the number of rows displayed, pass `-l N` per cell or set:
 ```
 
 Options go on the first line, optionally naming a variable to store the
-resulting DataFrame in. When a variable name is provided, it receives the full,
-un-truncated `DataFrame` (the `-l` display limit only affects what is shown on
-screen):
+resulting DataFrame in (by default, `%%sparksql` stores it in `_df`). The
+variable receives the full, un-truncated `DataFrame` (the `-l` display limit
+only affects what is shown on screen):
 
 ```python
 %%sparksql --cache --view result_view df

google/cloud/managed_spark_magics/sparksql.py

@@ -59,6 +59,7 @@
 EXPLAIN_MODES = ("simple", "extended", "codegen", "cost", "formatted")
 DISPLAY_MODES = ("colabsqlviz", "viz", "html")
 DEFAULT_HTML_LIMIT = 20
+DEFAULT_VARIABLE = "_df"
 
 _IDENT_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z")
 _PLACEHOLDER_RE = re.compile(r":([A-Za-z_][A-Za-z0-9_]*)\Z")
@@ -187,10 +188,10 @@ def _usage_error(msg: str) -> UsageError:
 def parse_options(line: str, line_form: bool) -> Tuple[Options, str]:
     """Parses magic options.
 
-    For the cell form, `line` holds only options and an optional variable name,
-    and the returned rest is always empty. For the line form, options end at
-    the first token that is not an option, and everything from there on is
-    returned verbatim as the SQL.
+    For the cell form, `line` holds only options and an optional variable name
+    (defaulting to `_df`), and the returned rest is always empty. For the line
+    form, options end at the first token that is not an option, and everything
+    from there on is returned verbatim as the SQL.
 
     This is hand-rolled rather than built on argparse or
     IPython.core.magic_arguments for two reasons:
@@ -310,6 +311,8 @@ def take_value(flag: str) -> str:
         if not rest.strip():
             raise _usage_error("No SQL given")
         opts.quiet = True
+    elif opts.variable is None:
+        opts.variable = DEFAULT_VARIABLE
     return opts, rest
 
 
@@ -1064,7 +1067,8 @@ def __init__(self, shell: Any, **kwargs: Any) -> None:
     def sparksql(self, line: str, cell: Optional[str] = None) -> Any:
         """Runs Spark SQL against the active SparkSession.
 
-        Cell form: displays the result of the last statement.
+        Cell form: displays the result of the last statement and stores the
+        DataFrame in `variable` (default: `_df`).
 
             %%sparksql [-s SESSION] [-d DISPLAY] [-l LIMIT] [-v VIEW] [-c] [-e]
                        [-q] [--explain [MODE]] [variable]
@@ -1077,7 +1081,8 @@ def sparksql(self, line: str, cell: Optional[str] = None) -> Any:
 
         Options:
           variable          Store the un-truncated DataFrame in this variable
-                            (cell form only; -l only limits displayed rows).
+                            (cell form only; default: _df; -l only limits
+                            displayed rows).
           -s, --session N   Use the SparkSession in variable N. Default: the
                             active session, else the only one in the notebook.
           -d, --display M   How to display results: colabsqlviz (or viz) for an

tests/local/test_sparksql_local.py

@@ -151,6 +151,7 @@ def test_select_renders_colabsqlviz_table(ip, displayed):
 
     ip.run_cell_magic("sparksql", "-l 2", "SELECT * FROM range(25)")
     assert len(displayed[1]._df) == 2
+    assert ip.user_ns["_df"].count() == 25
 
 
 def test_select_renders_html_table(ip, displayed):

tests/unit/managed_spark_magics/test_sparksql.py

@@ -34,7 +34,7 @@ class ParseOptionsTest(unittest.TestCase):
 
     def test_cell_defaults(self):
         opts, rest = sparksql.parse_options("", line_form=False)
-        self.assertEqual(opts, sparksql.Options())
+        self.assertEqual(opts, sparksql.Options(variable="_df"))
         self.assertEqual(rest, "")
 
     def test_cell_all_options(self):
@@ -78,6 +78,7 @@ def test_explain_mode_is_optional(self):
         self.assertEqual(opts.variable, "df")
         opts, _ = sparksql.parse_options("--explain=EXTENDED", line_form=False)
         self.assertEqual(opts.explain, "extended")
+        self.assertEqual(opts.variable, "_df")
 
     def test_errors(self):
         for line, msg in [
@@ -112,6 +113,7 @@ def test_line_form(self):
         )
         self.assertTrue(opts.cache)
         self.assertTrue(opts.quiet)
+        self.assertIsNone(opts.variable)
         # The SQL is everything from the first non-option token, verbatim.
         self.assertEqual(rest, "SELECT  -1 AS x  -- c")
 
@@ -793,6 +795,7 @@ def test_cell_runs_all_and_displays_last(self):
             ["CREATE VIEW v AS SELECT 1", "SELECT :n FROM {df}"],
         )
         df = self.shell.user_ns["out"]
+        self.assertNotIn("_df", self.shell.user_ns)
         self.assertEqual(df.args, {"n": 3})
         self.assertEqual(df.kwargs, {"df": input_df})
         self.display.assert_called_once_with(
@@ -802,6 +805,13 @@ def test_cell_runs_all_and_displays_last(self):
             explicit_display=False,
         )
 
+    def test_cell_stores_in_underscore_df_by_default(self):
+        result = self.magics.sparksql("", "SELECT 1")
+        self.assertIsNone(result)
+        df = self.shell.user_ns["_df"]
+        self.assertIsInstance(df, FakeDataFrame)
+        self.assertEqual(df.sql, "SELECT 1")
+
     def test_limit_option_and_config(self):
         self.magics.sparksql("-l 5", "SELECT 1")
         self.display.assert_called_with(
@@ -873,6 +883,7 @@ def test_line_form_returns_dataframe(self):
         self.assertIsInstance(df, FakeDataFrame)
         self.assertEqual(df.sql, "SELECT * FROM {x}")
         self.assertEqual(df.kwargs, {"x": input_df})
+        self.assertNotIn("_df", self.shell.user_ns)
         self.display.assert_not_called()
 
     def test_missing_variable(self):
@@ -909,6 +920,9 @@ def test_variable_not_bound_when_display_fails(self):
         with self.assertRaises(sparksql.SparkSqlError):
             self.magics.sparksql("out", "SELECT 1")
         self.assertNotIn("out", self.shell.user_ns)
+        with self.assertRaises(sparksql.SparkSqlError):
+            self.magics.sparksql("", "SELECT 1")
+        self.assertNotIn("_df", self.shell.user_ns)
 
     def test_error_in_later_statement(self):
         self.session.fail_on = "boom"