Skip to content

Commit 7990262

Browse files
perf: Prune projections more aggressively (#1398)
1 parent a914a5f commit 7990262

File tree

1 file changed

+10
-5
lines changed

1 file changed

+10
-5
lines changed

bigframes/core/rewrite/pruning.py

+10-5
Original file line numberDiff line numberDiff line change
@@ -79,12 +79,17 @@ def prune_selection_child(
7979
elif isinstance(child, bigframes.core.nodes.AdditiveNode):
8080
if not set(field.id for field in child.added_fields) & selection.consumed_ids:
8181
return selection.replace_child(child.additive_base)
82-
return selection.replace_child(
83-
child.replace_additive_base(
84-
prune_node(
85-
child.additive_base, selection.consumed_ids | child.referenced_ids
86-
)
82+
needed_ids = selection.consumed_ids | child.referenced_ids
83+
if isinstance(child, bigframes.core.nodes.ProjectionNode):
84+
# Projection expressions are independent, so can be individually removed from the node
85+
child = dataclasses.replace(
86+
child,
87+
assignments=tuple(
88+
(ex, id) for (ex, id) in child.assignments if id in needed_ids
89+
),
8790
)
91+
return selection.replace_child(
92+
child.replace_additive_base(prune_node(child.additive_base, needed_ids))
8893
)
8994
elif isinstance(child, bigframes.core.nodes.ConcatNode):
9095
indices = [

0 commit comments

Comments
 (0)