From 53978a8230848954e521ff2d767a5c982133cae8 Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Mon, 3 Nov 2025 22:05:08 +0000 Subject: [PATCH] refactor: fix remap variable errors for InNode --- bigframes/core/nodes.py | 2 +- bigframes/core/rewrite/identifiers.py | 4 +++- tests/system/small/engines/test_join.py | 4 ++-- tests/system/small/engines/test_slicing.py | 2 +- 4 files changed, 7 insertions(+), 5 deletions(-) diff --git a/bigframes/core/nodes.py b/bigframes/core/nodes.py index 9e0fcb3ace..553b41a631 100644 --- a/bigframes/core/nodes.py +++ b/bigframes/core/nodes.py @@ -1627,7 +1627,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 diff --git a/bigframes/core/rewrite/identifiers.py b/bigframes/core/rewrite/identifiers.py index e911d81895..2e31f07a79 100644 --- a/bigframes/core/rewrite/identifiers.py +++ b/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) diff --git a/tests/system/small/engines/test_join.py b/tests/system/small/engines/test_join.py index 91c199a437..7ea24a554d 100644 --- a/tests/system/small/engines/test_join.py +++ b/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"), [ diff --git a/tests/system/small/engines/test_slicing.py b/tests/system/small/engines/test_slicing.py index 7340ff145b..022758893d 100644 --- a/tests/system/small/engines/test_slicing.py +++ b/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"), [