From c28669d203e7056160de6b43a394fc9648c42a2c Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Thu, 13 Aug 2026 20:29:25 +0000 Subject: [PATCH 1/2] fix(bigframes): parenthesize OR expressions for AND chaining in SQLGlot --- .../core/compile/sqlglot/sqlglot_ir.py | 10 +++++++- .../test_compile_filter_w_multiple_or/out.sql | 25 +++++++++++++++++++ .../test_compile_filter_w_or_first/out.sql | 15 +++++++++++ .../test_compile_filter_w_or_second/out.sql | 15 +++++++++++ .../compile/sqlglot/test_compile_filter.py | 25 +++++++++++++++++++ 5 files changed, 89 insertions(+), 1 deletion(-) create mode 100644 packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_multiple_or/out.sql create mode 100644 packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_or_first/out.sql create mode 100644 packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_or_second/out.sql diff --git a/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py b/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py index b29a23cd84b8..92d9b02b3ad3 100644 --- a/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py +++ b/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py @@ -581,8 +581,16 @@ def _and(conditions: tuple[sge.Expression, ...]) -> typing.Optional[sge.Expressi if not conditions: return None + def check_and_parenthesize(expr: sge.Expression) -> sge.Expression: + if isinstance(expr, sge.Or): + return sge.paren(expr, copy=False) + return expr + return functools.reduce( - lambda left, right: sge.And(this=left, expression=right), conditions + lambda left, right: sge.And( + this=check_and_parenthesize(left), expression=check_and_parenthesize(right) + ), + conditions, ) diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_multiple_or/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_multiple_or/out.sql new file mode 100644 index 000000000000..cc059008d503 --- /dev/null +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_multiple_or/out.sql @@ -0,0 +1,25 @@ +SELECT + `rowindex`, + `rowindex` AS `rowindex_1`, + `int64_col`, + `float64_col`, + `string_col` +FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` +WHERE + ( + ( + `rowindex` = 1 + ) OR ( + `int64_col` = 2 + ) + ) + AND ( + IF(( + `rowindex` = 1 + ) OR ( + `int64_col` = 2 + ), `float64_col` > 0, NULL) + OR ( + `string_col` = 'a' + ) + ) \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_or_first/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_or_first/out.sql new file mode 100644 index 000000000000..c821a53b8a05 --- /dev/null +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_or_first/out.sql @@ -0,0 +1,15 @@ +SELECT + `rowindex`, + `rowindex` AS `rowindex_1`, + `int64_col`, + `string_col` +FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` +WHERE + ( + ( + `rowindex` = 1 + ) OR ( + `int64_col` = 2 + ) + ) + AND STARTS_WITH(`string_col`, 'H') \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_or_second/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_or_second/out.sql new file mode 100644 index 000000000000..68414ee9eb1f --- /dev/null +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter_w_or_second/out.sql @@ -0,0 +1,15 @@ +SELECT + `rowindex`, + `rowindex` AS `rowindex_1`, + `int64_col`, + `string_col` +FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` +WHERE + STARTS_WITH(`string_col`, 'H') + AND ( + ( + `rowindex` = 1 + ) OR ( + `int64_col` = 2 + ) + ) \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/test_compile_filter.py b/packages/bigframes/tests/unit/core/compile/sqlglot/test_compile_filter.py index 0afb5eb45b9d..c46204da09b8 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/test_compile_filter.py +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/test_compile_filter.py @@ -22,4 +22,29 @@ def test_compile_filter(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["rowindex", "int64_col"]] bf_filter = bf_df[bf_df["rowindex"] >= 1] + snapshot.assert_match(bf_filter.sql, "out.sql") + + +def test_compile_filter_w_or_first(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["rowindex", "int64_col", "string_col"]] + filtered = bf_df[(bf_df["rowindex"] == 1) | (bf_df["int64_col"] == 2)] + filtered = filtered[filtered["string_col"].str.startswith("H")] + + snapshot.assert_match(filtered.sql, "out.sql") + + +def test_compile_filter_w_or_second(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["rowindex", "int64_col", "string_col"]] + filtered = bf_df[bf_df["string_col"].str.startswith("H")] + filtered = filtered[(filtered["rowindex"] == 1) | (filtered["int64_col"] == 2)] + + snapshot.assert_match(filtered.sql, "out.sql") + + +def test_compile_filter_w_multiple_or(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["rowindex", "int64_col", "float64_col", "string_col"]] + filtered = bf_df[(bf_df["rowindex"] == 1) | (bf_df["int64_col"] == 2)] + filtered = filtered[(filtered["float64_col"] > 0) | (bf_df["string_col"] == "a")] + + snapshot.assert_match(filtered.sql, "out.sql") From aeaea98a50f786e9043542d4f4946593a806db67 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Thu, 13 Aug 2026 21:48:06 +0000 Subject: [PATCH 2/2] perform deep copy --- packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py b/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py index 92d9b02b3ad3..2bbcd1787ed9 100644 --- a/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py +++ b/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py @@ -583,7 +583,7 @@ def _and(conditions: tuple[sge.Expression, ...]) -> typing.Optional[sge.Expressi def check_and_parenthesize(expr: sge.Expression) -> sge.Expression: if isinstance(expr, sge.Or): - return sge.paren(expr, copy=False) + return sge.paren(expr) return expr return functools.reduce(