refactoring

Author: jialuoo

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

@@ -152,13 +152,17 @@ 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:
+def _get_remote_function_name(op):
     routine_ref = op.function_def.routine_ref
     # Quote project, dataset, and routine IDs to avoid keyword clashes.
-    func_name = (
+    return (
         f"`{routine_ref.project}`.`{routine_ref.dataset_id}`.`{routine_ref.routine_id}`"
     )
+
+
+@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:
@@ -175,23 +179,13 @@ def _(expr: TypedExpr, op: ops.RemoteFunctionOp) -> sge.Expression:
 def _(
     left: TypedExpr, right: TypedExpr, op: ops.BinaryRemoteFunctionOp
 ) -> 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_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:
-    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_name = _get_remote_function_name(op)
     return sge.func(func_name, *(operand.expr for operand in operands))
 
 

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

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from google.cloud import bigquery
 import pandas as pd
 import pytest
 
 from bigframes import dtypes
 from bigframes import operations as ops
 from bigframes.core import expression as ex
+from bigframes.functions import udf_def
 import bigframes.pandas as bpd
 from bigframes.testing import utils
 
@@ -170,10 +172,6 @@ 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"]]
     function_def = udf_def.BigqueryUdf(
         routine_ref=bigquery.RoutineReference.from_string(
@@ -206,10 +204,6 @@ def test_remote_function_op(scalar_types_df: bpd.DataFrame, snapshot):
 
 
 def test_binary_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", "float64_col"]]
     op = ops.BinaryRemoteFunctionOp(
         function_def=udf_def.BigqueryUdf(
@@ -243,10 +237,6 @@ def test_binary_remote_function_op(scalar_types_df: bpd.DataFrame, snapshot):
 
 
 def test_nary_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", "float64_col", "string_col"]]
     op = ops.NaryRemoteFunctionOp(
         function_def=udf_def.BigqueryUdf(