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 69fe317

Browse filesBrowse files
tswasttswast
authored
feat: expose DataFrame.bigquery in both pandas and bigframes DataFrames (#2533)
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 17ecc65 commit 69fe317
Copy full SHA for 69fe317

13 files changed

+395-120Lines changed: 395 additions & 120 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/bigquery/_operations/ai.py‎

Copy file name to clipboardExpand all lines: bigframes/bigquery/_operations/ai.py
+19Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -893,6 +893,25 @@ def forecast(
893893
and might have limited support. For more information, see the launch stage descriptions
894894
(https://cloud.google.com/products#product-launch-stages).
895895
896+
**Examples:**
897+
898+
Forecast using a pandas DataFrame:
899+
900+
>>> import pandas as pd
901+
>>> import bigframes.pandas as bpd
902+
>>> df = pd.DataFrame({"value": [1, 2, 3], "time": pd.to_datetime(["2020-01-01", "2020-01-02", "2020-01-03"])})
903+
>>> bpd.options.display.progress_bar = None # doctest: +SKIP
904+
>>> forecasted_pandas_df = df.bigquery.ai.forecast(data_col="value", timestamp_col="time", horizon=2) # doctest: +SKIP
905+
>>> type(forecasted_pandas_df) # doctest: +SKIP
906+
<class 'pandas.core.frame.DataFrame'>
907+
908+
Forecast using a BigFrames DataFrame:
909+
910+
>>> bf_df = bpd.DataFrame({"value": [1, 2, 3], "time": pd.to_datetime(["2020-01-01", "2020-01-02", "2020-01-03"])})
911+
>>> forecasted_bf_df = bf_df.bigquery.ai.forecast(data_col="value", timestamp_col="time", horizon=2) # doctest: +SKIP
912+
>>> type(forecasted_bf_df) # doctest: +SKIP
913+
<class 'bigframes.dataframe.DataFrame'>
914+
896915
Args:
897916
df (DataFrame):
898917
The dataframe that contains the data that you want to forecast. It could be either a BigFrames Dataframe or
Collapse file

‎bigframes/bigquery/_operations/sql.py‎

Copy file name to clipboardExpand all lines: bigframes/bigquery/_operations/sql.py
+19Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,25 @@ def sql_scalar(
7171
2 4.000000000
7272
dtype: decimal128(38, 9)[pyarrow]
7373
74+
You can also use the `.bigquery` DataFrame accessor to apply a SQL scalar function.
75+
76+
Compute SQL scalar using a pandas DataFrame:
77+
78+
>>> import pandas as pd
79+
>>> df = pd.DataFrame({"x": [1, 2, 3]})
80+
>>> bpd.options.display.progress_bar = None # doctest: +SKIP
81+
>>> pandas_s = df.bigquery.sql_scalar("POW({0}, 2)") # doctest: +SKIP
82+
>>> type(pandas_s) # doctest: +SKIP
83+
<class 'pandas.core.series.Series'>
84+
85+
Compute SQL scalar using a BigFrames DataFrame:
86+
87+
>>> bf_df = bpd.DataFrame({"x": [1, 2, 3]})
88+
>>> bf_s = bf_df.bigquery.sql_scalar("POW({0}, 2)") # doctest: +SKIP
89+
>>> type(bf_s) # doctest: +SKIP
90+
<class 'bigframes.series.Series'>
91+
92+
7493
Args:
7594
sql_template (str):
7695
A SQL format string with Python-style {0} placeholders for each of
Collapse file

‎bigframes/dataframe.py‎

Copy file name to clipboardExpand all lines: bigframes/dataframe.py
+34-7Lines changed: 34 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
overload,
3838
Sequence,
3939
Tuple,
40+
TYPE_CHECKING,
4041
TypeVar,
4142
Union,
4243
)
@@ -47,6 +48,8 @@
4748
import bigframes_vendored.pandas.pandas._typing as vendored_pandas_typing
4849
import google.api_core.exceptions
4950
import google.cloud.bigquery as bigquery
51+
import google.cloud.bigquery.job
52+
import google.cloud.bigquery.table
5053
import numpy
5154
import pandas
5255
from pandas.api import extensions as pd_ext
@@ -91,9 +94,10 @@
9194
import bigframes.session._io.bigquery
9295
import bigframes.session.execution_spec as ex_spec
9396

94-
if typing.TYPE_CHECKING:
97+
if TYPE_CHECKING:
9598
from _typeshed import SupportsRichComparison
9699

100+
import bigframes.extensions.bigframes.dataframe_accessor as bigquery_accessor
97101
import bigframes.session
98102

99103
SingleItemValue = Union[
@@ -144,7 +148,7 @@ def __init__(
144148
):
145149
global bigframes
146150

147-
self._query_job: Optional[bigquery.QueryJob] = None
151+
self._query_job: Optional[google.cloud.bigquery.job.QueryJob] = None
148152

149153
if copy is not None and not copy:
150154
raise ValueError(
@@ -376,6 +380,25 @@ def bqclient(self) -> bigframes.Session:
376380
def _session(self) -> bigframes.Session:
377381
return self._get_block().expr.session
378382

383+
@property
384+
def bigquery(
385+
self,
386+
) -> bigquery_accessor.BigframesBigQueryDataFrameAccessor:
387+
"""
388+
Accessor for BigQuery functionality.
389+
390+
Returns:
391+
bigframes.extensions.core.dataframe_accessor.BigQueryDataFrameAccessor:
392+
Accessor that exposes BigQuery functionality on a DataFrame,
393+
with method names closer to SQL.
394+
"""
395+
# Import the accessor here to avoid circular imports.
396+
import bigframes.extensions.bigframes.dataframe_accessor
397+
398+
return bigframes.extensions.bigframes.dataframe_accessor.BigframesBigQueryDataFrameAccessor(
399+
self
400+
)
401+
379402
@property
380403
def _has_index(self) -> bool:
381404
return len(self._block.index_columns) > 0
@@ -438,7 +461,9 @@ def _should_sql_have_index(self) -> bool:
438461
self.index.name is not None or len(self.index.names) > 1
439462
)
440463

441-
def _to_placeholder_table(self, dry_run: bool = False) -> bigquery.TableReference:
464+
def _to_placeholder_table(
465+
self, dry_run: bool = False
466+
) -> google.cloud.bigquery.table.TableReference:
442467
"""Compiles this DataFrame's expression tree to SQL and saves it to a
443468
(temporary) view or table (in the case of a dry run).
444469
"""
@@ -488,11 +513,11 @@ def sql(self) -> str:
488513
) from e
489514

490515
@property
491-
def query_job(self) -> Optional[bigquery.QueryJob]:
516+
def query_job(self) -> Optional[google.cloud.bigquery.job.QueryJob]:
492517
"""BigQuery job metadata for the most recent query.
493518
494519
Returns:
495-
None or google.cloud.bigquery.QueryJob:
520+
None or google.cloud.bigquery.job.QueryJob:
496521
The most recent `QueryJob
497522
<https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJob>`_.
498523
"""
@@ -606,7 +631,9 @@ def select_dtypes(self, include=None, exclude=None) -> DataFrame:
606631
)
607632
return DataFrame(self._block.select_columns(selected_columns))
608633

609-
def _set_internal_query_job(self, query_job: Optional[bigquery.QueryJob]):
634+
def _set_internal_query_job(
635+
self, query_job: Optional[google.cloud.bigquery.job.QueryJob]
636+
):
610637
self._query_job = query_job
611638

612639
def __getitem__(
@@ -1782,7 +1809,7 @@ def _to_pandas_batches(
17821809
allow_large_results=allow_large_results,
17831810
)
17841811

1785-
def _compute_dry_run(self) -> bigquery.QueryJob:
1812+
def _compute_dry_run(self) -> google.cloud.bigquery.job.QueryJob:
17861813
_, query_job = self._block._compute_dry_run()
17871814
return query_job
17881815

Collapse file

‎bigframes/extensions/bigframes/__init__.py‎

Copy file name to clipboard
+20Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
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 bigframes.extensions.bigframes.dataframe_accessor import (
16+
BigframesAIAccessor,
17+
BigframesBigQueryDataFrameAccessor,
18+
)
19+
20+
__all__ = ["BigframesAIAccessor", "BigframesBigQueryDataFrameAccessor"]
Collapse file

‎bigframes/extensions/bigframes/dataframe_accessor.py‎

Copy file name to clipboard
+71Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
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 __future__ import annotations
16+
17+
from typing import cast, TypeVar
18+
19+
from bigframes.core.logging import log_adapter
20+
import bigframes.dataframe
21+
import bigframes.extensions.core.dataframe_accessor as core_accessor
22+
import bigframes.series
23+
24+
T = TypeVar("T", bound="bigframes.dataframe.DataFrame")
25+
S = TypeVar("S", bound="bigframes.series.Series")
26+
27+
28+
@log_adapter.class_logger
29+
class BigframesAIAccessor(core_accessor.AIAccessor[T, S]):
30+
"""
31+
BigFrames DataFrame accessor for BigQuery AI functions.
32+
"""
33+
34+
def __init__(self, bf_obj: T):
35+
super().__init__(bf_obj)
36+
37+
def _bf_from_dataframe(
38+
self, session: bigframes.session.Session | None
39+
) -> bigframes.dataframe.DataFrame:
40+
return self._obj
41+
42+
def _to_dataframe(self, bf_df: bigframes.dataframe.DataFrame) -> T:
43+
return cast(T, bf_df)
44+
45+
def _to_series(self, bf_series: bigframes.series.Series) -> S:
46+
return cast(S, bf_series)
47+
48+
49+
@log_adapter.class_logger
50+
class BigframesBigQueryDataFrameAccessor(core_accessor.BigQueryDataFrameAccessor[T, S]):
51+
"""
52+
BigFrames DataFrame accessor for BigQuery DataFrames functionality.
53+
"""
54+
55+
def __init__(self, bf_obj: T):
56+
super().__init__(bf_obj)
57+
58+
@property
59+
def ai(self) -> BigframesAIAccessor:
60+
return BigframesAIAccessor(self._obj)
61+
62+
def _bf_from_dataframe(
63+
self, session: bigframes.session.Session | None
64+
) -> bigframes.dataframe.DataFrame:
65+
return self._obj
66+
67+
def _to_dataframe(self, bf_df: bigframes.dataframe.DataFrame) -> T:
68+
return cast(T, bf_df)
69+
70+
def _to_series(self, bf_series: bigframes.series.Series) -> S:
71+
return cast(S, bf_series)
Collapse file

‎bigframes/extensions/core/__init__.py‎

Copy file name to clipboard
+20Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
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 bigframes.extensions.core.dataframe_accessor import (
16+
AIAccessor,
17+
BigQueryDataFrameAccessor,
18+
)
19+
20+
__all__ = ["AIAccessor", "BigQueryDataFrameAccessor"]

0 commit comments

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