feat(bigframes): Defer unnamed @udf deployment until needed (#17217)

Author: TrevorBergeron

packages/bigframes/bigframes/core/blocks.py

@@ -1091,14 +1091,14 @@ def multi_apply_window_op(
 
     def multi_apply_unary_op(
         self,
-        op: Union[ops.UnaryOp, ex.Expression],
+        op: Union[ops.UnaryOp, ops.NaryOp, ex.Expression],
     ) -> Block:
-        if isinstance(op, ops.UnaryOp):
+        if isinstance(op, (ops.UnaryOp, ops.NaryOp)):
             input_varname = guid.generate_guid()
             expr = op.as_expr(ex.free_var(input_varname))
         else:
             input_varnames = op.free_variables
-            assert len(input_varnames) == 1
+            assert len(set(input_varnames)) == 1
             expr = op
             input_varname = input_varnames[0]
 

packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1034,44 +1034,8 @@ def timedelta_floor_op_impl(x: ibis_types.NumericValue):
     return ibis_api.case().when(x > ibis.literal(0), x.floor()).else_(x.ceil()).end()
 
 
-@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(
-        name=str(op.function_def.routine_ref), signature=ibis_py_sig
-    )
-    def udf(input): ...
-
-    x_transformed = udf(x)
-    if not op.apply_on_null:
-        return ibis_api.case().when(x.isnull(), x).else_(x_transformed).end()
-    return x_transformed
-
-
-@scalar_op_compiler.register_binary_op(ops.BinaryRemoteFunctionOp, pass_op=True)
-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(
-        name=str(op.function_def.routine_ref), signature=ibis_py_sig
-    )
-    def udf(input1, input2): ...
-
-    x_transformed = udf(x, y)
-    return x_transformed
-
-
-@scalar_op_compiler.register_nary_op(ops.NaryRemoteFunctionOp, pass_op=True)
-def nary_remote_function_op_impl(
-    *operands: ibis_types.Value, op: ops.NaryRemoteFunctionOp
-):
+@scalar_op_compiler.register_nary_op(ops.RemoteFunctionOp, pass_op=True)
+def remote_function_op_impl(*values: 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)
@@ -1084,8 +1048,7 @@ def nary_remote_function_op_impl(
     )
     def udf(*inputs): ...
 
-    result = udf(*operands)
-    return result
+    return udf(*values)
 
 
 @scalar_op_compiler.register_unary_op(ops.MapOp, pass_op=True)

packages/bigframes/bigframes/core/compile/sqlglot/expressions/generic_ops.py

@@ -186,33 +186,9 @@ def _get_remote_function_name(op):
     )
 
 
-@register_unary_op(ops.RemoteFunctionOp, pass_op=True)
-def _(expr: TypedExpr, op: ops.RemoteFunctionOp) -> sge.Expression:
-    func_name = _get_remote_function_name(op)
-    func = sge.func(func_name, expr.expr)
-
-    if not op.apply_on_null:
-        return sge.If(
-            this=sge.Is(this=expr.expr, expression=sge.Null()),
-            true=expr.expr,
-            false=func,
-        )
-
-    return func
-
-
-@register_binary_op(ops.BinaryRemoteFunctionOp, pass_op=True)
-def _(
-    left: TypedExpr, right: TypedExpr, op: ops.BinaryRemoteFunctionOp
-) -> sge.Expression:
-    func_name = _get_remote_function_name(op)
-    return sge.func(func_name, left.expr, right.expr)
-
-
-@register_nary_op(ops.NaryRemoteFunctionOp, pass_op=True)
-def _(*operands: TypedExpr, op: ops.NaryRemoteFunctionOp) -> sge.Expression:
-    func_name = _get_remote_function_name(op)
-    return sge.func(func_name, *(operand.expr for operand in operands))
+@register_nary_op(ops.RemoteFunctionOp, pass_op=True)
+def _(*values: TypedExpr, op: ops.RemoteFunctionOp) -> sge.Expression:
+    return sge.func(_get_remote_function_name(op), *(value.expr for value in values))
 
 
 @register_nary_op(ops.case_when_op)

packages/bigframes/bigframes/core/rewrite/udfs.py

@@ -32,55 +32,14 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
         func_def = expr.op.function_def
         devirtualized_expr = ops.RemoteFunctionOp(
             func_def.with_devirtualize(),
-            apply_on_null=expr.op.apply_on_null,
         ).as_expr(*expr.children)
         if isinstance(func_def.signature.output, udf_def.VirtualListTypeV1):
             return func_def.signature.output.out_expr(devirtualized_expr)
         else:
             return devirtualized_expr
 
 
-@dataclasses.dataclass
-class LowerBinaryRemoteFunctionRule(op_lowering.OpLoweringRule):
-    @property
-    def op(self) -> type[ops.ScalarOp]:
-        return ops.BinaryRemoteFunctionOp
-
-    def lower(self, expr: expression.OpExpression) -> expression.Expression:
-        assert isinstance(expr.op, ops.BinaryRemoteFunctionOp)
-        func_def = expr.op.function_def
-        devirtualized_expr = ops.BinaryRemoteFunctionOp(
-            func_def.with_devirtualize(),
-        ).as_expr(*expr.children)
-        if isinstance(func_def.signature.output, udf_def.VirtualListTypeV1):
-            return func_def.signature.output.out_expr(devirtualized_expr)
-        else:
-            return devirtualized_expr
-
-
-@dataclasses.dataclass
-class LowerNaryRemoteFunctionRule(op_lowering.OpLoweringRule):
-    @property
-    def op(self) -> type[ops.ScalarOp]:
-        return ops.NaryRemoteFunctionOp
-
-    def lower(self, expr: expression.OpExpression) -> expression.Expression:
-        assert isinstance(expr.op, ops.NaryRemoteFunctionOp)
-        func_def = expr.op.function_def
-        devirtualized_expr = ops.NaryRemoteFunctionOp(
-            func_def.with_devirtualize(),
-        ).as_expr(*expr.children)
-        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 = (
-    LowerRemoteFunctionRule(),
-    LowerBinaryRemoteFunctionRule(),
-    LowerNaryRemoteFunctionRule(),
-)
+UDF_LOWERING_RULES = (LowerRemoteFunctionRule(),)
 
 
 def lower_udfs(root: bigframe_node.BigFrameNode) -> bigframe_node.BigFrameNode:

packages/bigframes/bigframes/dataframe.py

@@ -4694,12 +4694,14 @@ def map(self, func, na_action: Optional[str] = None) -> DataFrame:
         if na_action not in {None, "ignore"}:
             raise ValueError(f"na_action={na_action} not supported")
 
-        # TODO(shobs): Support **kwargs
-        return self._apply_unary_op(
-            ops.RemoteFunctionOp(
-                function_def=func.udf_def, apply_on_null=(na_action is None)
+        expr = ops.func_to_op(func).as_expr(ex.free_var("input"))
+        if na_action == "ignore":
+            # True case, predicate, False case
+            expr = ops.where_op.as_expr(
+                expr, ops.notnull_op.as_expr(ex.free_var("input")), ex.const(None)
             )
-        )
+
+        return DataFrame(self._block.multi_apply_unary_op(expr))
 
     def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
         # In Bigframes BigQuery function, DataFrame '.apply' method is specifically
@@ -4770,17 +4772,11 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
                 )
 
                 # Apply the function
-                if args:
-                    result_series = rows_as_json_series._apply_nary_op(
-                        ops.NaryRemoteFunctionOp(function_def=func.udf_def),
-                        list(args),
-                    )
-                else:
-                    result_series = rows_as_json_series._apply_unary_op(
-                        ops.RemoteFunctionOp(
-                            function_def=func.udf_def, apply_on_null=True
-                        )
-                    )
+                result_series = rows_as_json_series._apply_nary_op(
+                    ops.func_to_op(func),
+                    list(args),
+                )
+
             else:
                 # This is a special case where we are providing not-pandas-like
                 # extension. If the bigquery function can take one or more
@@ -4838,7 +4834,7 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
                 series_list = [self[col] for col in self.columns]
                 op_list = series_list[1:] + list(args)
                 result_series = series_list[0]._apply_nary_op(
-                    ops.NaryRemoteFunctionOp(function_def=func.udf_def), op_list
+                    ops.func_to_op(func), op_list
                 )
             result_series.name = None