refactor(bigframes): Simplify @udf wrapper object (#16556)

TrevorBergeron · web-flow · commit 471eb13a1feb · 2026-05-08T16:25:31.000-07:00
diff --git a/packages/bigframes/bigframes/dataframe.py b/packages/bigframes/bigframes/dataframe.py
@@ -2851,7 +2851,7 @@ def _apply_callable(self, condition):
         """Executes the possible callable condition as needed."""
         if callable(condition):
             # When it's a bigframes function.
-            if hasattr(condition, "bigframes_bigquery_function"):
+            if isinstance(condition, bigframes.functions.Udf):
                 return self.apply(condition, axis=1)
 
             # When it's a plain Python function.
@@ -4685,7 +4685,7 @@ def _prepare_export(
         return array_value, id_overrides
 
     def map(self, func, na_action: Optional[str] = None) -> DataFrame:
-        if not isinstance(func, bigframes.functions.BigqueryCallableRoutine):
+        if not isinstance(func, bigframes.functions.Udf):
             raise TypeError("the first argument must be callable")
 
         if na_action not in {None, "ignore"}:
@@ -4709,18 +4709,12 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
             )
             warnings.warn(msg, category=bfe.FunctionAxisOnePreviewWarning)
 
-            if not isinstance(
-                func,
-                (
-                    bigframes.functions.BigqueryCallableRoutine,
-                    bigframes.functions.BigqueryCallableRowRoutine,
-                ),
-            ):
+            if not isinstance(func, bigframes.functions.Udf):
                 raise ValueError(
                     "For axis=1 a BigFrames BigQuery function must be used."
                 )
 
-            if func.is_row_processor:
+            if func.udf_def.signature.is_row_processor:
                 # Early check whether the dataframe dtypes are currently supported
                 # in the bigquery function
                 # NOTE: Keep in sync with the value converters used in the gcf code
@@ -4849,7 +4843,7 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
 
         # At this point column-wise or element-wise bigquery function operation will
         # be performed (not supported).
-        if hasattr(func, "bigframes_bigquery_function"):
+        if isinstance(func, bigframes.functions.Udf):
             raise formatter.create_exception_with_feedback_link(
                 NotImplementedError,
                 "BigFrames DataFrame '.apply()' does not support BigFrames "
diff --git a/packages/bigframes/bigframes/functions/__init__.py b/packages/bigframes/bigframes/functions/__init__.py
@@ -11,12 +11,8 @@
 # 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.functions.function import (
-    BigqueryCallableRoutine,
-    BigqueryCallableRowRoutine,
-)
+from bigframes.functions.function import Udf
 
 __all__ = [
-    "BigqueryCallableRoutine",
-    "BigqueryCallableRowRoutine",
+    "Udf",
 ]
diff --git a/packages/bigframes/bigframes/functions/_function_session.py b/packages/bigframes/bigframes/functions/_function_session.py
@@ -20,18 +20,18 @@
 import inspect
 import sys
 import threading
-import warnings
 from typing import (
-    TYPE_CHECKING,
     Any,
+    cast,
     Dict,
     Literal,
     Mapping,
     Optional,
     Sequence,
+    TYPE_CHECKING,
     Union,
-    cast,
 )
+import warnings
 
 import google.api_core.exceptions
 from google.cloud import (
@@ -41,9 +41,9 @@
     resourcemanager_v3,
 )
 
+from bigframes import clients
 import bigframes.exceptions as bfe
 import bigframes.formatting_helpers as bf_formatting
-from bigframes import clients
 from bigframes.functions import function as bq_functions
 from bigframes.functions import udf_def
 
@@ -630,25 +630,15 @@ def wrapper(func):
             if udf_sig.is_row_processor:
                 msg = bfe.format_message("input_types=Series is in preview.")
                 warnings.warn(msg, stacklevel=1, category=bfe.PreviewWarning)
-                return decorator(
-                    bq_functions.BigqueryCallableRowRoutine(
-                        udf_definition,
-                        session,
-                        cloud_function_ref=bigframes_cloud_function,
-                        local_func=func,
-                        is_managed=False,
-                    )
-                )
-            else:
-                return decorator(
-                    bq_functions.BigqueryCallableRoutine(
-                        udf_definition,
-                        session,
-                        cloud_function_ref=bigframes_cloud_function,
-                        local_func=func,
-                        is_managed=False,
-                    )
+            return decorator(
+                bq_functions.BigqueryCallableRoutine(
+                    udf_definition,
+                    session,
+                    cloud_function_ref=bigframes_cloud_function,
+                    local_func=func,
+                    is_managed=False,
                 )
+            )
 
         return wrapper
 
@@ -834,8 +824,9 @@ def wrapper(func):
                 bq_connection_manager,
                 session=session,  # type: ignore
             )
+            code_def = udf_def.CodeDef.from_func(func)
             config = udf_def.ManagedFunctionConfig(
-                code=udf_def.CodeDef.from_func(func),
+                code=code_def,
                 signature=udf_sig,
                 max_batching_rows=max_batching_rows,
                 container_cpu=container_cpu,
@@ -859,28 +850,18 @@ def wrapper(func):
                 signature=udf_sig,
             )
 
-            if not name:
-                self._update_temp_artifacts(full_rf_name, "")
-
-            decorator = functools.wraps(func)
             if udf_sig.is_row_processor:
                 msg = bfe.format_message("input_types=Series is in preview.")
                 warnings.warn(msg, stacklevel=1, category=bfe.PreviewWarning)
-                assert session is not None  # appease mypy
-                return decorator(
-                    bq_functions.BigqueryCallableRowRoutine(
-                        udf_definition, session, local_func=func, is_managed=True
-                    )
-                )
+
+            if not name:  # session-owned resource - will be cleaned up automatically
+                self._update_temp_artifacts(full_rf_name, "")
+                return bq_functions.UdfRoutine(func=func, _udf_def=udf_definition)
+
+            # user-managed permanent resource - will not be cleaned up automatically
             else:
-                assert session is not None  # appease mypy
-                return decorator(
-                    bq_functions.BigqueryCallableRoutine(
-                        udf_definition,
-                        session,
-                        local_func=func,
-                        is_managed=True,
-                    )
+                return bq_functions.BigqueryCallableRoutine(
+                    udf_definition, session, local_func=func, is_managed=True
                 )
 
         return wrapper
diff --git a/packages/bigframes/bigframes/functions/function.py b/packages/bigframes/bigframes/functions/function.py
@@ -15,11 +15,9 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Callable, Optional
+from typing import Callable, Optional, Protocol, runtime_checkable, TYPE_CHECKING
 
-if TYPE_CHECKING:
-    import bigframes.series
-    from bigframes.session import Session
+import dataclasses
 
 import google.api_core.exceptions
 from google.cloud import bigquery
@@ -28,6 +26,11 @@
 from bigframes.functions import _function_session as bff_session
 from bigframes.functions import function_typing, udf_def
 
+if TYPE_CHECKING:
+    import bigframes.core.col
+    from bigframes.session import Session
+    import bigframes.series
+
 logger = logging.getLogger(__name__)
 
 
@@ -90,13 +93,13 @@ def _try_import_routine(
 
 def _try_import_row_routine(
     routine: bigquery.Routine, session: bigframes.Session
-) -> BigqueryCallableRowRoutine:
+) -> BigqueryCallableRoutine:
     udf_def = _routine_as_udf_def(routine, is_row_processor=True)
 
     is_remote = (
         hasattr(routine, "remote_function_options") and routine.remote_function_options
     )
-    return BigqueryCallableRowRoutine(udf_def, session, is_managed=not is_remote)
+    return BigqueryCallableRoutine(udf_def, session, is_managed=not is_remote)
 
 
 def _routine_as_udf_def(
@@ -117,7 +120,6 @@ def _routine_as_udf_def(
         )
 
 
-# TODO(b/399894805): Support managed function.
 def read_gbq_function(
     function_name: str,
     *,
@@ -152,6 +154,18 @@ def read_gbq_function(
         return _try_import_routine(routine, session)
 
 
+@runtime_checkable
+class Udf(Protocol):
+    """
+    Protocol for all BigFrames user-defined functions.
+
+    Has @runtime_checkable so functions like df.apply() can dispatch UDFs with isinstance() checks.
+    """
+
+    @property
+    def udf_def(self) -> udf_def.BigqueryUdf: ...
+
+
 class BigqueryCallableRoutine:
     """
     A reference to a routine in the context of a session.
@@ -178,8 +192,8 @@ def __call__(self, *args, **kwargs):
         if self._local_fun:
             return self._local_fun(*args, **kwargs)
         # avoid circular imports
-        import bigframes.session._io.bigquery as bf_io_bigquery
         from bigframes.core.compile.sqlglot import sql as sg_sql
+        import bigframes.session._io.bigquery as bf_io_bigquery
 
         args_string = ", ".join([sg_sql.to_sql(sg_sql.literal(v)) for v in args])
         sql = f"SELECT `{str(self._udf_def.routine_ref)}`({args_string})"
@@ -202,7 +216,7 @@ def bigframes_remote_function(self):
 
     @property
     def is_row_processor(self) -> bool:
-        return False
+        return self.udf_def.signature.is_row_processor
 
     @property
     def udf_def(self) -> udf_def.BigqueryUdf:
@@ -225,75 +239,16 @@ def bigframes_bigquery_function_output_dtype(self):
         return self.udf_def.signature.output.emulating_type.bf_type
 
 
-class BigqueryCallableRowRoutine:
-    """
-    A reference to a routine in the context of a session.
-
-    Can be used both directly as a callable, or as an input to dataframe ops that take a callable.
-    """
-
-    def __init__(
-        self,
-        udf_def: udf_def.BigqueryUdf,
-        session: bigframes.Session,
-        *,
-        local_func: Optional[Callable] = None,
-        cloud_function_ref: Optional[str] = None,
-        is_managed: bool = False,
-    ):
-        assert udf_def.signature.is_row_processor
-        self._udf_def = udf_def
-        self._session = session
-        self._local_fun = local_func
-        self._cloud_function = cloud_function_ref
-        self._is_managed = is_managed
+@dataclasses.dataclass(frozen=True)
+class UdfRoutine:
+    func: Callable
+    # Try not to depend on this, bq managed function creation will be deferred later
+    # And this ref will be replaced with requirements rather to support lazy creation
+    _udf_def: udf_def.BigqueryUdf
 
     def __call__(self, *args, **kwargs):
-        if self._local_fun:
-            return self._local_fun(*args, **kwargs)
-        # avoid circular imports
-        import bigframes.session._io.bigquery as bf_io_bigquery
-        from bigframes.core.compile.sqlglot import sql as sg_sql
-
-        args_string = ", ".join([sg_sql.to_sql(sg_sql.literal(v)) for v in args])
-        sql = f"SELECT `{str(self._udf_def.routine_ref)}`({args_string})"
-        iter, job = bf_io_bigquery.start_query_with_client(
-            self._session.bqclient,
-            sql=sql,
-            query_with_job=True,
-            job_config=bigquery.QueryJobConfig(),
-            publisher=self._session._publisher,
-        )  # type: ignore
-        return list(iter.to_arrow().to_pydict().values())[0][0]
-
-    @property
-    def bigframes_bigquery_function(self) -> str:
-        return str(self._udf_def.routine_ref)
-
-    @property
-    def bigframes_remote_function(self):
-        return None if self._is_managed else str(self._udf_def.routine_ref)
-
-    @property
-    def is_row_processor(self) -> bool:
-        return True
+        return self.func(*args, **kwargs)
 
     @property
     def udf_def(self) -> udf_def.BigqueryUdf:
         return self._udf_def
-
-    @property
-    def bigframes_cloud_function(self) -> Optional[str]:
-        return self._cloud_function
-
-    @property
-    def input_dtypes(self):
-        return tuple(arg.bf_type for arg in self.udf_def.signature.inputs)
-
-    @property
-    def output_dtype(self):
-        return self.udf_def.signature.output.bf_type
-
-    @property
-    def bigframes_bigquery_function_output_dtype(self):
-        return self.udf_def.signature.output.emulating_type.bf_type
diff --git a/packages/bigframes/bigframes/functions/function_template.py b/packages/bigframes/bigframes/functions/function_template.py
@@ -43,7 +43,7 @@ def convert_from_bq_json(type_, arg):
     import base64
     import collections
 
-    converters = collections.defaultdict(lambda: (lambda value: value))  # type: ignore
+    converters = collections.defaultdict(lambda: lambda value: value)  # type: ignore
     converters["BYTES"] = base64.b64decode
     converter = converters[type_]
     return converter(arg) if arg is not None else None
@@ -53,7 +53,7 @@ def convert_to_bq_json(type_, arg):
     import base64
     import collections
 
-    converters = collections.defaultdict(lambda: (lambda value: value))  # type: ignore
+    converters = collections.defaultdict(lambda: lambda value: value)  # type: ignore
     converters["BYTES"] = lambda value: base64.b64encode(value).decode("utf-8")
     converter = converters[type_]
     return converter(arg) if arg is not None else None
diff --git a/packages/bigframes/bigframes/functions/udf_def.py b/packages/bigframes/bigframes/functions/udf_def.py
@@ -457,6 +457,14 @@ def stable_hash(self) -> bytes:
 
         return hash_val.digest()
 
+    def to_callable(self):
+        """
+        Reconstructs the python callable from the pickled code.
+
+        Assumption: package_requirements match local environment
+        """
+        return cloudpickle.loads(self.pickled_code)
+
 
 @dataclasses.dataclass(frozen=True)
 class ManagedFunctionConfig:
diff --git a/packages/bigframes/bigframes/series.py b/packages/bigframes/bigframes/series.py
diff --git a/packages/bigframes/tests/system/large/functions/test_managed_function.py b/packages/bigframes/tests/system/large/functions/test_managed_function.py
diff --git a/packages/bigframes/tests/unit/functions/test_remote_function.py b/packages/bigframes/tests/unit/functions/test_remote_function.py
