Skip to content

Commit dbe8e7e

Browse files
authored
chore: Migrate BinaryRemoteFunctionOp operator to SQLGlot (#2371)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes b/452130300
1 parent b023cb0 commit dbe8e7e

File tree

3 files changed

+64
-0
lines changed

3 files changed

+64
-0
lines changed

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,19 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
140140
return sge.Coalesce(this=left.expr, expressions=[right.expr])
141141

142142

143+
@register_binary_op(ops.BinaryRemoteFunctionOp, pass_op=True)
144+
def _(
145+
left: TypedExpr, right: TypedExpr, op: ops.BinaryRemoteFunctionOp
146+
) -> sge.Expression:
147+
routine_ref = op.function_def.routine_ref
148+
# Quote project, dataset, and routine IDs to avoid keyword clashes.
149+
func_name = (
150+
f"`{routine_ref.project}`.`{routine_ref.dataset_id}`.`{routine_ref.routine_id}`"
151+
)
152+
153+
return sge.func(func_name, left.expr, right.expr)
154+
155+
143156
@register_nary_op(ops.case_when_op)
144157
def _(*cases_and_outputs: TypedExpr) -> sge.Expression:
145158
# Need to upcast BOOL to INT if any output is numeric
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`float64_col`,
4+
`int64_col`
5+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
6+
), `bfcte_1` AS (
7+
SELECT
8+
*,
9+
`my_project`.`my_dataset`.`my_routine`(`int64_col`, `float64_col`) AS `bfcol_2`
10+
FROM `bfcte_0`
11+
)
12+
SELECT
13+
`bfcol_2` AS `int64_col`
14+
FROM `bfcte_1`

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,43 @@ def test_astype_json_invalid(
168168
)
169169

170170

171+
def test_binary_remote_function_op(scalar_types_df: bpd.DataFrame, snapshot):
172+
from google.cloud import bigquery
173+
174+
from bigframes.functions import udf_def
175+
176+
bf_df = scalar_types_df[["int64_col", "float64_col"]]
177+
op = ops.BinaryRemoteFunctionOp(
178+
function_def=udf_def.BigqueryUdf(
179+
routine_ref=bigquery.RoutineReference.from_string(
180+
"my_project.my_dataset.my_routine"
181+
),
182+
signature=udf_def.UdfSignature(
183+
input_types=(
184+
udf_def.UdfField(
185+
"x",
186+
bigquery.StandardSqlDataType(
187+
type_kind=bigquery.StandardSqlTypeNames.INT64
188+
),
189+
),
190+
udf_def.UdfField(
191+
"y",
192+
bigquery.StandardSqlDataType(
193+
type_kind=bigquery.StandardSqlTypeNames.FLOAT64
194+
),
195+
),
196+
),
197+
output_bq_type=bigquery.StandardSqlDataType(
198+
type_kind=bigquery.StandardSqlTypeNames.FLOAT64
199+
),
200+
),
201+
)
202+
)
203+
sql = utils._apply_binary_op(bf_df, op, "int64_col", "float64_col")
204+
205+
snapshot.assert_match(sql, "out.sql")
206+
207+
171208
def test_case_when_op(scalar_types_df: bpd.DataFrame, snapshot):
172209
ops_map = {
173210
"single_case": ops.case_when_op.as_expr(

0 commit comments

Comments
 (0)