|
| 1 | +# Copyright 2025 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +import dataclasses |
| 16 | +from typing import cast |
| 17 | + |
| 18 | +from bigframes.core import expression, nodes |
| 19 | + |
| 20 | + |
| 21 | +def defer_selection( |
| 22 | + root: nodes.BigFrameNode, |
| 23 | +) -> nodes.BigFrameNode: |
| 24 | + """ |
| 25 | + Defers SelectionNode operations in the tree, pulling them up. |
| 26 | +
|
| 27 | + In many cases, these nodes will be merged or eliminated entirely, simplifying the overall tree. |
| 28 | + """ |
| 29 | + return nodes.bottom_up(root, pull_up_select) |
| 30 | + |
| 31 | + |
| 32 | +def pull_up_select(node: nodes.BigFrameNode) -> nodes.BigFrameNode: |
| 33 | + if isinstance(node, nodes.LeafNode): |
| 34 | + return node |
| 35 | + if isinstance(node, nodes.JoinNode): |
| 36 | + return pull_up_selects_under_join(node) |
| 37 | + if isinstance(node, nodes.ConcatNode): |
| 38 | + return handle_selects_under_concat(node) |
| 39 | + if isinstance(node, nodes.UnaryNode): |
| 40 | + return pull_up_select_unary(node) |
| 41 | + # shouldn't hit this, but not worth crashing over |
| 42 | + return node |
| 43 | + |
| 44 | + |
| 45 | +def pull_up_select_unary(node: nodes.UnaryNode) -> nodes.BigFrameNode: |
| 46 | + child = node.child |
| 47 | + if not isinstance(child, nodes.SelectionNode): |
| 48 | + return node |
| 49 | + |
| 50 | + # Schema-preserving nodes |
| 51 | + if isinstance( |
| 52 | + node, |
| 53 | + ( |
| 54 | + nodes.ReversedNode, |
| 55 | + nodes.OrderByNode, |
| 56 | + nodes.SliceNode, |
| 57 | + nodes.FilterNode, |
| 58 | + nodes.RandomSampleNode, |
| 59 | + ), |
| 60 | + ): |
| 61 | + pushed_down_node: nodes.BigFrameNode = node.remap_refs( |
| 62 | + {id: ref.id for ref, id in child.input_output_pairs} |
| 63 | + ).replace_child(child.child) |
| 64 | + pulled_up_select = cast( |
| 65 | + nodes.SelectionNode, child.replace_child(pushed_down_node) |
| 66 | + ) |
| 67 | + return pulled_up_select |
| 68 | + elif isinstance( |
| 69 | + node, |
| 70 | + ( |
| 71 | + nodes.SelectionNode, |
| 72 | + nodes.ResultNode, |
| 73 | + ), |
| 74 | + ): |
| 75 | + return node.remap_refs( |
| 76 | + {id: ref.id for ref, id in child.input_output_pairs} |
| 77 | + ).replace_child(child.child) |
| 78 | + elif isinstance(node, nodes.AggregateNode): |
| 79 | + pushed_down_agg = node.remap_refs( |
| 80 | + {id: ref.id for ref, id in child.input_output_pairs} |
| 81 | + ).replace_child(child.child) |
| 82 | + new_selection = tuple( |
| 83 | + nodes.AliasedRef.identity(id).remap_refs( |
| 84 | + {id: ref.id for ref, id in child.input_output_pairs} |
| 85 | + ) |
| 86 | + for id in node.ids |
| 87 | + ) |
| 88 | + return nodes.SelectionNode(pushed_down_agg, new_selection) |
| 89 | + elif isinstance(node, nodes.ExplodeNode): |
| 90 | + pushed_down_node = node.remap_refs( |
| 91 | + {id: ref.id for ref, id in child.input_output_pairs} |
| 92 | + ).replace_child(child.child) |
| 93 | + pulled_up_select = cast( |
| 94 | + nodes.SelectionNode, child.replace_child(pushed_down_node) |
| 95 | + ) |
| 96 | + if node.offsets_col: |
| 97 | + pulled_up_select = dataclasses.replace( |
| 98 | + pulled_up_select, |
| 99 | + input_output_pairs=( |
| 100 | + *pulled_up_select.input_output_pairs, |
| 101 | + nodes.AliasedRef( |
| 102 | + expression.DerefOp(node.offsets_col), node.offsets_col |
| 103 | + ), |
| 104 | + ), |
| 105 | + ) |
| 106 | + return pulled_up_select |
| 107 | + elif isinstance(node, nodes.AdditiveNode): |
| 108 | + pushed_down_node = node.replace_additive_base(child.child).remap_refs( |
| 109 | + {id: ref.id for ref, id in child.input_output_pairs} |
| 110 | + ) |
| 111 | + new_selection = ( |
| 112 | + *child.input_output_pairs, |
| 113 | + *( |
| 114 | + nodes.AliasedRef(expression.DerefOp(col.id), col.id) |
| 115 | + for col in node.added_fields |
| 116 | + ), |
| 117 | + ) |
| 118 | + pulled_up_select = dataclasses.replace( |
| 119 | + child, child=pushed_down_node, input_output_pairs=new_selection |
| 120 | + ) |
| 121 | + return pulled_up_select |
| 122 | + # shouldn't hit this, but not worth crashing over |
| 123 | + return node |
| 124 | + |
| 125 | + |
| 126 | +def pull_up_selects_under_join(node: nodes.JoinNode) -> nodes.JoinNode: |
| 127 | + # Can in theory pull up selects here, but it is a bit dangerous, in particular or self-joins, when there are more transforms to do. |
| 128 | + # TODO: Safely pull up selects above join |
| 129 | + return node |
| 130 | + |
| 131 | + |
| 132 | +def handle_selects_under_concat(node: nodes.ConcatNode) -> nodes.ConcatNode: |
| 133 | + new_children = [] |
| 134 | + for child in node.child_nodes: |
| 135 | + # remove select if no-op |
| 136 | + if not isinstance(child, nodes.SelectionNode): |
| 137 | + new_children.append(child) |
| 138 | + else: |
| 139 | + inputs = (ref.id for ref in child.input_output_pairs) |
| 140 | + if inputs == tuple(child.child.ids): |
| 141 | + new_children.append(child.child) |
| 142 | + else: |
| 143 | + new_children.append(child) |
| 144 | + return dataclasses.replace(node, children=tuple(new_children)) |
0 commit comments