From 1a0b1119bd69d46f11fff02ccbdde7603a0cac25 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Thu, 22 Jan 2026 02:04:22 +0000 Subject: [PATCH 1/3] feat: Add a bigframes cell magic for ipython --- bigframes/__init__.py | 12 ++++ bigframes/_magics.py | 51 +++++++++++++++++ tests/system/small/test_magics.py | 95 +++++++++++++++++++++++++++++++ 3 files changed, 158 insertions(+) create mode 100644 bigframes/_magics.py create mode 100644 tests/system/small/test_magics.py diff --git a/bigframes/__init__.py b/bigframes/__init__.py index 240608ebc2d..3db0616524f 100644 --- a/bigframes/__init__.py +++ b/bigframes/__init__.py @@ -16,12 +16,24 @@ from bigframes._config import option_context, options from bigframes._config.bigquery_options import BigQueryOptions +from bigframes._magics import _cell_magic from bigframes.core.global_session import close_session, get_global_session import bigframes.enums as enums import bigframes.exceptions as exceptions from bigframes.session import connect, Session from bigframes.version import __version__ +_MAGIC_NAMES = ["bigframes"] + + +def load_ipython_extension(ipython): + """Called by IPython when this module is loaded as an IPython extension.""" + for magic_name in _MAGIC_NAMES: + ipython.register_magic_function( + _cell_magic, magic_kind="cell", magic_name=magic_name + ) + + __all__ = [ "options", "BigQueryOptions", diff --git a/bigframes/_magics.py b/bigframes/_magics.py new file mode 100644 index 00000000000..2c661e1aa8b --- /dev/null +++ b/bigframes/_magics.py @@ -0,0 +1,51 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from IPython.core import magic_arguments # type: ignore +from IPython.core.getipython import get_ipython +from IPython.display import display + +import bigframes.pandas + + +@magic_arguments.magic_arguments() +@magic_arguments.argument( + "destination_var", + nargs="?", + help=("If provided, save the output to this variable instead of displaying it."), +) +@magic_arguments.argument( + "--dry_run", + action="store_true", + default=False, + help=( + "Sets query to be a dry run to estimate costs. " + "Defaults to executing the query instead of dry run if this argument is not used." + "Does not work with engine 'bigframes'. " + ), +) +def _cell_magic(line, cell): + ipython = get_ipython() + args = magic_arguments.parse_argstring(_cell_magic, line) + if not cell: + print("Query is missing.") + pyformat_args = ipython.user_ns + dataframe = bigframes.pandas._read_gbq_colab( + cell, pyformat_args=pyformat_args, dry_run=args.dry_run + ) + if args.destination_var: + ipython.push({args.destination_var: dataframe}) + else: + display(dataframe) + return dataframe diff --git a/tests/system/small/test_magics.py b/tests/system/small/test_magics.py new file mode 100644 index 00000000000..d725dc8bbf3 --- /dev/null +++ b/tests/system/small/test_magics.py @@ -0,0 +1,95 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from IPython.testing.globalipapp import get_ipython +from IPython.utils.capture import capture_output +import pandas as pd +import pytest + +import bigframes +import bigframes.pandas as bpd + +MAGIC_NAME = "bigframes" + + +@pytest.fixture(scope="module") +def ip(): + """Provides a persistent IPython shell instance for the test session.""" + shell = get_ipython() + shell.extension_manager.load_extension("bigframes") + return shell + + +def test_magic_select_lit_to_var(ip): + bigframes.close_session() + + line = "dst_var" + cell_body = "SELECT 3" + + ip.run_cell_magic(MAGIC_NAME, line, cell_body) + + assert "dst_var" in ip.user_ns + result_df = ip.user_ns["dst_var"] + assert result_df.shape == (1, 1) + assert result_df.loc[0, 0] == 3 + + +def test_magic_select_lit_dry_run(ip): + bigframes.close_session() + + line = "dst_var --dry_run" + cell_body = "SELECT 3" + + ip.run_cell_magic(MAGIC_NAME, line, cell_body) + + assert "dst_var" in ip.user_ns + result_df = ip.user_ns["dst_var"] + assert result_df.totalBytesProcessed == 0 + + +def test_magic_select_lit_display(ip): + bigframes.close_session() + + cell_body = "SELECT 3" + + with capture_output() as io: + ip.run_cell_magic(MAGIC_NAME, "", cell_body) + assert len(io.outputs) > 0 + html_data = io.outputs[0].data["text/html"] + assert "[1 rows x 1 columns in total]" in html_data + + +def test_magic_select_interpolate(ip): + bigframes.close_session() + df = bpd.read_pandas( + pd.DataFrame({"col_a": [1, 2, 3, 4, 5, 6], "col_b": [1, 2, 1, 3, 1, 2]}) + ) + const_val = 1 + + ip.push({"df": df, "const_val": const_val}) + + query = """ + SELECT + SUM(col_a) AS total + FROM + {df} + WHERE col_b={const_val} + """ + + ip.run_cell_magic(MAGIC_NAME, "dst_var", query) + + assert "dst_var" in ip.user_ns + result_df = ip.user_ns["dst_var"] + assert result_df.shape == (1, 1) + assert result_df.loc[0, 0] == 9 From b599d81714ab1bb72745ea1ef708a0a8018bc244 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Thu, 22 Jan 2026 18:51:00 +0000 Subject: [PATCH 2/3] make robust to polars missing, change magics name --- bigframes/__init__.py | 2 +- bigframes/_magics.py | 9 +++++++-- bigframes/pandas/io/api.py | 9 +++++++-- tests/system/small/test_magics.py | 2 +- 4 files changed, 16 insertions(+), 6 deletions(-) diff --git a/bigframes/__init__.py b/bigframes/__init__.py index 3db0616524f..f282e092e4b 100644 --- a/bigframes/__init__.py +++ b/bigframes/__init__.py @@ -23,7 +23,7 @@ from bigframes.session import connect, Session from bigframes.version import __version__ -_MAGIC_NAMES = ["bigframes"] +_MAGIC_NAMES = ["bigquery_sql"] def load_ipython_extension(ipython): diff --git a/bigframes/_magics.py b/bigframes/_magics.py index 2c661e1aa8b..33205e3a372 100644 --- a/bigframes/_magics.py +++ b/bigframes/_magics.py @@ -40,6 +40,7 @@ def _cell_magic(line, cell): args = magic_arguments.parse_argstring(_cell_magic, line) if not cell: print("Query is missing.") + return pyformat_args = ipython.user_ns dataframe = bigframes.pandas._read_gbq_colab( cell, pyformat_args=pyformat_args, dry_run=args.dry_run @@ -47,5 +48,9 @@ def _cell_magic(line, cell): if args.destination_var: ipython.push({args.destination_var: dataframe}) else: - display(dataframe) - return dataframe + with bigframes.option_context( + "display.repr_mode", + "anywidget", + ): + display(dataframe) + return diff --git a/bigframes/pandas/io/api.py b/bigframes/pandas/io/api.py index 483bc5e530d..600fd9b9bce 100644 --- a/bigframes/pandas/io/api.py +++ b/bigframes/pandas/io/api.py @@ -49,6 +49,7 @@ import pyarrow as pa import bigframes._config as config +import bigframes._importing import bigframes.core.global_session as global_session import bigframes.core.indexes import bigframes.dataframe @@ -356,8 +357,12 @@ def _read_gbq_colab( with warnings.catch_warnings(): # Don't warning about Polars in SQL cell. # Related to b/437090788. - warnings.simplefilter("ignore", bigframes.exceptions.PreviewWarning) - config.options.bigquery.enable_polars_execution = True + try: + bigframes._importing.import_polars() + warnings.simplefilter("ignore", bigframes.exceptions.PreviewWarning) + config.options.bigquery.enable_polars_execution = True + except TypeError: + pass # don't fail if polars isn't available return global_session.with_default_session( bigframes.session.Session._read_gbq_colab, diff --git a/tests/system/small/test_magics.py b/tests/system/small/test_magics.py index d725dc8bbf3..aed759e4ac9 100644 --- a/tests/system/small/test_magics.py +++ b/tests/system/small/test_magics.py @@ -20,7 +20,7 @@ import bigframes import bigframes.pandas as bpd -MAGIC_NAME = "bigframes" +MAGIC_NAME = "bigquery_sql" @pytest.fixture(scope="module") From 0d764757516642caee6db16c10706e4dc6054e21 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Thu, 22 Jan 2026 21:13:05 +0000 Subject: [PATCH 3/3] fix polars import error catch --- bigframes/pandas/io/api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bigframes/pandas/io/api.py b/bigframes/pandas/io/api.py index 600fd9b9bce..b8ad6cbd0c4 100644 --- a/bigframes/pandas/io/api.py +++ b/bigframes/pandas/io/api.py @@ -361,7 +361,7 @@ def _read_gbq_colab( bigframes._importing.import_polars() warnings.simplefilter("ignore", bigframes.exceptions.PreviewWarning) config.options.bigquery.enable_polars_execution = True - except TypeError: + except ImportError: pass # don't fail if polars isn't available return global_session.with_default_session(