feat: expose DataFrame.bigquery in both pandas and bigframes DataFrames

Author: tswast

bigframes/dataframe.py

@@ -37,6 +37,7 @@
     overload,
     Sequence,
     Tuple,
+    TYPE_CHECKING,
     TypeVar,
     Union,
 )
@@ -47,6 +48,8 @@
 import bigframes_vendored.pandas.pandas._typing as vendored_pandas_typing
 import google.api_core.exceptions
 import google.cloud.bigquery as bigquery
+import google.cloud.bigquery.job
+import google.cloud.bigquery.table
 import numpy
 import pandas
 from pandas.api import extensions as pd_ext
@@ -91,9 +94,10 @@
 import bigframes.session._io.bigquery
 import bigframes.session.execution_spec as ex_spec
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     from _typeshed import SupportsRichComparison
 
+    import bigframes.extensions.bigframes.dataframe_accessor as bigquery_accessor
     import bigframes.session
 
     SingleItemValue = Union[
@@ -144,7 +148,7 @@ def __init__(
     ):
         global bigframes
 
-        self._query_job: Optional[bigquery.QueryJob] = None
+        self._query_job: Optional[google.cloud.bigquery.job.QueryJob] = None
 
         if copy is not None and not copy:
             raise ValueError(
@@ -376,6 +380,25 @@ def bqclient(self) -> bigframes.Session:
     def _session(self) -> bigframes.Session:
         return self._get_block().expr.session
 
+    @property
+    def bigquery(
+        self,
+    ) -> bigquery_accessor.BigframesBigQueryDataFrameAccessor:
+        """
+        Accessor for BigQuery functionality.
+
+        Returns:
+            bigframes.extensions.core.dataframe_accessor.BigQueryDataFrameAccessor:
+                Accessor that exposes BigQuery functionality on a DataFrame,
+                with method names closer to SQL.
+        """
+        # Import the accessor here to avoid circular imports.
+        import bigframes.extensions.bigframes.dataframe_accessor
+
+        return bigframes.extensions.bigframes.dataframe_accessor.BigframesBigQueryDataFrameAccessor(
+            self
+        )
+
     @property
     def _has_index(self) -> bool:
         return len(self._block.index_columns) > 0
@@ -438,7 +461,9 @@ def _should_sql_have_index(self) -> bool:
             self.index.name is not None or len(self.index.names) > 1
         )
 
-    def _to_placeholder_table(self, dry_run: bool = False) -> bigquery.TableReference:
+    def _to_placeholder_table(
+        self, dry_run: bool = False
+    ) -> google.cloud.bigquery.table.TableReference:
         """Compiles this DataFrame's expression tree to SQL and saves it to a
         (temporary) view or table (in the case of a dry run).
         """
@@ -488,11 +513,11 @@ def sql(self) -> str:
             ) from e
 
     @property
-    def query_job(self) -> Optional[bigquery.QueryJob]:
+    def query_job(self) -> Optional[google.cloud.bigquery.job.QueryJob]:
         """BigQuery job metadata for the most recent query.
 
         Returns:
-            None or google.cloud.bigquery.QueryJob:
+            None or google.cloud.bigquery.job.QueryJob:
                 The most recent `QueryJob
                 <https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJob>`_.
         """
@@ -606,7 +631,9 @@ def select_dtypes(self, include=None, exclude=None) -> DataFrame:
         )
         return DataFrame(self._block.select_columns(selected_columns))
 
-    def _set_internal_query_job(self, query_job: Optional[bigquery.QueryJob]):
+    def _set_internal_query_job(
+        self, query_job: Optional[google.cloud.bigquery.job.QueryJob]
+    ):
         self._query_job = query_job
 
     def __getitem__(
@@ -1782,7 +1809,7 @@ def _to_pandas_batches(
             allow_large_results=allow_large_results,
         )
 
-    def _compute_dry_run(self) -> bigquery.QueryJob:
+    def _compute_dry_run(self) -> google.cloud.bigquery.job.QueryJob:
         _, query_job = self._block._compute_dry_run()
         return query_job
 

bigframes/extensions/bigframes/__init__.py

@@ -0,0 +1,20 @@
+# 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 bigframes.extensions.bigframes.dataframe_accessor import (
+    BigframesAIAccessor,
+    BigframesBigQueryDataFrameAccessor,
+)
+
+__all__ = ["BigframesAIAccessor", "BigframesBigQueryDataFrameAccessor"]

bigframes/extensions/bigframes/dataframe_accessor.py

@@ -0,0 +1,71 @@
+# 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 __future__ import annotations
+
+from typing import cast, TypeVar
+
+from bigframes.core.logging import log_adapter
+import bigframes.dataframe
+import bigframes.extensions.core.dataframe_accessor as core_accessor
+import bigframes.series
+
+T = TypeVar("T", bound="bigframes.dataframe.DataFrame")
+S = TypeVar("S", bound="bigframes.series.Series")
+
+
+@log_adapter.class_logger
+class BigframesAIAccessor(core_accessor.AIAccessor[T, S]):
+    """
+    BigFrames DataFrame accessor for BigQuery AI functions.
+    """
+
+    def __init__(self, bf_obj: T):
+        super().__init__(bf_obj)
+
+    def _bf_from_dataframe(
+        self, session: bigframes.session.Session | None
+    ) -> bigframes.dataframe.DataFrame:
+        return self._obj
+
+    def _to_dataframe(self, bf_df: bigframes.dataframe.DataFrame) -> T:
+        return cast(T, bf_df)
+
+    def _to_series(self, bf_series: bigframes.series.Series) -> S:
+        return cast(S, bf_series)
+
+
+@log_adapter.class_logger
+class BigframesBigQueryDataFrameAccessor(core_accessor.BigQueryDataFrameAccessor[T, S]):
+    """
+    BigFrames DataFrame accessor for BigQuery DataFrames functionality.
+    """
+
+    def __init__(self, bf_obj: T):
+        super().__init__(bf_obj)
+
+    @property
+    def ai(self) -> BigframesAIAccessor:
+        return BigframesAIAccessor(self._obj)
+
+    def _bf_from_dataframe(
+        self, session: bigframes.session.Session | None
+    ) -> bigframes.dataframe.DataFrame:
+        return self._obj
+
+    def _to_dataframe(self, bf_df: bigframes.dataframe.DataFrame) -> T:
+        return cast(T, bf_df)
+
+    def _to_series(self, bf_series: bigframes.series.Series) -> S:
+        return cast(S, bf_series)

bigframes/extensions/core/__init__.py

@@ -0,0 +1,20 @@
+# 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 bigframes.extensions.core.dataframe_accessor import (
+    AIAccessor,
+    BigQueryDataFrameAccessor,
+)
+
+__all__ = ["AIAccessor", "BigQueryDataFrameAccessor"]