From 4f79f409e6271cf6609a48186e997bac4973088c Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Thu, 13 Feb 2025 22:29:45 +0000 Subject: [PATCH] perf: Prune projections more aggressively --- bigframes/core/rewrite/pruning.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/bigframes/core/rewrite/pruning.py b/bigframes/core/rewrite/pruning.py index 7e40137f3e..5a94f2aa40 100644 --- a/bigframes/core/rewrite/pruning.py +++ b/bigframes/core/rewrite/pruning.py @@ -79,12 +79,17 @@ def prune_selection_child( elif isinstance(child, bigframes.core.nodes.AdditiveNode): if not set(field.id for field in child.added_fields) & selection.consumed_ids: return selection.replace_child(child.additive_base) - return selection.replace_child( - child.replace_additive_base( - prune_node( - child.additive_base, selection.consumed_ids | child.referenced_ids - ) + needed_ids = selection.consumed_ids | child.referenced_ids + if isinstance(child, bigframes.core.nodes.ProjectionNode): + # Projection expressions are independent, so can be individually removed from the node + child = dataclasses.replace( + child, + assignments=tuple( + (ex, id) for (ex, id) in child.assignments if id in needed_ids + ), ) + return selection.replace_child( + child.replace_additive_base(prune_node(child.additive_base, needed_ids)) ) elif isinstance(child, bigframes.core.nodes.ConcatNode): indices = [