Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings
This repository was archived by the owner on May 7, 2026. It is now read-only.

Commit e6de52d

Browse filesBrowse files
TrevorBergeronTrevorBergeron
authored
feat: Add a bigframes cell magic for ipython (#2395)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> 🦕
1 parent 318752a commit e6de52d
Copy full SHA for e6de52d

4 files changed

+177-2Lines changed: 177 additions & 2 deletions

File tree

Expand file treeCollapse file tree
Open diff view settings
Filter options
Expand file treeCollapse file tree
Open diff view settings
Collapse file

‎bigframes/__init__.py‎

Copy file name to clipboardExpand all lines: bigframes/__init__.py
+14Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,20 @@
2222
from bigframes.session import connect, Session
2323
from bigframes.version import __version__
2424

25+
_MAGIC_NAMES = ["bqsql"]
26+
27+
28+
def load_ipython_extension(ipython):
29+
"""Called by IPython when this module is loaded as an IPython extension."""
30+
# Requires IPython to be installed for import to succeed
31+
from bigframes._magics import _cell_magic
32+
33+
for magic_name in _MAGIC_NAMES:
34+
ipython.register_magic_function(
35+
_cell_magic, magic_kind="cell", magic_name=magic_name
36+
)
37+
38+
2539
__all__ = [
2640
"options",
2741
"BigQueryOptions",
Collapse file

‎bigframes/_magics.py‎

Copy file name to clipboard
+56Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from IPython.core import magic_arguments # type: ignore
16+
from IPython.core.getipython import get_ipython
17+
from IPython.display import display
18+
19+
import bigframes.pandas
20+
21+
22+
@magic_arguments.magic_arguments()
23+
@magic_arguments.argument(
24+
"destination_var",
25+
nargs="?",
26+
help=("If provided, save the output to this variable instead of displaying it."),
27+
)
28+
@magic_arguments.argument(
29+
"--dry_run",
30+
action="store_true",
31+
default=False,
32+
help=(
33+
"Sets query to be a dry run to estimate costs. "
34+
"Defaults to executing the query instead of dry run if this argument is not used."
35+
"Does not work with engine 'bigframes'. "
36+
),
37+
)
38+
def _cell_magic(line, cell):
39+
ipython = get_ipython()
40+
args = magic_arguments.parse_argstring(_cell_magic, line)
41+
if not cell:
42+
print("Query is missing.")
43+
return
44+
pyformat_args = ipython.user_ns
45+
dataframe = bigframes.pandas._read_gbq_colab(
46+
cell, pyformat_args=pyformat_args, dry_run=args.dry_run
47+
)
48+
if args.destination_var:
49+
ipython.push({args.destination_var: dataframe})
50+
else:
51+
with bigframes.option_context(
52+
"display.repr_mode",
53+
"anywidget",
54+
):
55+
display(dataframe)
56+
return
Collapse file

‎bigframes/pandas/io/api.py‎

Copy file name to clipboardExpand all lines: bigframes/pandas/io/api.py
+7-2Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
import pyarrow as pa
5050

5151
import bigframes._config as config
52+
import bigframes._importing
5253
import bigframes.core.global_session as global_session
5354
import bigframes.core.indexes
5455
import bigframes.dataframe
@@ -356,8 +357,12 @@ def _read_gbq_colab(
356357
with warnings.catch_warnings():
357358
# Don't warning about Polars in SQL cell.
358359
# Related to b/437090788.
359-
warnings.simplefilter("ignore", bigframes.exceptions.PreviewWarning)
360-
config.options.bigquery.enable_polars_execution = True
360+
try:
361+
bigframes._importing.import_polars()
362+
warnings.simplefilter("ignore", bigframes.exceptions.PreviewWarning)
363+
config.options.bigquery.enable_polars_execution = True
364+
except ImportError:
365+
pass # don't fail if polars isn't available
361366

362367
return global_session.with_default_session(
363368
bigframes.session.Session._read_gbq_colab,
Collapse file

‎tests/system/small/test_magics.py‎

Copy file name to clipboard
+100Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
import pandas as pd
15+
import pytest
16+
17+
import bigframes
18+
import bigframes.pandas as bpd
19+
20+
IPython = pytest.importorskip("IPython")
21+
22+
23+
MAGIC_NAME = "bqsql"
24+
25+
26+
@pytest.fixture(scope="module")
27+
def ip():
28+
"""Provides a persistent IPython shell instance for the test session."""
29+
from IPython.testing.globalipapp import get_ipython
30+
31+
shell = get_ipython()
32+
shell.extension_manager.load_extension("bigframes")
33+
return shell
34+
35+
36+
def test_magic_select_lit_to_var(ip):
37+
bigframes.close_session()
38+
39+
line = "dst_var"
40+
cell_body = "SELECT 3"
41+
42+
ip.run_cell_magic(MAGIC_NAME, line, cell_body)
43+
44+
assert "dst_var" in ip.user_ns
45+
result_df = ip.user_ns["dst_var"]
46+
assert result_df.shape == (1, 1)
47+
assert result_df.loc[0, 0] == 3
48+
49+
50+
def test_magic_select_lit_dry_run(ip):
51+
bigframes.close_session()
52+
53+
line = "dst_var --dry_run"
54+
cell_body = "SELECT 3"
55+
56+
ip.run_cell_magic(MAGIC_NAME, line, cell_body)
57+
58+
assert "dst_var" in ip.user_ns
59+
result_df = ip.user_ns["dst_var"]
60+
assert result_df.totalBytesProcessed == 0
61+
62+
63+
def test_magic_select_lit_display(ip):
64+
from IPython.utils.capture import capture_output
65+
66+
bigframes.close_session()
67+
68+
cell_body = "SELECT 3"
69+
70+
with capture_output() as io:
71+
ip.run_cell_magic(MAGIC_NAME, "", cell_body)
72+
assert len(io.outputs) > 0
73+
# Check that the output has data, regardless of the format (html, plain, etc)
74+
available_formats = io.outputs[0].data.keys()
75+
assert len(available_formats) > 0
76+
77+
78+
def test_magic_select_interpolate(ip):
79+
bigframes.close_session()
80+
df = bpd.read_pandas(
81+
pd.DataFrame({"col_a": [1, 2, 3, 4, 5, 6], "col_b": [1, 2, 1, 3, 1, 2]})
82+
)
83+
const_val = 1
84+
85+
ip.push({"df": df, "const_val": const_val})
86+
87+
query = """
88+
SELECT
89+
SUM(col_a) AS total
90+
FROM
91+
{df}
92+
WHERE col_b={const_val}
93+
"""
94+
95+
ip.run_cell_magic(MAGIC_NAME, "dst_var", query)
96+
97+
assert "dst_var" in ip.user_ns
98+
result_df = ip.user_ns["dst_var"]
99+
assert result_df.shape == (1, 1)
100+
assert result_df.loc[0, 0] == 9

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.