simplify udf wrapper guarantees

Author: TrevorBergeron

bigframes/dataframe.py

@@ -2843,7 +2843,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.
@@ -4662,13 +4662,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,
-                bigframes.functions.UdfRoutine,
-            ),
-        ):
+        if not isinstance(func, bigframes.functions.Udf):
             raise TypeError("the first argument must be callable")
 
         if na_action not in {None, "ignore"}:
@@ -4692,13 +4686,7 @@ 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.UdfRoutine,
-                ),
-            ):
+            if not isinstance(func, bigframes.functions.Udf):
                 raise ValueError(
                     "For axis=1 a BigFrames BigQuery function must be used."
                 )
@@ -4832,7 +4820,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 "

bigframes/functions/__init__.py

@@ -11,9 +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, UdfRoutine
+from bigframes.functions.function import Udf
 
 __all__ = [
-    "BigqueryCallableRoutine",
-    "UdfRoutine",
+    "Udf",
 ]

bigframes/functions/_function_session.py

@@ -851,14 +851,19 @@ def wrapper(func):
                 signature=udf_sig,
             )
 
-            if not name:
-                self._update_temp_artifacts(full_rf_name, "")
-
             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 bq_functions.UdfRoutine(func=func, _udf_def=udf_definition)
+            if not name:  # session-managed resource
+                self._update_temp_artifacts(full_rf_name, "")
+                return bq_functions.UdfRoutine(func=func, _udf_def=udf_definition)
+
+            # user-managed permanent resource
+            else:
+                return bq_functions.BigqueryCallableRoutine(
+                    udf_definition, session, local_func=func, is_managed=True
+                )
 
         return wrapper
 

bigframes/functions/function.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Callable, Optional, TYPE_CHECKING
+from typing import Callable, Optional, Protocol, runtime_checkable, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from bigframes.session import Session
@@ -31,6 +31,9 @@
 from bigframes.functions import _function_session as bff_session
 from bigframes.functions import function_typing, udf_def
 
+if TYPE_CHECKING:
+    import bigframes.core.col
+
 logger = logging.getLogger(__name__)
 
 
@@ -154,6 +157,13 @@ def read_gbq_function(
         return _try_import_routine(routine, session)
 
 
+@runtime_checkable
+class Udf(Protocol):
+    @property
+    def udf_def(self) -> udf_def.BigqueryUdf:
+        ...
+
+
 class BigqueryCallableRoutine:
     """
     A reference to a routine in the context of a session.

bigframes/series.py

@@ -1558,7 +1558,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)
             # When it's a plain Python function.
             else:
@@ -2031,10 +2031,7 @@ def apply(
 
         if isinstance(
             func,
-            (
-                bigframes.functions.BigqueryCallableRoutine,
-                bigframes.functions.UdfRoutine,
-            ),
+            (bigframes.functions.Udf,),
         ):
             # We are working with bigquery function at this point
             if args:
@@ -2098,10 +2095,7 @@ def combine(
 
         if isinstance(
             func,
-            (
-                bigframes.functions.BigqueryCallableRoutine,
-                bigframes.functions.UdfRoutine,
-            ),
+            bigframes.functions.Udf,
         ):
             result_series = self._apply_binary_op(
                 other, ops.BinaryRemoteFunctionOp(function_def=func.udf_def)