feat(bigframes): add AI TVFs to the pandas bq accessor (googleapis#17402)

Also updated the tests to fully utilize the mocking framework.

internal issue: b/517233441

Author: sycai

packages/bigframes/bigframes/extensions/bigframes/series_accessor.py

@@ -20,41 +20,68 @@
 
 from typing import Optional, TypeVar, cast
 
-import bigframes.extensions.core.series_accessor as core_accessor
-import bigframes.series
-import bigframes.session
+from bigframes import dataframe, series, session
 from bigframes.core.logging import log_adapter
+from bigframes.extensions.core import series_accessor as core_accessor
 
-S = TypeVar("S", bound="bigframes.series.Series")
+T = TypeVar("T", bound="dataframe.DataFrame")
+S = TypeVar("S", bound="series.Series")
 
 
 @log_adapter.class_logger
-class BigframesBigQuerySeriesAccessor(core_accessor.BigQuerySeriesAccessor[S]):
+class BigframesBigQuerySeriesAccessor(core_accessor.BigQuerySeriesAccessor[T, S]):
     def __init__(self, bf_obj: S):
         super().__init__(bf_obj)
 
     def _bf_from_series(
-        self, session: Optional[bigframes.session.Session] = None
-    ) -> bigframes.series.Series:
+        self, session: Optional[session.Session] = None
+    ) -> series.Series:
         return self._obj
 
-    def _to_series(self, bf_series: bigframes.series.Series) -> S:
+    def _to_dataframe(self, bf_df: dataframe.DataFrame) -> T:
+        return cast(T, bf_df)
+
+    def _to_series(self, bf_series: series.Series) -> S:
         return cast(S, bf_series)
 
     @property
-    def aead(self) -> BigframesAeadSeriesAccessor[S]:
+    def ai(self) -> BigframesAiSeriesAccessor[T, S]:
+        return BigframesAiSeriesAccessor(self._obj)
+
+    @property
+    def aead(self) -> BigframesAeadSeriesAccessor[T, S]:
         return BigframesAeadSeriesAccessor(self._obj)
 
 
 @log_adapter.class_logger
-class BigframesAeadSeriesAccessor(core_accessor.AeadSeriesAccessor[S]):
+class BigframesAiSeriesAccessor(core_accessor.AiSeriesAccessor[T, S]):
     def __init__(self, bf_obj: S):
         super().__init__(bf_obj)
 
     def _bf_from_series(
-        self, session: Optional[bigframes.session.Session] = None
-    ) -> bigframes.series.Series:
+        self, session: Optional[session.Session] = None
+    ) -> series.Series:
         return self._obj
 
-    def _to_series(self, bf_series: bigframes.series.Series) -> S:
+    def _to_dataframe(self, bf_df: dataframe.DataFrame) -> T:
+        return cast(T, bf_df)
+
+    def _to_series(self, bf_series: series.Series) -> S:
+        return cast(S, bf_series)
+
+
+@log_adapter.class_logger
+class BigframesAeadSeriesAccessor(core_accessor.AeadSeriesAccessor[T, S]):
+    def __init__(self, bf_obj: S):
+        super().__init__(bf_obj)
+
+    def _bf_from_series(
+        self, session: Optional[session.Session] = None
+    ) -> series.Series:
+        return self._obj
+
+    def _to_dataframe(self, bf_df: dataframe.DataFrame) -> T:
+        return cast(T, bf_df)
+
+    def _to_series(self, bf_series: series.Series) -> S:
         return cast(S, bf_series)

packages/bigframes/bigframes/extensions/core/abstract_series_accessor.py

@@ -0,0 +1,50 @@
+# 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.
+#
+# DO NOT MODIFY THIS FILE DIRECTLY.
+# This file was generated by the script: scripts/generate_bigframes_bigquery.py
+#
+
+from __future__ import annotations
+
+import abc
+from typing import (
+    Generic,
+    Optional,
+    TypeVar,
+)
+
+from bigframes import dataframe, series, session
+
+T = TypeVar("T")
+S = TypeVar("S")
+
+
+class AbstractBigQuerySeriesAccessor(abc.ABC, Generic[T, S]):
+    def __init__(self, obj: S):
+        self._obj = obj
+
+    @abc.abstractmethod
+    def _bf_from_series(
+        self, session: Optional[session.Session] = None
+    ) -> series.Series:
+        """Convert the accessor's object to a BigFrames Series."""
+
+    @abc.abstractmethod
+    def _to_dataframe(self, bf_df: dataframe.DataFrame) -> T:
+        """Convert a BigFrames DataFrame to the accessor's object type."""
+
+    @abc.abstractmethod
+    def _to_series(self, bf_series: series.Series) -> S:
+        """Convert a BigFrames Series to the accessor's object type."""