File tree 1 file changed +10
-5
lines changed
1 file changed +10
-5
lines changed Original file line number Diff line number Diff line change @@ -79,12 +79,17 @@ def prune_selection_child(
79
79
elif isinstance (child , bigframes .core .nodes .AdditiveNode ):
80
80
if not set (field .id for field in child .added_fields ) & selection .consumed_ids :
81
81
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
+ ),
87
90
)
91
+ return selection .replace_child (
92
+ child .replace_additive_base (prune_node (child .additive_base , needed_ids ))
88
93
)
89
94
elif isinstance (child , bigframes .core .nodes .ConcatNode ):
90
95
indices = [
You can’t perform that action at this time.
0 commit comments