refactor: fix remap variable error in InNode

Author: chelsea-lin

bigframes/core/nodes.py

@@ -1681,7 +1681,7 @@ class ResultNode(UnaryNode):
     # TODO: CTE definitions
 
     def _validate(self):
-        for ref, name in self.output_cols:
+        for ref, _ in self.output_cols:
             assert ref.id in self.child.ids
 
     @property

bigframes/core/rewrite/identifiers.py

@@ -57,8 +57,10 @@ def remap_variables(
     new_root = root.transform_children(lambda node: remapped_children[node])
 
     # Step 3: Transform the current node using the mappings from its children.
+    # "reversed" is required for InNode so that in case of a duplicate column ID,
+    # the left child's mapping is the one that's kept.
     downstream_mappings: dict[identifiers.ColumnId, identifiers.ColumnId] = {
-        k: v for mapping in new_child_mappings for k, v in mapping.items()
+        k: v for mapping in reversed(new_child_mappings) for k, v in mapping.items()
     }
     if isinstance(new_root, nodes.InNode):
         new_root = typing.cast(nodes.InNode, new_root)

tests/system/small/engines/test_join.py

@@ -55,7 +55,7 @@ def test_engines_join_on_coerced_key(
     assert_equivalence_execution(result.node, REFERENCE_ENGINE, engine)
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
 @pytest.mark.parametrize("join_type", ["left", "inner", "right", "outer"])
 def test_engines_join_multi_key(
     scalars_array_value: array_value.ArrayValue,
@@ -90,7 +90,7 @@ def test_engines_cross_join(
     assert_equivalence_execution(result.node, REFERENCE_ENGINE, engine)
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
 @pytest.mark.parametrize(
     ("left_key", "right_key"),
     [

tests/system/small/engines/test_slicing.py

@@ -24,7 +24,7 @@
 REFERENCE_ENGINE = polars_executor.PolarsExecutor()
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
 @pytest.mark.parametrize(
     ("start", "stop", "step"),
     [

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

@@ -0,0 +1,35 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`,
+    `float64_col` AS `bfcol_1`,
+    `rowindex` AS `bfcol_2`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    `bfcol_2` AS `bfcol_3`,
+    `bfcol_0` AS `bfcol_4`,
+    `bfcol_1` AS `bfcol_5`
+  FROM `bfcte_0`
+), `bfcte_2` AS (
+  SELECT
+    `bfcte_1`.*,
+    EXISTS(
+      SELECT
+        1
+      FROM (
+        SELECT
+          `float64_col` AS `bfcol_6`
+        FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+      ) AS `bft_0`
+      WHERE
+        COALESCE(`bfcte_1`.`bfcol_4`, 0) = COALESCE(`bft_0`.`bfcol_6`, 0)
+        AND COALESCE(`bfcte_1`.`bfcol_4`, 1) = COALESCE(`bft_0`.`bfcol_6`, 1)
+    ) AS `bfcol_7`
+  FROM `bfcte_1`
+)
+SELECT
+  `bfcol_3` AS `bfuid_col_1`,
+  `bfcol_4` AS `int64_col`,
+  `bfcol_5` AS `float64_col`,
+  `bfcol_7` AS `bfuid_col_2`
+FROM `bfcte_2`

tests/unit/core/compile/sqlglot/test_compile_isin.py

@@ -37,3 +37,13 @@ def test_compile_isin_not_nullable(scalar_types_df: bpd.DataFrame, snapshot):
         scalar_types_df["rowindex_2"].isin(scalar_types_df["rowindex_2"]).to_frame()
     )
     snapshot.assert_match(bf_isin.sql, "out.sql")
+
+
+def test_compile_isin_for_array_value(scalar_types_df: bpd.DataFrame, snapshot):
+    scalars_array_value = scalar_types_df[["int64_col", "float64_col"]]._block.expr
+    result, _ = scalars_array_value.isin(
+        scalars_array_value, lcol="int64_col", rcol="float64_col"
+    )
+    sql = result.session._executor.to_sql(result, enable_cache=False)
+
+    snapshot.assert_match(sql, "out.sql")