refactor: fix test_isin tests

Author: chelsea-lin

bigframes/core/compile/sqlglot/expressions/comparison_ops.py

@@ -31,17 +31,23 @@
 @register_unary_op(ops.IsInOp, pass_op=True)
 def _(expr: TypedExpr, op: ops.IsInOp) -> sge.Expression:
     values = []
-    is_numeric_expr = dtypes.is_numeric(expr.dtype)
+    is_numeric_expr = dtypes.is_numeric(expr.dtype, include_bool=False)
     for value in op.values:
-        if value is None:
+        if _is_null(value):
             continue
         dtype = dtypes.bigframes_type(type(value))
-        if expr.dtype == dtype or is_numeric_expr and dtypes.is_numeric(dtype):
+        if (
+            expr.dtype == dtype
+            or is_numeric_expr
+            and dtypes.is_numeric(dtype, include_bool=False)
+        ):
             values.append(sge.convert(value))
 
     if op.match_nulls:
         contains_nulls = any(_is_null(value) for value in op.values)
         if contains_nulls:
+            if len(values) == 0:
+                return sge.Is(this=expr.expr, expression=sge.Null())
             return sge.Is(this=expr.expr, expression=sge.Null()) | sge.In(
                 this=expr.expr, expressions=values
             )

tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_is_in/out.sql

@@ -1,32 +1,35 @@
 WITH `bfcte_0` AS (
   SELECT
+    `bool_col`,
     `float64_col`,
     `int64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
     *,
-    COALESCE(`int64_col` IN (1, 2, 3), FALSE) AS `bfcol_2`,
-    (
-      `int64_col` IS NULL
-    ) OR `int64_col` IN (123456) AS `bfcol_3`,
-    COALESCE(`int64_col` IN (1.0, 2.0, 3.0), FALSE) AS `bfcol_4`,
-    FALSE AS `bfcol_5`,
-    COALESCE(`int64_col` IN (2.5, 3), FALSE) AS `bfcol_6`,
+    COALESCE(`bool_col` IN (TRUE, FALSE), FALSE) AS `bfcol_3`,
+    COALESCE(`int64_col` IN (1, 2, 3), FALSE) AS `bfcol_4`,
+    `int64_col` IS NULL AS `bfcol_5`,
+    COALESCE(`int64_col` IN (1.0, 2.0, 3.0), FALSE) AS `bfcol_6`,
     FALSE AS `bfcol_7`,
-    COALESCE(`int64_col` IN (123456), FALSE) AS `bfcol_8`,
+    COALESCE(`int64_col` IN (2.5, 3), FALSE) AS `bfcol_8`,
+    FALSE AS `bfcol_9`,
+    FALSE AS `bfcol_10`,
+    COALESCE(`int64_col` IN (123456), FALSE) AS `bfcol_11`,
     (
       `float64_col` IS NULL
-    ) OR `float64_col` IN (1, 2, 3) AS `bfcol_9`
+    ) OR `float64_col` IN (1, 2, 3) AS `bfcol_12`
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_2` AS `ints`,
-  `bfcol_3` AS `ints_w_null`,
-  `bfcol_4` AS `floats`,
-  `bfcol_5` AS `strings`,
-  `bfcol_6` AS `mixed`,
-  `bfcol_7` AS `empty`,
-  `bfcol_8` AS `ints_wo_match_nulls`,
-  `bfcol_9` AS `float_in_ints`
+  `bfcol_3` AS `bools`,
+  `bfcol_4` AS `ints`,
+  `bfcol_5` AS `ints_w_null`,
+  `bfcol_6` AS `floats`,
+  `bfcol_7` AS `strings`,
+  `bfcol_8` AS `mixed`,
+  `bfcol_9` AS `empty`,
+  `bfcol_10` AS `empty_wo_match_nulls`,
+  `bfcol_11` AS `ints_wo_match_nulls`,
+  `bfcol_12` AS `float_in_ints`
 FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/test_comparison_ops.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pandas as pd
 import pytest
 
 from bigframes import operations as ops
@@ -22,18 +23,23 @@
 
 
 def test_is_in(scalar_types_df: bpd.DataFrame, snapshot):
+    bool_col = "bool_col"
     int_col = "int64_col"
     float_col = "float64_col"
-    bf_df = scalar_types_df[[int_col, float_col]]
+    bf_df = scalar_types_df[[bool_col, int_col, float_col]]
     ops_map = {
+        "bools": ops.IsInOp(values=(True, False)).as_expr(bool_col),
         "ints": ops.IsInOp(values=(1, 2, 3)).as_expr(int_col),
-        "ints_w_null": ops.IsInOp(values=(None, 123456)).as_expr(int_col),
+        "ints_w_null": ops.IsInOp(values=(None, pd.NA)).as_expr(int_col),
         "floats": ops.IsInOp(values=(1.0, 2.0, 3.0), match_nulls=False).as_expr(
             int_col
         ),
         "strings": ops.IsInOp(values=("1.0", "2.0")).as_expr(int_col),
         "mixed": ops.IsInOp(values=("1.0", 2.5, 3)).as_expr(int_col),
         "empty": ops.IsInOp(values=()).as_expr(int_col),
+        "empty_wo_match_nulls": ops.IsInOp(values=(), match_nulls=False).as_expr(
+            int_col
+        ),
         "ints_wo_match_nulls": ops.IsInOp(
             values=(None, 123456), match_nulls=False
         ).as_expr(int_col),