diff --git a/bigframes/core/compile/sqlglot/expressions/generic_ops.py b/bigframes/core/compile/sqlglot/expressions/generic_ops.py index 4a2a5fb213..7e6dceebc1 100644 --- a/bigframes/core/compile/sqlglot/expressions/generic_ops.py +++ b/bigframes/core/compile/sqlglot/expressions/generic_ops.py @@ -140,6 +140,28 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: return sge.Coalesce(this=left.expr, expressions=[right.expr]) +@register_unary_op(ops.RemoteFunctionOp, pass_op=True) +def _(expr: TypedExpr, op: ops.RemoteFunctionOp) -> sge.Expression: + routine_ref = op.function_def.routine_ref + # Quote project, dataset, and routine IDs to avoid keyword clashes. + func_name = ( + f"`{routine_ref.project}`.`{routine_ref.dataset_id}`.`{routine_ref.routine_id}`" + ) + func = sge.func(func_name, expr.expr) + + if not op.apply_on_null: + return sge.Case( + ifs=[ + sge.If( + this=sge.Is(this=expr.expr, expression=sge.Null()), true=expr.expr + ) + ], + default=func, + ) + + return func + + @register_binary_op(ops.BinaryRemoteFunctionOp, pass_op=True) def _( left: TypedExpr, right: TypedExpr, op: ops.BinaryRemoteFunctionOp diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_remote_function_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_remote_function_op/out.sql new file mode 100644 index 0000000000..7f81d0bb30 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_remote_function_op/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `int64_col` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + `my_project`.`my_dataset`.`my_routine`(`int64_col`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `int64_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_generic_ops.py b/tests/unit/core/compile/sqlglot/expressions/test_generic_ops.py index 5657874eb5..14e3bb2030 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_generic_ops.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_generic_ops.py @@ -168,6 +168,37 @@ def test_astype_json_invalid( ) +def test_remote_function_op(scalar_types_df: bpd.DataFrame, snapshot): + from google.cloud import bigquery + + from bigframes.functions import udf_def + + bf_df = scalar_types_df[["int64_col"]] + op = ops.RemoteFunctionOp( + function_def=udf_def.BigqueryUdf( + routine_ref=bigquery.RoutineReference.from_string( + "my_project.my_dataset.my_routine" + ), + signature=udf_def.UdfSignature( + input_types=( + udf_def.UdfField( + "x", + bigquery.StandardSqlDataType( + type_kind=bigquery.StandardSqlTypeNames.INT64 + ), + ), + ), + output_bq_type=bigquery.StandardSqlDataType( + type_kind=bigquery.StandardSqlTypeNames.FLOAT64 + ), + ), + ), + apply_on_null=True, + ) + sql = utils._apply_ops_to_sql(bf_df, [op.as_expr("int64_col")], ["int64_col"]) + snapshot.assert_match(sql, "out.sql") + + def test_binary_remote_function_op(scalar_types_df: bpd.DataFrame, snapshot): from google.cloud import bigquery