Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

Commit daea7c9

Browse files
sycaisycai
authored
Merge branch 'main' into sycai_td_sample
2 parents fe18046 + 10b2a38 commit daea7c9

8 files changed

Lines changed: 187 additions & 3 deletions

File tree

‎bigframes/core/compile/sqlglot/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import bigframes.core.compile.sqlglot.expressions.ai_ops # noqa: F401
1818
import bigframes.core.compile.sqlglot.expressions.array_ops # noqa: F401
1919
import bigframes.core.compile.sqlglot.expressions.blob_ops # noqa: F401
20+
import bigframes.core.compile.sqlglot.expressions.bool_ops # noqa: F401
2021
import bigframes.core.compile.sqlglot.expressions.comparison_ops # noqa: F401
2122
import bigframes.core.compile.sqlglot.expressions.date_ops # noqa: F401
2223
import bigframes.core.compile.sqlglot.expressions.datetime_ops # noqa: F401

‎bigframes/core/compile/sqlglot/expressions/bool_ops.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
import sqlglot.expressions as sge
18+
19+
from bigframes import dtypes
20+
from bigframes import operations as ops
21+
from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
22+
import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
23+
24+
register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op
25+
26+
27+
@register_binary_op(ops.and_op)
28+
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
29+
if left.dtype == dtypes.BOOL_DTYPE and right.dtype == dtypes.BOOL_DTYPE:
30+
return sge.And(this=left.expr, expression=right.expr)
31+
return sge.BitwiseAnd(this=left.expr, expression=right.expr)
32+
33+
34+
@register_binary_op(ops.or_op)
35+
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
36+
if left.dtype == dtypes.BOOL_DTYPE and right.dtype == dtypes.BOOL_DTYPE:
37+
return sge.Or(this=left.expr, expression=right.expr)
38+
return sge.BitwiseOr(this=left.expr, expression=right.expr)
39+
40+
41+
@register_binary_op(ops.xor_op)
42+
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
43+
if left.dtype == dtypes.BOOL_DTYPE and right.dtype == dtypes.BOOL_DTYPE:
44+
left_expr = sge.And(this=left.expr, expression=sge.Not(this=right.expr))
45+
right_expr = sge.And(this=sge.Not(this=left.expr), expression=right.expr)
46+
return sge.Or(this=left_expr, expression=right_expr)
47+
return sge.BitwiseXor(this=left.expr, expression=right.expr)

