fixes

Author: TrevorBergeron

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1037,6 +1037,7 @@ def timedelta_floor_op_impl(x: ibis_types.NumericValue):
 @scalar_op_compiler.register_unary_op(ops.RemoteFunctionOp, pass_op=True)
 def remote_function_op_impl(x: ibis_types.Value, op: ops.RemoteFunctionOp):
     udf_sig = op.function_def.signature
+    assert not udf_sig.is_virtual  # should have been devirtualized in lowering pass
     ibis_py_sig = (tuple(arg.py_type for arg in udf_sig.inputs), udf_sig.output.py_type)
 
     @ibis_udf.scalar.builtin(
@@ -1056,6 +1057,7 @@ def binary_remote_function_op_impl(
     x: ibis_types.Value, y: ibis_types.Value, op: ops.BinaryRemoteFunctionOp
 ):
     udf_sig = op.function_def.signature
+    assert not udf_sig.is_virtual  # should have been devirtualized in lowering pass
     ibis_py_sig = (tuple(arg.py_type for arg in udf_sig.inputs), udf_sig.output.py_type)
 
     @ibis_udf.scalar.builtin(
@@ -1073,6 +1075,7 @@ def nary_remote_function_op_impl(
     *operands: ibis_types.Value, op: ops.NaryRemoteFunctionOp
 ):
     udf_sig = op.function_def.signature
+    assert not udf_sig.is_virtual  # should have been devirtualized in lowering pass
     ibis_py_sig = (tuple(arg.py_type for arg in udf_sig.inputs), udf_sig.output.py_type)
     arg_names = tuple(arg.name for arg in udf_sig.inputs)
 

bigframes/core/rewrite/udfs.py

@@ -30,14 +30,14 @@ def op(self) -> type[ops.ScalarOp]:
     def lower(self, expr: expression.OpExpression) -> expression.Expression:
         assert isinstance(expr.op, ops.RemoteFunctionOp)
         func_def = expr.op.function_def
-        if isinstance(func_def.signature.output, udf_def.DirectScalarType):
-            return expr.op.as_expr(*expr.children)
-        assert isinstance(func_def.signature.output, udf_def.VirtualListTypeV1)
         devirtualized_expr = ops.RemoteFunctionOp(
             func_def.with_devirtualize(),
             apply_on_null=expr.op.apply_on_null,
         ).as_expr(*expr.children)
-        return func_def.signature.output.out_expr(devirtualized_expr)
+        if isinstance(func_def.signature.output, udf_def.VirtualListTypeV1):
+            return func_def.signature.output.out_expr(devirtualized_expr)
+        else:
+            return devirtualized_expr
 
 
 @dataclasses.dataclass
@@ -49,14 +49,13 @@ def op(self) -> type[ops.ScalarOp]:
     def lower(self, expr: expression.OpExpression) -> expression.Expression:
         assert isinstance(expr.op, ops.BinaryRemoteFunctionOp)
         func_def = expr.op.function_def
-
-        if isinstance(func_def.signature.output, udf_def.DirectScalarType):
-            return expr.op.as_expr(*expr.children)
-        assert isinstance(func_def.signature.output, udf_def.VirtualListTypeV1)
         devirtualized_expr = ops.BinaryRemoteFunctionOp(
             func_def.with_devirtualize(),
         ).as_expr(*expr.children)
-        return func_def.signature.output.out_expr(devirtualized_expr)
+        if isinstance(func_def.signature.output, udf_def.VirtualListTypeV1):
+            return func_def.signature.output.out_expr(devirtualized_expr)
+        else:
+            return devirtualized_expr
 
 
 @dataclasses.dataclass
@@ -68,13 +67,13 @@ def op(self) -> type[ops.ScalarOp]:
     def lower(self, expr: expression.OpExpression) -> expression.Expression:
         assert isinstance(expr.op, ops.NaryRemoteFunctionOp)
         func_def = expr.op.function_def
-        if isinstance(func_def.signature.output, udf_def.DirectScalarType):
-            return expr.op.as_expr(*expr.children)
-        assert isinstance(func_def.signature.output, udf_def.VirtualListTypeV1)
         devirtualized_expr = ops.NaryRemoteFunctionOp(
             func_def.with_devirtualize(),
         ).as_expr(*expr.children)
-        return func_def.signature.output.out_expr(devirtualized_expr)
+        if isinstance(func_def.signature.output, udf_def.VirtualListTypeV1):
+            return func_def.signature.output.out_expr(devirtualized_expr)
+        else:
+            return devirtualized_expr
 
 
 UDF_LOWERING_RULES = (

bigframes/functions/_function_client.py

@@ -177,7 +177,6 @@ def create_bq_remote_function(
 
         # Create BQ function
         # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_remote_function_2
-        bq_function_return_type = udf_def.signature.output.sql_type
 
         remote_function_options = {
             "endpoint": udf_def.endpoint,
@@ -203,7 +202,7 @@ def create_bq_remote_function(
         bq_function_name_escaped = bigframes.core.sql.identifier(sql_func_legal_name)
         create_function_ddl = f"""
             CREATE OR REPLACE FUNCTION `{self._gcp_project_id}.{self._bq_dataset}`.{bq_function_name_escaped}({udf_def.signature.to_sql_input_signature()})
-            RETURNS {bq_function_return_type}
+            RETURNS {udf_def.signature.with_devirtualize().output.sql_type}
             REMOTE WITH CONNECTION `{self._gcp_project_id}.{self._bq_location}.{self._bq_connection_id}`
             OPTIONS ({remote_function_options_str})"""
 
@@ -658,6 +657,7 @@ def provision_bq_remote_function(
             connection_id=self._bq_connection_id,
             max_batching_rows=max_batching_rows,
             signature=func_signature,
+            bq_metadata=func_signature.protocol_metadata,
         )
         remote_function_name = name or get_bigframes_function_name(
             intended_rf_spec,

bigframes/functions/_function_session.py

@@ -596,7 +596,9 @@ def wrapper(func):
                 session=session,  # type: ignore
             )
 
-            udf_sig = udf_def.UdfSignature.from_py_signature(py_sig)
+            udf_sig = udf_def.UdfSignature.from_py_signature(
+                py_sig
+            ).to_remote_function_compatible()
 
             (
                 rf_name,

bigframes/functions/function.py

@@ -235,7 +235,7 @@ def __init__(
         cloud_function_ref: Optional[str] = None,
         is_managed: bool = False,
     ):
-        assert self.udf_def.signature.is_row_processor
+        assert udf_def.signature.is_row_processor
         self._udf_def = udf_def
         self._session = session
         self._local_fun = local_func

bigframes/functions/function_typing.py

@@ -81,7 +81,7 @@ def __init__(self, type_, supported_types):
 
 
 def sdk_type_from_python_type(
-    t: type, allow_lists: bool = False
+    t: type, allow_lists: bool = True
 ) -> bigquery.StandardSqlDataType:
     if (get_origin(t) is list) and allow_lists:
         return sdk_array_output_type_from_python_type(t)

bigframes/functions/udf_def.py

@@ -92,6 +92,12 @@ def stable_hash(self) -> bytes:
 
 @dataclasses.dataclass(frozen=True)
 class DirectScalarType:
+    """
+    Represents a scalar value that is passed directly to the remote function.
+
+    For these values, BigQuery handles the serialization and deserialization without any additional processing.
+    """
+
     _py_type: type
 
     @property
@@ -119,16 +125,26 @@ def stable_hash(self) -> bytes:
     def from_sdk_type(cls, sdk_type: bigquery.StandardSqlDataType) -> DirectScalarType:
         return cls(function_typing.sdk_type_to_py_type(sdk_type))
 
+    @property
+    def emulating_type(self) -> DirectScalarType:
+        return self
+
 
 @dataclasses.dataclass(frozen=True)
 class VirtualListTypeV1:
+    """
+    Represents a list of scalar values that is emulated as a JSON array string in the remote function.
+
+    Only works as output paramter right now where array -> string in function runtime, and then string -> array in SQL post-processing (defined in out_expr()).
+    """
+
     _PROTOCOL_ID = "virtual_list_v1"
 
     inner_dtype: DirectScalarType
 
     @property
     def py_type(self) -> Type[list[Any]]:
-        return list[function_typing.sdk_type_to_py_type(self.inner_dtype)]  # type: ignore
+        return list[self.inner_dtype.py_type]  # type: ignore
 
     # TODO: Specify emulating type and mapping expressions between said types
     @property
@@ -163,6 +179,8 @@ def stable_hash(self) -> bytes:
 class RowSeriesInputFieldV1:
     """
     Used to handle functions that logically take a series as an input, but handled via a string protocol in the remote function.
+
+    For these, the serialization is dependent on index metadata, which must be provided by the caller.
     """
 
     _PROTOCOL_ID = "row_series_input_v1"
@@ -180,6 +198,11 @@ def bf_type(self) -> bigframes.dtypes.Dtype:
     def sql_type(self) -> str:
         return "STRING"
 
+    @property
+    def emulating_type(self) -> DirectScalarType:
+        # Regardless of list inner type, string is used to emulate the list in the remote function.
+        return DirectScalarType(str)
+
     def stable_hash(self) -> bytes:
         hash_val = hashlib.md5()
         hash_val.update(self._PROTOCOL_ID.encode())
@@ -196,14 +219,14 @@ class UdfSignature:
     output: DirectScalarType | VirtualListTypeV1
 
     def __post_init__(self):
-        if any(isinstance(arg, RowSeriesInputFieldV1) for arg in self.inputs):
-            if len(self.inputs) != 1:
-                raise ValueError("Row processor functions must have exactly one input.")
         assert all(isinstance(arg, UdfArg) for arg in self.inputs)
         assert isinstance(self.output, (DirectScalarType, VirtualListTypeV1))
 
     def to_sql_input_signature(self) -> str:
-        return ",".join(f"{field.name} {field.sql_type}" for field in self.inputs)
+        return ",".join(
+            f"{field.name} {field.sql_type}"
+            for field in self.with_devirtualize().inputs
+        )
 
     @property
     def protocol_metadata(self) -> str:
@@ -214,16 +237,20 @@ def protocol_metadata(self) -> str:
             python_output_type=self.output.py_type
         )
 
+    @property
+    def is_virtual(self) -> bool:
+        dtypes = (self.output,) + tuple(arg.dtype for arg in self.inputs)
+        return not all(isinstance(dtype, DirectScalarType) for dtype in dtypes)
+
     @property
     def is_row_processor(self) -> bool:
-        return any(isinstance(arg, RowSeriesInputFieldV1) for arg in self.inputs)
+        return any(isinstance(arg.dtype, RowSeriesInputFieldV1) for arg in self.inputs)
 
     def with_devirtualize(self) -> UdfSignature:
-        if isinstance(self.output, DirectScalarType):
-            return self
-        assert isinstance(self.output, VirtualListTypeV1)
         return UdfSignature(
-            inputs=self.inputs,
+            inputs=tuple(
+                UdfArg(arg.name, arg.dtype.emulating_type) for arg in self.inputs
+            ),
             output=self.output.emulating_type,
         )
 
@@ -245,7 +272,7 @@ def from_routine(cls, routine: bigquery.Routine) -> UdfSignature:
         if python_output_type := bigframes.functions._utils.get_python_output_type_from_bigframes_metadata(
             routine.description
         ):
-            if routine.return_type is None or bq_return_type.type_kind != "STRING":
+            if bq_return_type.type_kind != "STRING":
                 raise bf_formatting.create_exception_with_feedback_link(
                     TypeError,
                     "An explicit output_type should be provided only for a BigQuery function with STRING output.",
@@ -280,7 +307,7 @@ def from_py_signature(cls, signature: inspect.Signature):
                     ValueError,
                     "'input_types' was not set and parameter "
                     f"'{parameter.name}' is missing a type annotation. "
-                    "Types are required to use @remote_function.",
+                    "Types are required to use udfs.",
                 )
 
             input_types.append(UdfArg.from_py_param(parameter))
@@ -290,18 +317,22 @@ def from_py_signature(cls, signature: inspect.Signature):
                 ValueError,
                 "'output_type' was not set and function is missing a "
                 "return type annotation. Types are required to use "
-                "@remote_function.",
+                "udfs.",
             )
 
-        if get_origin(signature.return_annotation) is list:
-            inner_py_type = get_args(signature.return_annotation)[0]
-            virtual_list_output_type = VirtualListTypeV1(
-                DirectScalarType(inner_py_type)
-            )
-            return cls(tuple(input_types), virtual_list_output_type)
-        else:
-            direct_output_type = DirectScalarType(signature.return_annotation)
-            return cls(tuple(input_types), direct_output_type)
+        output_type = DirectScalarType(signature.return_annotation)
+        return cls(tuple(input_types), output_type)
+
+    def to_remote_function_compatible(self) -> UdfSignature:
+        # need to virtualize list outputs
+        if isinstance(self.output, DirectScalarType):
+            if get_origin(self.output.py_type) is list:
+                inner_py_type = get_args(self.output.py_type)[0]
+                return UdfSignature(
+                    inputs=self.inputs,
+                    output=VirtualListTypeV1(DirectScalarType(inner_py_type)),
+                )
+        return self
 
     def stable_hash(self) -> bytes:
         hash_val = hashlib.md5()
@@ -321,6 +352,8 @@ class BigqueryUdf:
     signature: UdfSignature
 
     def with_devirtualize(self) -> BigqueryUdf:
+        if not self.signature.is_virtual:
+            return self
         return BigqueryUdf(
             routine_ref=self.routine_ref,
             signature=self.signature.with_devirtualize(),