For isin, use subquery rather than cte for the right selections

Author: chelsea-lin

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -347,12 +347,9 @@ def isin_join(
         left_cte_name = sge.to_identifier(
             next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
         )
-        right_cte_name = sge.to_identifier(
-            next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
-        )
 
         left_select = _select_to_cte(self.expr, left_cte_name)
-        right_select = _select_to_cte(right.expr, right_cte_name)
+        right_select = right.expr
 
         left_ctes = left_select.args.pop("with", [])
         right_ctes = right_select.args.pop("with", [])
@@ -362,25 +359,28 @@ def isin_join(
             sge.Column(this=conditions[0].expr, table=left_cte_name),
             conditions[0].dtype,
         )
-        right_condition = typed_expr.TypedExpr(
-            sge.Column(this=conditions[1].expr, table=right_cte_name),
-            conditions[1].dtype,
-        )
 
         new_column: sge.Expression
         if joins_nulls:
+            right_table_name = sge.to_identifier(
+                next(self.uid_gen.get_uid_stream("bft_")), quoted=self.quoted
+            )
+            right_condition = typed_expr.TypedExpr(
+                sge.Column(this=conditions[1].expr, table=right_table_name),
+                conditions[1].dtype,
+            )
             new_column = sge.Exists(
                 this=sge.Select()
                 .select(sge.convert(1))
-                .from_(sge.Table(this=right_cte_name))
+                .from_(sge.Alias(this=right_select.subquery(), alias=right_table_name))
                 .where(
                     _join_condition(left_condition, right_condition, joins_nulls=True)
                 )
             )
         else:
             new_column = sge.In(
                 this=left_condition.expr,
-                expressions=[right_condition.expr],
+                expressions=[right_select.subquery()],
             )
 
         new_column = sge.Alias(

tests/unit/core/compile/sqlglot/snapshots/test_compile_isin/test_compile_isin/out.sql

@@ -13,25 +13,25 @@ WITH `bfcte_1` AS (
     `int64_too` AS `bfcol_4`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_3` AS (
-  SELECT
-    `bfcol_4`
-  FROM `bfcte_0`
-  GROUP BY
-    `bfcol_4`
-), `bfcte_4` AS (
   SELECT
     `bfcte_2`.*,
     EXISTS(
       SELECT
         1
-      FROM `bfcte_3`
+      FROM (
+        SELECT
+          `bfcol_4`
+        FROM `bfcte_0`
+        GROUP BY
+          `bfcol_4`
+      ) AS `bft_0`
       WHERE
-        COALESCE(`bfcte_2`.`bfcol_3`, 0) = COALESCE(`bfcte_3`.`bfcol_4`, 0)
-        AND COALESCE(`bfcte_2`.`bfcol_3`, 1) = COALESCE(`bfcte_3`.`bfcol_4`, 1)
+        COALESCE(`bfcte_2`.`bfcol_3`, 0) = COALESCE(`bft_0`.`bfcol_4`, 0)
+        AND COALESCE(`bfcte_2`.`bfcol_3`, 1) = COALESCE(`bft_0`.`bfcol_4`, 1)
     ) AS `bfcol_5`
   FROM `bfcte_2`
 )
 SELECT
   `bfcol_2` AS `rowindex`,
   `bfcol_5` AS `int64_col`
-FROM `bfcte_4`
+FROM `bfcte_3`

tests/unit/core/compile/sqlglot/snapshots/test_compile_isin/test_compile_isin_not_nullable/out.sql

@@ -13,18 +13,18 @@ WITH `bfcte_1` AS (
     `rowindex_2` AS `bfcol_4`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_3` AS (
-  SELECT
-    `bfcol_4`
-  FROM `bfcte_0`
-  GROUP BY
-    `bfcol_4`
-), `bfcte_4` AS (
   SELECT
     `bfcte_2`.*,
-    `bfcte_2`.`bfcol_3` IN (`bfcte_3`.`bfcol_4`) AS `bfcol_5`
+    `bfcte_2`.`bfcol_3` IN ((
+        SELECT
+          `bfcol_4`
+        FROM `bfcte_0`
+        GROUP BY
+          `bfcol_4`
+    )) AS `bfcol_5`
   FROM `bfcte_2`
 )
 SELECT
   `bfcol_2` AS `rowindex`,
   `bfcol_5` AS `rowindex_2`
-FROM `bfcte_4`
+FROM `bfcte_3`