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.
14 changes: 14 additions & 0 deletions 14 bigframes/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,20 @@
from bigframes.session import connect, Session
from bigframes.version import __version__

_MAGIC_NAMES = ["bqsql"]


def load_ipython_extension(ipython):
"""Called by IPython when this module is loaded as an IPython extension."""
# Requires IPython to be installed for import to succeed
from bigframes._magics import _cell_magic

for magic_name in _MAGIC_NAMES:
ipython.register_magic_function(
_cell_magic, magic_kind="cell", magic_name=magic_name
)


__all__ = [
"options",
"BigQueryOptions",
Expand Down
56 changes: 56 additions & 0 deletions 56 bigframes/_magics.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
# 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.")
return
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:
with bigframes.option_context(
"display.repr_mode",
"anywidget",
):
display(dataframe)
return
9 changes: 7 additions & 2 deletions 9 bigframes/pandas/io/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 ImportError:
pass # don't fail if polars isn't available

return global_session.with_default_session(
bigframes.session.Session._read_gbq_colab,
Expand Down
100 changes: 100 additions & 0 deletions 100 tests/system/small/test_magics.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
# 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.
import pandas as pd
import pytest

import bigframes
import bigframes.pandas as bpd

IPython = pytest.importorskip("IPython")


MAGIC_NAME = "bqsql"


@pytest.fixture(scope="module")
def ip():
"""Provides a persistent IPython shell instance for the test session."""
from IPython.testing.globalipapp import get_ipython

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):
from IPython.utils.capture import capture_output

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
# Check that the output has data, regardless of the format (html, plain, etc)
available_formats = io.outputs[0].data.keys()
assert len(available_formats) > 0


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
Loading
Morty Proxy This is a proxified and sanitized view of the page, visit original site.