diff --git a/docs/sparksql.md b/docs/sparksql.md index 34d1107..463e76d 100644 --- a/docs/sparksql.md +++ b/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 diff --git a/google/cloud/managed_spark_magics/sparksql.py b/google/cloud/managed_spark_magics/sparksql.py index fd4c1f0..7c2ff34 100644 --- a/google/cloud/managed_spark_magics/sparksql.py +++ b/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 diff --git a/tests/local/test_sparksql_local.py b/tests/local/test_sparksql_local.py index b52315d..a004804 100644 --- a/tests/local/test_sparksql_local.py +++ b/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): diff --git a/tests/unit/managed_spark_magics/test_sparksql.py b/tests/unit/managed_spark_magics/test_sparksql.py index 6183c67..6c88e54 100644 --- a/tests/unit/managed_spark_magics/test_sparksql.py +++ b/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"