diff --git a/bigframes/core/compile/sqlglot/__init__.py b/bigframes/core/compile/sqlglot/__init__.py index 1fc22e1af6..4ceb4118cd 100644 --- a/bigframes/core/compile/sqlglot/__init__.py +++ b/bigframes/core/compile/sqlglot/__init__.py @@ -17,6 +17,7 @@ import bigframes.core.compile.sqlglot.expressions.ai_ops # noqa: F401 import bigframes.core.compile.sqlglot.expressions.array_ops # noqa: F401 import bigframes.core.compile.sqlglot.expressions.blob_ops # noqa: F401 +import bigframes.core.compile.sqlglot.expressions.bool_ops # noqa: F401 import bigframes.core.compile.sqlglot.expressions.comparison_ops # noqa: F401 import bigframes.core.compile.sqlglot.expressions.date_ops # noqa: F401 import bigframes.core.compile.sqlglot.expressions.datetime_ops # noqa: F401 diff --git a/bigframes/core/compile/sqlglot/expressions/bool_ops.py b/bigframes/core/compile/sqlglot/expressions/bool_ops.py new file mode 100644 index 0000000000..41076b666a --- /dev/null +++ b/bigframes/core/compile/sqlglot/expressions/bool_ops.py @@ -0,0 +1,47 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import sqlglot.expressions as sge + +from bigframes import dtypes +from bigframes import operations as ops +from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr +import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler + +register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op + + +@register_binary_op(ops.and_op) +def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: + if left.dtype == dtypes.BOOL_DTYPE and right.dtype == dtypes.BOOL_DTYPE: + return sge.And(this=left.expr, expression=right.expr) + return sge.BitwiseAnd(this=left.expr, expression=right.expr) + + +@register_binary_op(ops.or_op) +def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: + if left.dtype == dtypes.BOOL_DTYPE and right.dtype == dtypes.BOOL_DTYPE: + return sge.Or(this=left.expr, expression=right.expr) + return sge.BitwiseOr(this=left.expr, expression=right.expr) + + +@register_binary_op(ops.xor_op) +def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: + if left.dtype == dtypes.BOOL_DTYPE and right.dtype == dtypes.BOOL_DTYPE: + left_expr = sge.And(this=left.expr, expression=sge.Not(this=right.expr)) + right_expr = sge.And(this=sge.Not(this=left.expr), expression=right.expr) + return sge.Or(this=left_expr, expression=right_expr) + return sge.BitwiseXor(this=left.expr, expression=right.expr) diff --git a/bigframes/operations/type.py b/bigframes/operations/type.py index 020bd0ea57..6542233081 100644 --- a/bigframes/operations/type.py +++ b/bigframes/operations/type.py @@ -204,7 +204,7 @@ def output_type( raise TypeError(f"Type {right_type} is not binary") if left_type != right_type: raise TypeError( - "Bitwise operands {left_type} and {right_type} do not match" + f"Bitwise operands {left_type} and {right_type} do not match" ) return left_type @@ -222,7 +222,7 @@ def output_type( raise TypeError(f"Type {right_type} is not array-like") if left_type != right_type: raise TypeError( - "Vector op operands {left_type} and {right_type} do not match" + f"Vector op operands {left_type} and {right_type} do not match" ) return bigframes.dtypes.FLOAT_DTYPE diff --git a/tests/system/small/engines/test_bool_ops.py b/tests/system/small/engines/test_bool_ops.py index 065a43c209..a77d52b356 100644 --- a/tests/system/small/engines/test_bool_ops.py +++ b/tests/system/small/engines/test_bool_ops.py @@ -46,7 +46,7 @@ def apply_op_pairwise( return new_arr -@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True) +@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True) @pytest.mark.parametrize( "op", [ diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_and_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_and_op/out.sql new file mode 100644 index 0000000000..42c5847401 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_and_op/out.sql @@ -0,0 +1,31 @@ +WITH `bfcte_0` AS ( + SELECT + `bool_col` AS `bfcol_0`, + `int64_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`, + `bfcol_1` & `bfcol_1` 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`, + `bfcol_7` AND `bfcol_7` AS `bfcol_18` + FROM `bfcte_1` +) +SELECT + `bfcol_14` AS `rowindex`, + `bfcol_15` AS `bool_col`, + `bfcol_16` AS `int64_col`, + `bfcol_17` AS `int_and_int`, + `bfcol_18` AS `bool_and_bool` +FROM `bfcte_2` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_or_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_or_op/out.sql new file mode 100644 index 0000000000..d1e7bd1822 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_or_op/out.sql @@ -0,0 +1,31 @@ +WITH `bfcte_0` AS ( + SELECT + `bool_col` AS `bfcol_0`, + `int64_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`, + `bfcol_1` | `bfcol_1` 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`, + `bfcol_7` OR `bfcol_7` AS `bfcol_18` + FROM `bfcte_1` +) +SELECT + `bfcol_14` AS `rowindex`, + `bfcol_15` AS `bool_col`, + `bfcol_16` AS `int64_col`, + `bfcol_17` AS `int_and_int`, + `bfcol_18` AS `bool_and_bool` +FROM `bfcte_2` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_xor_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_xor_op/out.sql new file mode 100644 index 0000000000..7d5f74ede7 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_xor_op/out.sql @@ -0,0 +1,31 @@ +WITH `bfcte_0` AS ( + SELECT + `bool_col` AS `bfcol_0`, + `int64_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`, + `bfcol_1` ^ `bfcol_1` 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`, + `bfcol_7` AND NOT `bfcol_7` OR NOT `bfcol_7` AND `bfcol_7` AS `bfcol_18` + FROM `bfcte_1` +) +SELECT + `bfcol_14` AS `rowindex`, + `bfcol_15` AS `bool_col`, + `bfcol_16` AS `int64_col`, + `bfcol_17` AS `int_and_int`, + `bfcol_18` AS `bool_and_bool` +FROM `bfcte_2` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_bool_ops.py b/tests/unit/core/compile/sqlglot/expressions/test_bool_ops.py new file mode 100644 index 0000000000..08b60d6ddf --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/test_bool_ops.py @@ -0,0 +1,43 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +import bigframes.pandas as bpd + +pytest.importorskip("pytest_snapshot") + + +def test_and_op(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["bool_col", "int64_col"]] + + bf_df["int_and_int"] = bf_df["int64_col"] & bf_df["int64_col"] + bf_df["bool_and_bool"] = bf_df["bool_col"] & bf_df["bool_col"] + snapshot.assert_match(bf_df.sql, "out.sql") + + +def test_or_op(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["bool_col", "int64_col"]] + + bf_df["int_and_int"] = bf_df["int64_col"] | bf_df["int64_col"] + bf_df["bool_and_bool"] = bf_df["bool_col"] | bf_df["bool_col"] + snapshot.assert_match(bf_df.sql, "out.sql") + + +def test_xor_op(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["bool_col", "int64_col"]] + + bf_df["int_and_int"] = bf_df["int64_col"] ^ bf_df["int64_col"] + bf_df["bool_and_bool"] = bf_df["bool_col"] ^ bf_df["bool_col"] + snapshot.assert_match(bf_df.sql, "out.sql")