Skip to content

Commit

Permalink
perf: Prune projections more aggressively (#1398)
Browse files Browse the repository at this point in the history
  • Loading branch information
TrevorBergeron authored Feb 14, 2025
1 parent a914a5f commit 7990262
Showing 1 changed file with 10 additions and 5 deletions.
15 changes: 10 additions & 5 deletions bigframes/core/rewrite/pruning.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down

0 comments on commit 7990262

Please sign in to comment.