‎bigframes/operations/type.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -204,7 +204,7 @@ def output_type(
204204
raise TypeError(f"Type {right_type} is not binary")
205205
if left_type != right_type:
206206
raise TypeError(
207-
"Bitwise operands {left_type} and {right_type} do not match"
207+
f"Bitwise operands {left_type} and {right_type} do not match"
208208
)
209209
return left_type
210210

@@ -222,7 +222,7 @@ def output_type(
222222
raise TypeError(f"Type {right_type} is not array-like")
223223
if left_type != right_type:
224224
raise TypeError(
225-
"Vector op operands {left_type} and {right_type} do not match"
225+
f"Vector op operands {left_type} and {right_type} do not match"
226226
)
227227
return bigframes.dtypes.FLOAT_DTYPE
228228

‎tests/system/small/engines/test_bool_ops.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ def apply_op_pairwise(
4646
return new_arr
4747

4848

49-
@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
49+
@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
5050
@pytest.mark.parametrize(
5151
"op",
5252
[

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_and_op/out.sql‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`bool_col` AS `bfcol_0`,
4+
`int64_col` AS `bfcol_1`,
5+
`rowindex` AS `bfcol_2`
6+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
7+
), `bfcte_1` AS (
8+
SELECT
9+
*,
10+
`bfcol_2` AS `bfcol_6`,
11+
`bfcol_0` AS `bfcol_7`,
12+
`bfcol_1` AS `bfcol_8`,
13+
`bfcol_1` & `bfcol_1` AS `bfcol_9`
14+
FROM `bfcte_0`
15+
), `bfcte_2` AS (
16+
SELECT
17+
*,
18+
`bfcol_6` AS `bfcol_14`,
19+
`bfcol_7` AS `bfcol_15`,
20+
`bfcol_8` AS `bfcol_16`,
21+
`bfcol_9` AS `bfcol_17`,
22+
`bfcol_7` AND `bfcol_7` AS `bfcol_18`
23+
FROM `bfcte_1`
24+
)
25+
SELECT
26+
`bfcol_14` AS `rowindex`,
27+
`bfcol_15` AS `bool_col`,
28+
`bfcol_16` AS `int64_col`,
29+
`bfcol_17` AS `int_and_int`,
30+
`bfcol_18` AS `bool_and_bool`
31+
FROM `bfcte_2`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_or_op/out.sql‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`bool_col` AS `bfcol_0`,
4+
`int64_col` AS `bfcol_1`,
5+
`rowindex` AS `bfcol_2`
6+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
7+
), `bfcte_1` AS (
8+
SELECT
9+
*,
10+
`bfcol_2` AS `bfcol_6`,
11+
`bfcol_0` AS `bfcol_7`,
12+
`bfcol_1` AS `bfcol_8`,
13+
`bfcol_1` | `bfcol_1` AS `bfcol_9`
14+
FROM `bfcte_0`
15+
), `bfcte_2` AS (
16+
SELECT
17+
*,
18+
`bfcol_6` AS `bfcol_14`,
19+
`bfcol_7` AS `bfcol_15`,
20+
`bfcol_8` AS `bfcol_16`,
21+
`bfcol_9` AS `bfcol_17`,
22+
`bfcol_7` OR `bfcol_7` AS `bfcol_18`
23+
FROM `bfcte_1`
24+
)
25+
SELECT
26+
`bfcol_14` AS `rowindex`,
27+
`bfcol_15` AS `bool_col`,
28+
`bfcol_16` AS `int64_col`,
29+
`bfcol_17` AS `int_and_int`,
30+
`bfcol_18` AS `bool_and_bool`
31+
FROM `bfcte_2`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_xor_op/out.sql‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`bool_col` AS `bfcol_0`,
4+
`int64_col` AS `bfcol_1`,
5+
`rowindex` AS `bfcol_2`
6+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
7+
), `bfcte_1` AS (
8+
SELECT
9+
*,
10+
`bfcol_2` AS `bfcol_6`,
11+
`bfcol_0` AS `bfcol_7`,
12+
`bfcol_1` AS `bfcol_8`,
13+
`bfcol_1` ^ `bfcol_1` AS `bfcol_9`
14+
FROM `bfcte_0`
15+
), `bfcte_2` AS (
16+
SELECT
17+
*,
18+
`bfcol_6` AS `bfcol_14`,
19+
`bfcol_7` AS `bfcol_15`,
20+
`bfcol_8` AS `bfcol_16`,
21+
`bfcol_9` AS `bfcol_17`,
22+
`bfcol_7` AND NOT `bfcol_7` OR NOT `bfcol_7` AND `bfcol_7` AS `bfcol_18`
23+
FROM `bfcte_1`
24+
)
25+
SELECT
26+
`bfcol_14` AS `rowindex`,
27+
`bfcol_15` AS `bool_col`,
28+
`bfcol_16` AS `int64_col`,
29+
`bfcol_17` AS `int_and_int`,
30+
`bfcol_18` AS `bool_and_bool`
31+
FROM `bfcte_2`

‎tests/unit/core/compile/sqlglot/expressions/test_bool_ops.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import pytest
16+
17+
import bigframes.pandas as bpd
18+
19+
pytest.importorskip("pytest_snapshot")
20+
21+
22+
def test_and_op(scalar_types_df: bpd.DataFrame, snapshot):
23+
bf_df = scalar_types_df[["bool_col", "int64_col"]]
24+
25+
bf_df["int_and_int"] = bf_df["int64_col"] & bf_df["int64_col"]
26+
bf_df["bool_and_bool"] = bf_df["bool_col"] & bf_df["bool_col"]
27+
snapshot.assert_match(bf_df.sql, "out.sql")
28+
29+
30+
def test_or_op(scalar_types_df: bpd.DataFrame, snapshot):
31+
bf_df = scalar_types_df[["bool_col", "int64_col"]]
32+
33+
bf_df["int_and_int"] = bf_df["int64_col"] | bf_df["int64_col"]
34+
bf_df["bool_and_bool"] = bf_df["bool_col"] | bf_df["bool_col"]
35+
snapshot.assert_match(bf_df.sql, "out.sql")
36+
37+
38+
def test_xor_op(scalar_types_df: bpd.DataFrame, snapshot):
39+
bf_df = scalar_types_df[["bool_col", "int64_col"]]
40+
41+
bf_df["int_and_int"] = bf_df["int64_col"] ^ bf_df["int64_col"]
42+
bf_df["bool_and_bool"] = bf_df["bool_col"] ^ bf_df["bool_col"]
43+
snapshot.assert_match(bf_df.sql, "out.sql")

0 commit comments

Comments
 (0)