refactor: Ensure only valid IDs are propagated in identifier remapping

Corrected remap_variables to only propagate column IDs that are actually present in the current node's output fields. This prevents parent nodes from seeing internal or leaked column IDs from child nodes, which was specifically problematic for aggregate nodes.

Added unit tests to verify correct propagation for AggregateNode with and without grouping.

Author: chelsea-lin

bigframes/core/rewrite/identifiers.py

@@ -82,12 +82,10 @@ def remap_variables(
     new_root._validate()
 
     # Step 5: Determine which mappings to propagate up to the parent.
-    if root.defines_namespace:
-        # If a node defines a new namespace (e.g., a join), mappings from its
-        # children are not visible to its parents.
-        mappings_for_parent = node_defined_mappings
-    else:
-        # Otherwise, pass up the combined mappings from children and the current node.
-        mappings_for_parent = downstream_mappings | node_defined_mappings
+    propagated_mappings = {
+        old_id: new_id
+        for old_id, new_id in (downstream_mappings | node_defined_mappings).items()
+        if old_id in root.ids
+    }
 
-    return new_root, mappings_for_parent
+    return new_root, propagated_mappings

tests/unit/core/rewrite/test_identifiers.py

@@ -14,10 +14,12 @@
 import typing
 
 import bigframes.core as core
+import bigframes.core.agg_expressions as agg_ex
 import bigframes.core.expression as ex
 import bigframes.core.identifiers as identifiers
 import bigframes.core.nodes as nodes
 import bigframes.core.rewrite.identifiers as id_rewrite
+import bigframes.operations.aggregations as agg_ops
 
 
 def test_remap_variables_single_node(leaf):
@@ -51,6 +53,51 @@ def test_remap_variables_projection(leaf):
     assert set(mapping.values()) == {identifiers.ColumnId(f"id_{i}") for i in range(3)}
 
 
+def test_remap_variables_aggregate(leaf):
+    # Aggregation: sum(col_a) AS sum_a
+    # Group by nothing
+    agg_op = agg_ex.UnaryAggregation(
+        op=agg_ops.sum_op,
+        arg=ex.DerefOp(leaf.fields[0].id),
+    )
+    node = nodes.AggregateNode(
+        child=leaf,
+        aggregations=((agg_op, identifiers.ColumnId("sum_a")),),
+        by_column_ids=(),
+    )
+
+    id_generator = (identifiers.ColumnId(f"id_{i}") for i in range(100))
+    _, mapping = id_rewrite.remap_variables(node, id_generator)
+
+    # leaf has 2 columns: col_a, col_b
+    # AggregateNode defines 1 column: sum_a
+    # Output of AggregateNode should only be sum_a
+    assert len(mapping) == 1
+    assert identifiers.ColumnId("sum_a") in mapping
+
+
+def test_remap_variables_aggregate_with_grouping(leaf):
+    # Aggregation: sum(col_b) AS sum_b
+    # Group by col_a
+    agg_op = agg_ex.UnaryAggregation(
+        op=agg_ops.sum_op,
+        arg=ex.DerefOp(leaf.fields[1].id),
+    )
+    node = nodes.AggregateNode(
+        child=leaf,
+        aggregations=((agg_op, identifiers.ColumnId("sum_b")),),
+        by_column_ids=(ex.DerefOp(leaf.fields[0].id),),
+    )
+
+    id_generator = (identifiers.ColumnId(f"id_{i}") for i in range(100))
+    _, mapping = id_rewrite.remap_variables(node, id_generator)
+
+    # Output should have 2 columns: col_a (grouping) and sum_b (agg)
+    assert len(mapping) == 2
+    assert leaf.fields[0].id in mapping
+    assert identifiers.ColumnId("sum_b") in mapping
+
+
 def test_remap_variables_nested_join_stability(leaf, fake_session, table):
     # Create two more distinct leaf nodes
     leaf2_uncached = core.ArrayValue.from_table(