Skip to content

Commit 69fe317

Browse files
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

File tree

13 files changed

+395
-120
lines changed

13 files changed

+395
-120
lines changed

bigframes/bigquery/_operations/ai.py

Lines 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

bigframes/bigquery/_operations/sql.py

Lines 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

bigframes/dataframe.py

Lines 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

Lines 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"]
Lines 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)
Lines 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)