fix: Respect remote function config changes even if logic unchanged (#2512)

Author: TrevorBergeron

bigframes/core/compile/ibis_compiler/ibis_compiler.py

@@ -88,6 +88,7 @@ def _replace_unsupported_ops(node: nodes.BigFrameNode):
 node=nodes.bottom_up(node, rewrites.rewrite_slice)
 node=nodes.bottom_up(node, rewrites.rewrite_timedelta_expressions)
 node=nodes.bottom_up(node, rewrites.rewrite_range_rolling)
+node=nodes.bottom_up(node, rewrites.lower_udfs)
 returnnode
 
 

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1037,7 +1037,8 @@ def timedelta_floor_op_impl(x: ibis_types.NumericValue):
 @scalar_op_compiler.register_unary_op(ops.RemoteFunctionOp, pass_op=True)
 defremote_function_op_impl(x: ibis_types.Value, op: ops.RemoteFunctionOp):
 udf_sig=op.function_def.signature
-ibis_py_sig= (udf_sig.py_input_types, udf_sig.py_output_type)
+assertnotudf_sig.is_virtual# should have been devirtualized in lowering pass
+ibis_py_sig= (tuple(arg.py_typeforarginudf_sig.inputs), udf_sig.output.py_type)
 
 @ibis_udf.scalar.builtin(
 name=str(op.function_def.routine_ref), signature=ibis_py_sig
@@ -1056,7 +1057,8 @@ def binary_remote_function_op_impl(
 x: ibis_types.Value, y: ibis_types.Value, op: ops.BinaryRemoteFunctionOp
 ):
 udf_sig=op.function_def.signature
-ibis_py_sig= (udf_sig.py_input_types, udf_sig.py_output_type)
+assertnotudf_sig.is_virtual# should have been devirtualized in lowering pass
+ibis_py_sig= (tuple(arg.py_typeforarginudf_sig.inputs), udf_sig.output.py_type)
 
 @ibis_udf.scalar.builtin(
 name=str(op.function_def.routine_ref), signature=ibis_py_sig
@@ -1073,8 +1075,9 @@ def nary_remote_function_op_impl(
 *operands: ibis_types.Value, op: ops.NaryRemoteFunctionOp
 ):
 udf_sig=op.function_def.signature
-ibis_py_sig= (udf_sig.py_input_types, udf_sig.py_output_type)
-arg_names=tuple(arg.nameforarginudf_sig.input_types)
+assertnotudf_sig.is_virtual# should have been devirtualized in lowering pass
+ibis_py_sig= (tuple(arg.py_typeforarginudf_sig.inputs), udf_sig.output.py_type)
+arg_names=tuple(arg.nameforarginudf_sig.inputs)
 
 @ibis_udf.scalar.builtin(
 name=str(op.function_def.routine_ref),
@@ -1153,6 +1156,13 @@ def array_reduce_op_impl(x: ibis_types.Value, op: ops.ArrayReduceOp):
  )
 
 
+@scalar_op_compiler.register_unary_op(ops.ArrayMapOp, pass_op=True)
+defarray_map_op_impl(x: ibis_types.Value, op: ops.ArrayMapOp):
+returntyping.cast(ibis_types.ArrayValue, x).map(
+lambdaarr_vals: scalar_op_compiler.compile_row_op(op.map_op, (arr_vals,))
+ )
+
+
 # JSON Ops
 @scalar_op_compiler.register_binary_op(ops.JSONSet, pass_op=True)
 defjson_set_op_impl(x: ibis_types.Value, y: ibis_types.Value, op: ops.JSONSet):

bigframes/core/compile/sqlglot/compiler.py

@@ -369,4 +369,5 @@ def compile_aggregate(
 def_replace_unsupported_ops(node: nodes.BigFrameNode):
 node=nodes.bottom_up(node, rewrite.rewrite_slice)
 node=nodes.bottom_up(node, rewrite.rewrite_range_rolling)
+node=nodes.bottom_up(node, rewrite.lower_udfs)
 returnnode

bigframes/core/compile/sqlglot/expressions/array_ops.py

@@ -73,6 +73,28 @@ def _(expr: TypedExpr, op: ops.ArrayReduceOp) -> sge.Expression:
  )
 
 
+@register_unary_op(ops.ArrayMapOp, pass_op=True)
+def_(expr: TypedExpr, op: ops.ArrayMapOp) ->sge.Expression:
+sub_expr=sg.to_identifier("bf_arr_map_uid")
+sub_type=dtypes.get_array_inner_type(expr.dtype)
+
+# TODO: Expression should be provided instead of invoking compiler manually
+map_expr=expression_compiler.expression_compiler.compile_row_op(
+op.map_op, (TypedExpr(sub_expr, sub_type),)
+ )
+
+returnsge.array(
+sge.select(map_expr)
+ .from_(
+sge.Unnest(
+expressions=[expr.expr],
+alias=sge.TableAlias(columns=[sub_expr]),
+ )
+ )
+ .subquery()
+ )
+
+
 @register_unary_op(ops.ArraySliceOp, pass_op=True)
 def_(expr: TypedExpr, op: ops.ArraySliceOp) ->sge.Expression:
 ifexpr.dtype==dtypes.STRING_DTYPE:

bigframes/core/compile/sqlglot/sql/__init__.py

@@ -15,7 +15,6 @@
 
 frombigframes.core.compile.sqlglot.sql.baseimport (
 cast,
-escape_chars,
 identifier,
 is_null_literal,
 literal,
@@ -28,7 +27,6 @@
 __all__= [
 # From base.py
 "cast",
-"escape_chars",
 "identifier",
 "is_null_literal",
 "literal",

bigframes/core/compile/sqlglot/sql/base.py

@@ -136,29 +136,6 @@ def table(table: bigquery.TableReference) -> sge.Table:
  )
 
 
-defescape_chars(value: str):
-"""Escapes all special characters"""
-# TODO: Reuse literal's escaping logic instead of re-implementing it here.
-# https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#string_and_bytes_literals
-trans_table=str.maketrans(
- {
-"\a": r"\a",
-"\b": r"\b",
-"\f": r"\f",
-"\n": r"\n",
-"\r": r"\r",
-"\t": r"\t",
-"\v": r"\v",
-"\\": r"\\",
-"?": r"\?",
-'"': r"\"",
-"'": r"\'",
-"`": r"\`",
- }
- )
-returnvalue.translate(trans_table)
-
-
 defis_null_literal(expr: sge.Expression) ->bool:
 """Checks if the given expression is a NULL literal."""
 ifisinstance(expr, sge.Null):

bigframes/core/rewrite/__init__.py

@@ -27,6 +27,7 @@
 frombigframes.core.rewrite.select_pullupimportdefer_selection
 frombigframes.core.rewrite.slicesimportpull_out_limit, pull_up_limits, rewrite_slice
 frombigframes.core.rewrite.timedeltasimportrewrite_timedelta_expressions
+frombigframes.core.rewrite.udfsimportlower_udfs
 frombigframes.core.rewrite.windowsimport (
 pull_out_window_order,
 rewrite_range_rolling,
@@ -53,4 +54,5 @@
 "pull_out_window_order",
 "defer_selection",
 "simplify_complex_windows",
+"lower_udfs",
 ]

bigframes/core/rewrite/udfs.py

@@ -0,0 +1,87 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 __future__ importannotations
+
+importdataclasses
+
+frombigframes.coreimportbigframe_node, expression
+frombigframes.core.rewriteimportop_lowering
+importbigframes.functions.udf_defasudf_def
+importbigframes.operationsasops
+
+
+@dataclasses.dataclass
+classLowerRemoteFunctionRule(op_lowering.OpLoweringRule):
+@property
+defop(self) ->type[ops.ScalarOp]:
+returnops.RemoteFunctionOp
+
+deflower(self, expr: expression.OpExpression) ->expression.Expression:
+assertisinstance(expr.op, ops.RemoteFunctionOp)
+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)
+ifisinstance(func_def.signature.output, udf_def.VirtualListTypeV1):
+returnfunc_def.signature.output.out_expr(devirtualized_expr)
+else:
+returndevirtualized_expr
+
+
+@dataclasses.dataclass
+classLowerBinaryRemoteFunctionRule(op_lowering.OpLoweringRule):
+@property
+defop(self) ->type[ops.ScalarOp]:
+returnops.BinaryRemoteFunctionOp
+
+deflower(self, expr: expression.OpExpression) ->expression.Expression:
+assertisinstance(expr.op, ops.BinaryRemoteFunctionOp)
+func_def=expr.op.function_def
+devirtualized_expr=ops.BinaryRemoteFunctionOp(
+func_def.with_devirtualize(),
+ ).as_expr(*expr.children)
+ifisinstance(func_def.signature.output, udf_def.VirtualListTypeV1):
+returnfunc_def.signature.output.out_expr(devirtualized_expr)
+else:
+returndevirtualized_expr
+
+
+@dataclasses.dataclass
+classLowerNaryRemoteFunctionRule(op_lowering.OpLoweringRule):
+@property
+defop(self) ->type[ops.ScalarOp]:
+returnops.NaryRemoteFunctionOp
+
+deflower(self, expr: expression.OpExpression) ->expression.Expression:
+assertisinstance(expr.op, ops.NaryRemoteFunctionOp)
+func_def=expr.op.function_def
+devirtualized_expr=ops.NaryRemoteFunctionOp(
+func_def.with_devirtualize(),
+ ).as_expr(*expr.children)
+ifisinstance(func_def.signature.output, udf_def.VirtualListTypeV1):
+returnfunc_def.signature.output.out_expr(devirtualized_expr)
+else:
+returndevirtualized_expr
+
+
+UDF_LOWERING_RULES= (
+LowerRemoteFunctionRule(),
+LowerBinaryRemoteFunctionRule(),
+LowerNaryRemoteFunctionRule(),
+)
+
+
+deflower_udfs(root: bigframe_node.BigFrameNode) ->bigframe_node.BigFrameNode:
+returnop_lowering.lower_ops(root, rules=UDF_LOWERING_RULES)

bigframes/core/sql/__init__.py

@@ -48,6 +48,35 @@
 to_wkt=dumps
 
 
+defidentifier(name: str) ->str:
+iflen(name) >256:
+raiseValueError("Identifier must be less than 256 characters")
+returnf"`{escape_chars(name)}`"
+
+
+defescape_chars(value: str):
+"""Escapes all special characters"""
+# TODO: Reuse literal's escaping logic instead of re-implementing it here.
+# https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#string_and_bytes_literals
+trans_table=str.maketrans(
+ {
+"\a": r"\a",
+"\b": r"\b",
+"\f": r"\f",
+"\n": r"\n",
+"\r": r"\r",
+"\t": r"\t",
+"\v": r"\v",
+"\\": r"\\",
+"?": r"\?",
+'"': r"\"",
+"'": r"\'",
+"`": r"\`",
+ }
+ )
+returnvalue.translate(trans_table)
+
+
 defmulti_literal(*values: Any):
 literal_strings= [sql.to_sql(sql.literal(i)) foriinvalues]
 return"("+", ".join(literal_strings) +")"

bigframes/dataframe.py

@@ -4748,7 +4748,9 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
 # compatible with the data types of the input params.
 # 3. The order of the columns in the dataframe must correspond
 # to the order of the input params in the function.
-udf_input_dtypes=func.udf_def.signature.bf_input_types
+udf_input_dtypes=tuple(
+arg.bf_typeforarginfunc.udf_def.signature.inputs
+ )
 ifnotargsandlen(udf_input_dtypes) !=len(self.columns):
 raiseValueError(
 f"Parameter count mismatch: BigFrames BigQuery function"
@@ -4793,7 +4795,6 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
  )
 result_series.name=None
 
-result_series=func._post_process_series(result_series)
 returnresult_series
 
 # At this point column-wise or element-wise bigquery function operation will