diff --git a/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py b/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py index b29a23cd84b8..2bbcd1787ed9 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) + 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")