From 4ee74b010c16268b473851d830d9182d71835545 Mon Sep 17 00:00:00 2001 From: jialuo Date: Tue, 28 Oct 2025 21:13:03 +0000 Subject: [PATCH] chore: Migrate round_op operator to SQLGlot This commit migrates the `round_op` operator from the Ibis compiler to the SQLGlot compiler. --- .../sqlglot/expressions/numeric_ops.py | 8 ++ .../test_numeric_ops/test_round/out.sql | 81 +++++++++++++++++++ .../sqlglot/expressions/test_numeric_ops.py | 14 ++++ 3 files changed, 103 insertions(+) create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_round/out.sql diff --git a/bigframes/core/compile/sqlglot/expressions/numeric_ops.py b/bigframes/core/compile/sqlglot/expressions/numeric_ops.py index 3bbe2623ea..b7420ae549 100644 --- a/bigframes/core/compile/sqlglot/expressions/numeric_ops.py +++ b/bigframes/core/compile/sqlglot/expressions/numeric_ops.py @@ -376,6 +376,14 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: return result +@register_binary_op(ops.round_op) +def _(expr: TypedExpr, n_digits: TypedExpr) -> sge.Expression: + rounded = sge.Round(this=expr.expr, decimals=n_digits.expr) + if expr.dtype == dtypes.INT_DTYPE: + return sge.Cast(this=rounded, to="INT64") + return rounded + + @register_binary_op(ops.sub_op) def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: if dtypes.is_numeric(left.dtype) and dtypes.is_numeric(right.dtype): diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_round/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_round/out.sql new file mode 100644 index 0000000000..8513c8d63f --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_round/out.sql @@ -0,0 +1,81 @@ +WITH `bfcte_0` AS ( + SELECT + `int64_col` AS `bfcol_0`, + `float64_col` AS `bfcol_1`, + `rowindex` AS `bfcol_2` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + `bfcol_2` AS `bfcol_6`, + `bfcol_0` AS `bfcol_7`, + `bfcol_1` AS `bfcol_8`, + CAST(ROUND(`bfcol_0`, 0) AS INT64) AS `bfcol_9` + FROM `bfcte_0` +), `bfcte_2` AS ( + SELECT + *, + `bfcol_6` AS `bfcol_14`, + `bfcol_7` AS `bfcol_15`, + `bfcol_8` AS `bfcol_16`, + `bfcol_9` AS `bfcol_17`, + CAST(ROUND(`bfcol_7`, 1) AS INT64) AS `bfcol_18` + FROM `bfcte_1` +), `bfcte_3` AS ( + SELECT + *, + `bfcol_14` AS `bfcol_24`, + `bfcol_15` AS `bfcol_25`, + `bfcol_16` AS `bfcol_26`, + `bfcol_17` AS `bfcol_27`, + `bfcol_18` AS `bfcol_28`, + CAST(ROUND(`bfcol_15`, -1) AS INT64) AS `bfcol_29` + FROM `bfcte_2` +), `bfcte_4` AS ( + SELECT + *, + `bfcol_24` AS `bfcol_36`, + `bfcol_25` AS `bfcol_37`, + `bfcol_26` AS `bfcol_38`, + `bfcol_27` AS `bfcol_39`, + `bfcol_28` AS `bfcol_40`, + `bfcol_29` AS `bfcol_41`, + ROUND(`bfcol_26`, 0) AS `bfcol_42` + FROM `bfcte_3` +), `bfcte_5` AS ( + SELECT + *, + `bfcol_36` AS `bfcol_50`, + `bfcol_37` AS `bfcol_51`, + `bfcol_38` AS `bfcol_52`, + `bfcol_39` AS `bfcol_53`, + `bfcol_40` AS `bfcol_54`, + `bfcol_41` AS `bfcol_55`, + `bfcol_42` AS `bfcol_56`, + ROUND(`bfcol_38`, 1) AS `bfcol_57` + FROM `bfcte_4` +), `bfcte_6` AS ( + SELECT + *, + `bfcol_50` AS `bfcol_66`, + `bfcol_51` AS `bfcol_67`, + `bfcol_52` AS `bfcol_68`, + `bfcol_53` AS `bfcol_69`, + `bfcol_54` AS `bfcol_70`, + `bfcol_55` AS `bfcol_71`, + `bfcol_56` AS `bfcol_72`, + `bfcol_57` AS `bfcol_73`, + ROUND(`bfcol_52`, -1) AS `bfcol_74` + FROM `bfcte_5` +) +SELECT + `bfcol_66` AS `rowindex`, + `bfcol_67` AS `int64_col`, + `bfcol_68` AS `float64_col`, + `bfcol_69` AS `int_round_0`, + `bfcol_70` AS `int_round_1`, + `bfcol_71` AS `int_round_m1`, + `bfcol_72` AS `float_round_0`, + `bfcol_73` AS `float_round_1`, + `bfcol_74` AS `float_round_m1` +FROM `bfcte_6` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py b/tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py index 59726da73b..53590f144f 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py @@ -166,6 +166,20 @@ def test_pos(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") +def test_round(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["int64_col", "float64_col"]] + + bf_df["int_round_0"] = bf_df["int64_col"].round(0) + bf_df["int_round_1"] = bf_df["int64_col"].round(1) + bf_df["int_round_m1"] = bf_df["int64_col"].round(-1) + + bf_df["float_round_0"] = bf_df["float64_col"].round(0) + bf_df["float_round_1"] = bf_df["float64_col"].round(1) + bf_df["float_round_m1"] = bf_df["float64_col"].round(-1) + + snapshot.assert_match(bf_df.sql, "out.sql") + + def test_sqrt(scalar_types_df: bpd.DataFrame, snapshot): col_name = "float64_col" bf_df = scalar_types_df[[col_name]]