refactor: Create memoized tree traversal helpers (#1198)

Author: TrevorBergeron

bigframes/core/compile/compiler.py

@@ -82,7 +82,7 @@ def _preprocess(self, node: nodes.BigFrameNode):
         if self.enable_pruning:
             used_fields = frozenset(field.id for field in node.fields)
             node = node.prune(used_fields)
-        node = functools.cache(rewrites.replace_slice_ops)(node)
+        node = nodes.bottom_up(node, rewrites.rewrite_slice)
         if self.enable_densify_ids:
             original_names = [id.name for id in node.ids]
             node, _ = rewrites.remap_variables(

bigframes/core/compile/polars/compiler.py

@@ -173,7 +173,7 @@ def compile(self, array_value: bigframes.core.ArrayValue) -> pl.LazyFrame:
 
         # TODO: Create standard way to configure BFET -> BFET rewrites
         # Polars has incomplete slice support in lazy mode
-        node = bigframes.core.rewrite.replace_slice_ops(array_value.node)
+        node = nodes.bottom_up(array_value.node, bigframes.core.rewrite.rewrite_slice)
         return self.compile_node(node)
 
     @functools.singledispatchmethod

bigframes/core/nodes.py

@@ -1486,3 +1486,56 @@ def remap_refs(
     ) -> BigFrameNode:
         new_ids = tuple(id.remap_column_refs(mappings) for id in self.column_ids)
         return dataclasses.replace(self, column_ids=new_ids)  # type: ignore
+
+
+# Tree operators
+def top_down(
+    root: BigFrameNode,
+    transform: Callable[[BigFrameNode], BigFrameNode],
+    *,
+    memoize=False,
+    validate=False,
+) -> BigFrameNode:
+    """
+    Perform a top-down transformation of the BigFrameNode tree.
+
+    If memoize=True, recursive calls are memoized within the scope of the traversal only.
+    """
+
+    def top_down_internal(root: BigFrameNode) -> BigFrameNode:
+        return transform(root).transform_children(top_down_internal)
+
+    if memoize:
+        # MUST reassign to the same name or caching won't work recursively
+        top_down_internal = functools.cache(top_down_internal)
+
+    result = top_down_internal(root)
+    if validate:
+        result.validate_tree()
+    return result
+
+
+def bottom_up(
+    root: BigFrameNode,
+    transform: Callable[[BigFrameNode], BigFrameNode],
+    *,
+    memoize=False,
+    validate=False,
+) -> BigFrameNode:
+    """
+    Perform a bottom-up transformation of the BigFrameNode tree.
+
+    If memoize=True, recursive calls are memoized within the scope of the traversal only.
+    """
+
+    def bottom_up_internal(root: BigFrameNode) -> BigFrameNode:
+        return transform(root.transform_children(bottom_up_internal))
+
+    if memoize:
+        # MUST reassign to the same name or caching won't work recursively
+        bottom_up_internal = functools.cache(bottom_up_internal)
+
+    result = bottom_up_internal(root)
+    if validate:
+        result.validate_tree()
+    return result

bigframes/core/rewrite/__init__.py

@@ -15,12 +15,12 @@
 from bigframes.core.rewrite.identifiers import remap_variables
 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
+from bigframes.core.rewrite.slices import pullup_limit_from_slice, rewrite_slice
 
 __all__ = [
     "legacy_join_as_projection",
     "try_row_join",
-    "replace_slice_ops",
+    "rewrite_slice",
     "pullup_limit_from_slice",
     "remap_variables",
 ]

bigframes/core/rewrite/slices.py

@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 import functools
-from typing import cast, Optional, Sequence, Tuple
+from typing import Optional, Sequence, Tuple
 
 import bigframes.core.expression as scalar_exprs
 import bigframes.core.guid as guids
@@ -55,18 +55,11 @@ def pullup_limit_from_slice(
     return root, None
 
 
-def replace_slice_ops(root: nodes.BigFrameNode) -> nodes.BigFrameNode:
-    # TODO: we want to pull up some slices into limit op if near root.
-    if isinstance(root, nodes.SliceNode):
-        root = root.transform_children(replace_slice_ops)
-        return rewrite_slice(cast(nodes.SliceNode, root))
-    else:
-        return root.transform_children(replace_slice_ops)
+def rewrite_slice(node: nodes.BigFrameNode):
+    if not isinstance(node, nodes.SliceNode):
+        return node
 
-
-def rewrite_slice(node: nodes.SliceNode):
     slice_def = (node.start, node.stop, node.step)
-
     # no-op (eg. df[::1])
     if slices.is_noop(slice_def, node.child.row_count):
         return node.child

bigframes/core/tree_properties.py

@@ -88,7 +88,7 @@ def select_cache_target(
 
     @functools.cache
     def _with_caching(subtree: nodes.BigFrameNode) -> nodes.BigFrameNode:
-        return replace_nodes(subtree, cache)
+        return nodes.top_down(subtree, lambda x: cache.get(x, x), memoize=True)
 
     def _combine_counts(
         left: Dict[nodes.BigFrameNode, int], right: Dict[nodes.BigFrameNode, int]
@@ -162,17 +162,3 @@ def _node_counts_inner(
 
     counts = [_node_counts_inner(root) for root in forest]
     return functools.reduce(_combine_counts, counts, empty_counts)
-
-
-def replace_nodes(
-    root: nodes.BigFrameNode,
-    replacements: dict[nodes.BigFrameNode, nodes.BigFrameNode],
-):
-    @functools.cache
-    def apply_substition(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
-        if node in replacements.keys():
-            return replacements[node]
-        else:
-            return node.transform_children(apply_substition)
-
-    return apply_substition(root)

bigframes/session/executor.py

@@ -529,7 +529,9 @@ def _wait_on_job(
         return results_iterator
 
     def replace_cached_subtrees(self, node: nodes.BigFrameNode) -> nodes.BigFrameNode:
-        return tree_properties.replace_nodes(node, (dict(self._cached_executions)))
+        return nodes.top_down(
+            node, lambda x: self._cached_executions.get(x, x), memoize=True
+        )
 
     def _is_trivially_executable(self, array_value: bigframes.core.ArrayValue):
         """