refactor: support and, or, xor for sqlglot compiler (#2117)

Author: chelsea-lin

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

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)

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
 

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",
     [

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`

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`

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`

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")