fix: Fix row count local execution bug

Author: TrevorBergeron

bigframes/core/rewrite/pruning.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import dataclasses
 import functools
+import itertools
 import typing
 
 from bigframes.core import identifiers, nodes
@@ -51,24 +52,22 @@ def prune_columns(node: nodes.BigFrameNode):
     if isinstance(node, nodes.SelectionNode):
         result = prune_selection_child(node)
     elif isinstance(node, nodes.ResultNode):
-        result = node.replace_child(
-            prune_node(
-                node.child, node.consumed_ids or frozenset(list(node.child.ids)[0:1])
-            )
-        )
+        result = node.replace_child(prune_node(node.child, node.consumed_ids))
     elif isinstance(node, nodes.AggregateNode):
-        result = node.replace_child(
-            prune_node(
-                node.child, node.consumed_ids or frozenset(list(node.child.ids)[0:1])
-            )
-        )
+        result = node.replace_child(prune_node(node.child, node.consumed_ids))
     elif isinstance(node, nodes.InNode):
         result = dataclasses.replace(
             node,
             right_child=prune_node(node.right_child, frozenset([node.right_col.id])),
         )
     else:
         result = node
+
+    if len(set(result.ids)) == 0:
+        raise ValueError()
+    for child in result.child_nodes:
+        if len(set(child.ids)) == 0:
+            raise ValueError()
     return result
 
 
@@ -149,9 +148,13 @@ def prune_node(
     if not (set(node.ids) - ids):
         return node
     else:
+        # If no child ids are needed, probably a size op or numbering op above, keep a single column always
+        ids_to_keep = tuple(id for id in node.ids if id in ids) or tuple(
+            itertools.islice(node.ids, 0, 1)
+        )
         return nodes.SelectionNode(
             node,
-            tuple(nodes.AliasedRef.identity(id) for id in node.ids if id in ids),
+            tuple(nodes.AliasedRef.identity(id) for id in ids_to_keep),
         )
 
 

tests/system/small/engines/test_aggregation.py

@@ -48,6 +48,25 @@ def apply_agg_to_all_valid(
     return new_arr
 
 
+@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
+def test_engines_aggregate_post_filter_size(
+    scalars_array_value: array_value.ArrayValue,
+    engine,
+):
+    w_offsets, offsets_id = (
+        scalars_array_value.select_columns(("bool_col", "string_col"))
+        .filter(expression.deref("bool_col"))
+        .promote_offsets()
+    )
+    plan = (
+        w_offsets.select_columns((offsets_id, "bool_col", "string_col"))
+        .row_count()
+        .node
+    )
+
+    assert_equivalence_execution(plan, REFERENCE_ENGINE, engine)
+
+
 @pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
 def test_engines_aggregate_size(
     scalars_array_value: array_value.ArrayValue,