Skip to content

Commit 3702f56

Browse files
authored
chore: Migrate manhattan_distance_op operator to SQLGlot (#2254)
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/447388852 🦕
1 parent 86ed01b commit 3702f56

File tree

3 files changed

+38
-0
lines changed

3 files changed

+38
-0
lines changed

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -350,6 +350,13 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
350350
return result
351351

352352

353+
@register_binary_op(ops.manhattan_distance_op)
354+
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
355+
return sge.func(
356+
"ML.DISTANCE", left.expr, right.expr, sge.Literal.string("MANHATTAN")
357+
)
358+
359+
353360
@register_binary_op(ops.mod_op)
354361
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
355362
# In BigQuery returned value has the same sign as X. In pandas, the sign of y is used, so we need to flip the result if sign(x) != sign(y)
Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`float_list_col`,
4+
`numeric_list_col`
5+
FROM `bigframes-dev`.`sqlglot_test`.`repeated_types`
6+
), `bfcte_1` AS (
7+
SELECT
8+
*,
9+
ML.DISTANCE(`float_list_col`, `float_list_col`, 'MANHATTAN') AS `bfcol_2`,
10+
ML.DISTANCE(`numeric_list_col`, `numeric_list_col`, 'MANHATTAN') AS `bfcol_3`
11+
FROM `bfcte_0`
12+
)
13+
SELECT
14+
`bfcol_2` AS `float_list_col`,
15+
`bfcol_3` AS `numeric_list_col`
16+
FROM `bfcte_1`

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,21 @@ def test_floordiv_timedelta(scalar_types_df: bpd.DataFrame, snapshot):
354354
snapshot.assert_match(bf_df.sql, "out.sql")
355355

356356

357+
def test_manhattan_distance(repeated_types_df: bpd.DataFrame, snapshot):
358+
col_names = ["float_list_col", "numeric_list_col"]
359+
bf_df = repeated_types_df[col_names]
360+
361+
sql = utils._apply_ops_to_sql(
362+
bf_df,
363+
[
364+
ops.manhattan_distance_op.as_expr("float_list_col", "float_list_col"),
365+
ops.manhattan_distance_op.as_expr("numeric_list_col", "numeric_list_col"),
366+
],
367+
["float_list_col", "numeric_list_col"],
368+
)
369+
snapshot.assert_match(sql, "out.sql")
370+
371+
357372
def test_mul_numeric(scalar_types_df: bpd.DataFrame, snapshot):
358373
bf_df = scalar_types_df[["int64_col", "bool_col"]]
359374

0 commit comments

Comments
 (0)