fix mores tests

Author: TrevorBergeron

bigframes/functions/_function_client.py

@@ -687,19 +687,17 @@ def get_remote_function_specs(
         self, remote_function_name: str
     ) -> udf_def.RemoteFunctionConfig | None:
         """Check whether a remote function already exists for the udf."""
-        routines = self._bq_client.list_routines(
-            f"{self._gcp_project_id}.{self._bq_dataset}"
-        )
         try:
-            for routine in routines:
-                routine = cast(bigquery.Routine, routine)
-                if routine.reference.routine_id == remote_function_name:
-                    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
+            routine = self._bq_client.get_routine(
+                f"{self._gcp_project_id}.{self._bq_dataset}.{remote_function_name}"
+            )
+            if routine.reference.routine_id == remote_function_name:
+                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

@@ -91,16 +91,21 @@ def _try_import_routine(
 def _try_import_row_routine(
     routine: bigquery.Routine, session: bigframes.Session
 ) -> BigqueryCallableRowRoutine:
-    udf_def = _routine_as_udf_def(routine)
+    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)
 
 
-def _routine_as_udf_def(routine: bigquery.Routine) -> udf_def.BigqueryUdf:
+def _routine_as_udf_def(
+    routine: bigquery.Routine, is_row_processor: bool = False
+) -> udf_def.BigqueryUdf:
     try:
-        return udf_def.BigqueryUdf.from_routine(routine)
+        return udf_def.BigqueryUdf.from_routine(
+            routine, is_row_processor=is_row_processor
+        )
     except udf_def.ReturnTypeMissingError:
         raise bf_formatting.create_exception_with_feedback_link(
             ValueError, "Function return type must be specified."
@@ -140,6 +145,7 @@ def read_gbq_function(
             ValueError, f"Unknown function '{routine_ref}'."
         )
 
+    # TODO(493293086): Deprecate is_row_processor.
     if is_row_processor:
         return _try_import_row_routine(routine, session)
     else:

bigframes/functions/udf_def.py

@@ -255,8 +255,11 @@ def with_devirtualize(self) -> UdfSignature:
             output=self.output.emulating_type,
         )
 
+    # TODO(493293086): Deprecate is_row_processor.
     @classmethod
-    def from_routine(cls, routine: bigquery.Routine) -> UdfSignature:
+    def from_routine(
+        cls, routine: bigquery.Routine, is_row_processor: bool = False
+    ) -> UdfSignature:
         import bigframes.functions._utils
 
         ## Handle return type
@@ -291,7 +294,20 @@ def from_routine(cls, routine: bigquery.Routine) -> UdfSignature:
 
         ## Handle input types
         udf_fields = []
+
+        if is_row_processor:
+            if len(routine.arguments) != 1:
+                raise ValueError(
+                    "Row processor functions must have exactly one input argument."
+                )
+
         for argument in routine.arguments:
+            if is_row_processor:
+                if argument.data_type.type_kind != "STRING":
+                    raise ValueError(
+                        "Row processor functions must have STRING input type."
+                    )
+                udf_fields.append(UdfArg(argument.name, RowSeriesInputFieldV1()))
             udf_fields.append(UdfArg.from_sdk(argument))
 
         return cls(
@@ -361,8 +377,12 @@ def with_devirtualize(self) -> BigqueryUdf:
         )
 
     @classmethod
-    def from_routine(cls, routine: bigquery.Routine) -> BigqueryUdf:
-        signature = UdfSignature.from_routine(routine)
+    def from_routine(
+        cls, routine: bigquery.Routine, is_row_processor: bool = False
+    ) -> BigqueryUdf:
+        signature = UdfSignature.from_routine(
+            routine, is_row_processor=is_row_processor
+        )
         return cls(routine.reference, signature=signature)
 
 

tests/system/large/functions/test_remote_function.py

@@ -1723,7 +1723,7 @@ def square(x):
         gcf = session.cloudfunctionsclient.get_function(
             name=square_remote_2.bigframes_cloud_function
         )
-        assert gcf.service_config.available_cpu == 2.0
+        assert float(gcf.service_config.available_cpu) == 2.0
     finally:
         # clean up the gcp assets created for the remote function
         if square_remote is not None: