resolve the comments

Author: jialuoo

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

@@ -150,13 +150,10 @@ def _(expr: TypedExpr, op: ops.RemoteFunctionOp) -> sge.Expression:
     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 sge.If(
+            this=sge.Is(this=expr.expr, expression=sge.Null()),
+            true=expr.expr,
+            false=func,
         )
 
     return func

tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_remote_function_op/out.sql

@@ -5,9 +5,15 @@ WITH `bfcte_0` AS (
 ), `bfcte_1` AS (
   SELECT
     *,
-    `my_project`.`my_dataset`.`my_routine`(`int64_col`) AS `bfcol_1`
+    `my_project`.`my_dataset`.`my_routine`(`int64_col`) AS `bfcol_1`,
+    IF(
+      `int64_col` IS NULL,
+      `int64_col`,
+      `my_project`.`my_dataset`.`my_routine`(`int64_col`)
+    ) AS `bfcol_2`
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_1` AS `int64_col`
+  `bfcol_1` AS `apply_on_null_true`,
+  `bfcol_2` AS `apply_on_null_false`
 FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/test_generic_ops.py

@@ -174,28 +174,33 @@ def test_remote_function_op(scalar_types_df: bpd.DataFrame, snapshot):
     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
-                        ),
+    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
-                ),
+            ),
+            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"])
+    ops_map = {
+        "apply_on_null_true": ops.RemoteFunctionOp(
+            function_def=function_def, apply_on_null=True
+        ).as_expr("int64_col"),
+        "apply_on_null_false": ops.RemoteFunctionOp(
+            function_def=function_def, apply_on_null=False
+        ).as_expr("int64_col"),
+    }
+    sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
     snapshot.assert_match(sql, "out.sql")