fix some more stuff

Author: TrevorBergeron

bigframes/functions/_function_client.py

@@ -694,7 +694,12 @@ def get_remote_function_specs(
             for routine in routines:
                 routine = cast(bigquery.Routine, routine)
                 if routine.reference.routine_id == remote_function_name:
-                    return udf_def.RemoteFunctionConfig.from_bq_routine(routine)
+                    try:
+                        return udf_def.RemoteFunctionConfig.from_bq_routine(routine)
+                    except udf_def.ReturnTypeMissingError:
+                        # The remote function exists, but it's missing a return type.
+                        # Something is wrong with the function, so we should replace it.
+                        return None
         except google.api_core.exceptions.NotFound:
             # The dataset might not exist, in which case the remote function doesn't, either.
             # Note: list_routines doesn't make an API request until we iterate on the response object.

bigframes/functions/function.py

@@ -160,14 +160,10 @@ def __init__(
         *,
         local_func: Optional[Callable] = None,
         cloud_function_ref: Optional[str] = None,
-        post_routine: Optional[
-            Callable[[bigframes.series.Series], bigframes.series.Series]
-        ] = None,
         is_managed: bool = False,
     ):
         self._udf_def = udf_def
         self._session = session
-        self._post_routine = post_routine
         self._local_fun = local_func
         self._cloud_function = cloud_function_ref
         self._is_managed = is_managed

bigframes/functions/function_template.py

@@ -28,8 +28,8 @@
 
 
 # Placeholder variables for testing.
-input_sql_types = ("STRING",)
-output_sql_type = "STRING"
+input_types = ("STRING",)
+output_type = "STRING"
 
 
 # Convert inputs to BigQuery JSON. See:
@@ -152,7 +152,7 @@ def udf(*args):
 # }
 # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#input_format
 def udf_http(request):
-    global input_sql_types, output_sql_type
+    global input_types, output_type
     import json
     import traceback
 
@@ -164,7 +164,7 @@ def udf_http(request):
         replies = []
         for call in calls:
             reply = convert_to_bq_json(
-                output_sql_type, udf(*convert_call(input_sql_types, call))
+                output_type, udf(*convert_call(input_types, call))
             )
             if type(reply) is list:
                 # Since the BQ remote function does not support array yet,
@@ -178,7 +178,7 @@ def udf_http(request):
 
 
 def udf_http_row_processor(request):
-    global output_sql_type
+    global output_type
     import json
     import math
     import traceback
@@ -192,7 +192,7 @@ def udf_http_row_processor(request):
         replies = []
         for call in calls:
             reply = convert_to_bq_json(
-                output_sql_type, udf(get_pd_series(call[0]), *call[1:])
+                output_type, udf(get_pd_series(call[0]), *call[1:])
             )
             if type(reply) is list:
                 # Since the BQ remote function does not support array yet,
@@ -254,6 +254,9 @@ def generate_cloud_function_main_code(
     # Pickle the udf with all its dependencies
     udf_code_file, udf_pickle_file = generate_udf_code(code_def, directory)
 
+    input_types = tuple(arg.sql_type for arg in udf_signature.inputs)
+    output_type = udf_signature.output.sql_type
+
     code_blocks = [
         f"""\
 import cloudpickle
@@ -263,8 +266,8 @@ def generate_cloud_function_main_code(
 with open("{udf_pickle_file}", "rb") as f:
     udf = cloudpickle.load(f)
 
-input_types = {repr(input_sql_types)}
-output_type = {repr(output_sql_type)}
+input_types = {repr(input_types)}
+output_type = {repr(output_type)}
 """
     ]
 

bigframes/functions/udf_def.py

@@ -106,9 +106,9 @@ def bf_type(self) -> bigframes.dtypes.Dtype:
 
     @property
     def sql_type(self) -> str:
-        type_kind = function_typing.sdk_type_from_python_type(self._py_type)
-        assert type_kind is not None
-        return type_kind.name
+        sdk_type = function_typing.sdk_type_from_python_type(self._py_type)
+        assert sdk_type.type_kind is not None
+        return sdk_type.type_kind.name
 
     def stable_hash(self) -> bytes:
         hash_val = hashlib.md5()
@@ -233,7 +233,9 @@ def from_routine(cls, routine: bigquery.Routine) -> UdfSignature:
 
         ## Handle return type
         if routine.return_type is None:
-            raise ReturnTypeMissingError
+            raise ReturnTypeMissingError(
+                f"Routine {routine} has no return type. Routine properties: {routine._properties}"
+            )
 
         bq_return_type = cast(bigquery.StandardSqlDataType, routine.return_type)