refactor: Share descendent logic between old and new aligner (#1199)

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -446,7 +446,7 @@ def try_row_join(
         other_node, r_mapping = self.prepare_join_names(other)
         import bigframes.core.rewrite
 
-        result_node = bigframes.core.rewrite.try_join_as_projection(
+        result_node = bigframes.core.rewrite.try_row_join(
             self.node, other_node, conditions
         )
         if result_node is None:

bigframes/core/nodes.py

@@ -83,10 +83,6 @@ def child_nodes(self) -> typing.Sequence[BigFrameNode]:
         """Direct children of this node"""
         return tuple([])
 
-    @property
-    def projection_base(self) -> BigFrameNode:
-        return self
-
     @property
     @abc.abstractmethod
     def row_count(self) -> typing.Optional[int]:
@@ -918,10 +914,6 @@ def row_count(self) -> Optional[int]:
     def node_defined_ids(self) -> Tuple[bfet_ids.ColumnId, ...]:
         return (self.col_id,)
 
-    @property
-    def projection_base(self) -> BigFrameNode:
-        return self.child.projection_base
-
     @property
     def added_fields(self) -> Tuple[Field, ...]:
         return (Field(self.col_id, bigframes.dtypes.INT_DTYPE),)
@@ -1095,10 +1087,6 @@ def variables_introduced(self) -> int:
     def defines_namespace(self) -> bool:
         return True
 
-    @property
-    def projection_base(self) -> BigFrameNode:
-        return self.child.projection_base
-
     @property
     def row_count(self) -> Optional[int]:
         return self.child.row_count
@@ -1173,10 +1161,6 @@ def variables_introduced(self) -> int:
     def row_count(self) -> Optional[int]:
         return self.child.row_count
 
-    @property
-    def projection_base(self) -> BigFrameNode:
-        return self.child.projection_base
-
     @property
     def node_defined_ids(self) -> Tuple[bfet_ids.ColumnId, ...]:
         return tuple(id for _, id in self.assignments)
@@ -1361,10 +1345,6 @@ def fields(self) -> Iterable[Field]:
     def variables_introduced(self) -> int:
         return 1
 
-    @property
-    def projection_base(self) -> BigFrameNode:
-        return self.child.projection_base
-
     @property
     def added_fields(self) -> Tuple[Field, ...]:
         return (self.added_field,)

bigframes/core/rewrite/__init__.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 
 from bigframes.core.rewrite.identifiers import remap_variables
-from bigframes.core.rewrite.implicit_align import try_join_as_projection
+from bigframes.core.rewrite.implicit_align import try_row_join
 from bigframes.core.rewrite.legacy_align import legacy_join_as_projection
 from bigframes.core.rewrite.slices import pullup_limit_from_slice, replace_slice_ops
 
 __all__ = [
     "legacy_join_as_projection",
-    "try_join_as_projection",
+    "try_row_join",
     "replace_slice_ops",
     "pullup_limit_from_slice",
     "remap_variables",

bigframes/core/rewrite/implicit_align.py

@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import Optional, Tuple
+from typing import Iterable, Optional, Tuple
 
 import bigframes.core.expression
 import bigframes.core.guid
@@ -30,6 +30,11 @@
     bigframes.core.nodes.WindowOpNode,
     bigframes.core.nodes.PromoteOffsetsNode,
 )
+# Combination of selects and additive nodes can be merged as an explicit keyless "row join"
+ALIGNABLE_NODES = (
+    *ADDITIVE_NODES,
+    bigframes.core.nodes.SelectionNode,
+)
 
 
 @dataclasses.dataclass(frozen=True)
@@ -70,49 +75,29 @@ def get_expression_spec(
                 bigframes.core.nodes.PromoteOffsetsNode,
             ),
         ):
-            # we don't yet have a way of normalizing window ops into a ExpressionSpec, which only
-            # handles normalizing scalar expressions at the moment.
-            pass
+            if set(expression.column_references).isdisjoint(
+                field.id for field in curr_node.added_fields
+            ):
+                # we don't yet have a way of normalizing window ops into a ExpressionSpec, which only
+                # handles normalizing scalar expressions at the moment.
+                pass
+            else:
+                return ExpressionSpec(expression, curr_node)
         else:
             return ExpressionSpec(expression, curr_node)
         curr_node = curr_node.child
 
 
-def _linearize_trees(
-    base_tree: bigframes.core.nodes.BigFrameNode,
-    append_tree: bigframes.core.nodes.BigFrameNode,
-) -> bigframes.core.nodes.BigFrameNode:
-    """Linearize two divergent tree who only diverge through different additive nodes."""
-    assert append_tree.projection_base == base_tree.projection_base
-    # base case: append tree does not have any additive nodes to linearize
-    if append_tree == append_tree.projection_base:
-        return base_tree
-    else:
-        assert isinstance(append_tree, ADDITIVE_NODES)
-        return append_tree.replace_child(_linearize_trees(base_tree, append_tree.child))
-
-
-def combine_nodes(
-    l_node: bigframes.core.nodes.BigFrameNode,
-    r_node: bigframes.core.nodes.BigFrameNode,
-) -> bigframes.core.nodes.BigFrameNode:
-    assert l_node.projection_base == r_node.projection_base
-    l_node, l_selection = pull_up_selection(l_node)
-    r_node, r_selection = pull_up_selection(
-        r_node, rename_vars=True
-    )  # Rename only right vars to avoid collisions with left vars
-    combined_selection = (*l_selection, *r_selection)
-    merged_node = _linearize_trees(l_node, r_node)
-    return bigframes.core.nodes.SelectionNode(merged_node, combined_selection)
-
-
-def try_join_as_projection(
+def try_row_join(
     l_node: bigframes.core.nodes.BigFrameNode,
     r_node: bigframes.core.nodes.BigFrameNode,
     join_keys: Tuple[Tuple[str, str], ...],
 ) -> Optional[bigframes.core.nodes.BigFrameNode]:
     """Joins the two nodes"""
-    if l_node.projection_base != r_node.projection_base:
+    divergent_node = first_shared_descendent(
+        l_node, r_node, descendable_types=ALIGNABLE_NODES
+    )
+    if divergent_node is None:
         return None
     # check join keys are equivalent by normalizing the expressions as much as posisble
     # instead of just comparing ids
@@ -124,11 +109,35 @@ def try_join_as_projection(
             r_node, right_id
         ):
             return None
-    return combine_nodes(l_node, r_node)
+
+    l_node, l_selection = pull_up_selection(l_node, stop=divergent_node)
+    r_node, r_selection = pull_up_selection(
+        r_node, stop=divergent_node, rename_vars=True
+    )  # Rename only right vars to avoid collisions with left vars
+    combined_selection = (*l_selection, *r_selection)
+
+    def _linearize_trees(
+        base_tree: bigframes.core.nodes.BigFrameNode,
+        append_tree: bigframes.core.nodes.BigFrameNode,
+    ) -> bigframes.core.nodes.BigFrameNode:
+        """Linearize two divergent tree who only diverge through different additive nodes."""
+        # base case: append tree does not have any divergent nodes to linearize
+        if append_tree == divergent_node:
+            return base_tree
+        else:
+            assert isinstance(append_tree, ADDITIVE_NODES)
+            return append_tree.replace_child(
+                _linearize_trees(base_tree, append_tree.child)
+            )
+
+    merged_node = _linearize_trees(l_node, r_node)
+    return bigframes.core.nodes.SelectionNode(merged_node, combined_selection)
 
 
 def pull_up_selection(
-    node: bigframes.core.nodes.BigFrameNode, rename_vars: bool = False
+    node: bigframes.core.nodes.BigFrameNode,
+    stop: bigframes.core.nodes.BigFrameNode,
+    rename_vars: bool = False,
 ) -> Tuple[
     bigframes.core.nodes.BigFrameNode,
     Tuple[
@@ -147,14 +156,14 @@ def pull_up_selection(
     Returns:
         BigFrameNode, Selections
     """
-    if node == node.projection_base:  # base case
+    if node == stop:  # base case
         return node, tuple(
             (bigframes.core.expression.DerefOp(field.id), field.id)
             for field in node.fields
         )
     assert isinstance(node, (bigframes.core.nodes.SelectionNode, *ADDITIVE_NODES))
     child_node, child_selections = pull_up_selection(
-        node.child, rename_vars=rename_vars
+        node.child, stop, rename_vars=rename_vars
     )
     mapping = {out: ref.id for ref, out in child_selections}
     if isinstance(node, ADDITIVE_NODES):
@@ -188,3 +197,30 @@ def pull_up_selection(
         )
         return child_node, new_selection
     raise ValueError(f"Couldn't pull up select from node: {node}")
+
+
+## Traversal helpers
+def first_shared_descendent(
+    left: bigframes.core.nodes.BigFrameNode,
+    right: bigframes.core.nodes.BigFrameNode,
+    descendable_types: Tuple[type[bigframes.core.nodes.UnaryNode], ...],
+) -> Optional[bigframes.core.nodes.BigFrameNode]:
+    l_path = tuple(descend(left, descendable_types))
+    r_path = tuple(descend(right, descendable_types))
+    if l_path[-1] != r_path[-1]:
+        return None
+
+    for l_node, r_node in zip(l_path[-len(r_path) :], r_path[-len(l_path) :]):
+        if l_node == r_node:
+            return l_node
+    # should be impossible, as l_path[-1] == r_path[-1]
+    raise ValueError()
+
+
+def descend(
+    root: bigframes.core.nodes.BigFrameNode,
+    descendable_types: Tuple[type[bigframes.core.nodes.UnaryNode], ...],
+) -> Iterable[bigframes.core.nodes.BigFrameNode]:
+    yield root
+    if isinstance(root, descendable_types):
+        yield from descend(root.child, descendable_types)

bigframes/core/rewrite/legacy_align.py

@@ -23,16 +23,17 @@
 import bigframes.core.join_def as join_defs
 import bigframes.core.nodes as nodes
 import bigframes.core.ordering as order
+import bigframes.core.rewrite.implicit_align
 import bigframes.operations as ops
 
 Selection = Tuple[Tuple[scalar_exprs.Expression, ids.ColumnId], ...]
 
-REWRITABLE_NODE_TYPES = (
-    nodes.SelectionNode,
-    nodes.ProjectionNode,
-    nodes.FilterNode,
-    nodes.ReversedNode,
-    nodes.OrderByNode,
+LEGACY_REWRITER_NODES = (
+    bigframes.core.nodes.ProjectionNode,
+    bigframes.core.nodes.SelectionNode,
+    bigframes.core.nodes.ReversedNode,
+    bigframes.core.nodes.OrderByNode,
+    bigframes.core.nodes.FilterNode,
 )
 
 
@@ -51,9 +52,7 @@ def from_node_span(
         cls, node: nodes.BigFrameNode, target: nodes.BigFrameNode
     ) -> SquashedSelect:
         if node == target:
-            selection = tuple(
-                (scalar_exprs.DerefOp(id), id) for id in get_node_column_ids(node)
-            )
+            selection = tuple((scalar_exprs.DerefOp(id), id) for id in node.ids)
             return cls(node, selection, None, ())
 
         if isinstance(node, nodes.SelectionNode):
@@ -357,27 +356,10 @@ def decompose_conjunction(
         return (expr,)
 
 
-def get_node_column_ids(node: nodes.BigFrameNode) -> Tuple[ids.ColumnId, ...]:
-    return tuple(field.id for field in node.fields)
-
-
 def common_selection_root(
     l_tree: nodes.BigFrameNode, r_tree: nodes.BigFrameNode
 ) -> Optional[nodes.BigFrameNode]:
     """Find common subtree between join subtrees"""
-    l_node = l_tree
-    l_nodes: set[nodes.BigFrameNode] = set()
-    while isinstance(l_node, REWRITABLE_NODE_TYPES):
-        l_nodes.add(l_node)
-        l_node = l_node.child
-    l_nodes.add(l_node)
-
-    r_node = r_tree
-    while isinstance(r_node, REWRITABLE_NODE_TYPES):
-        if r_node in l_nodes:
-            return r_node
-        r_node = r_node.child
-
-    if r_node in l_nodes:
-        return r_node
-    return None
+    return bigframes.core.rewrite.implicit_align.first_shared_descendent(
+        l_tree, r_tree, descendable_types=LEGACY_REWRITER_NODES
+    )