refactor: caching is now a session property (#697)

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -184,7 +184,7 @@ def project_to_id(self, expression: ex.Expression, output_id: str):
                 child=self.node,
                 assignments=tuple(exprs),
             )
-        ).merge_projections()
+        )
 
     def assign(self, source_id: str, destination_id: str) -> ArrayValue:
         if destination_id in self.column_ids:  # Mutate case
@@ -209,7 +209,7 @@ def assign(self, source_id: str, destination_id: str) -> ArrayValue:
                 child=self.node,
                 assignments=tuple(exprs),
             )
-        ).merge_projections()
+        )
 
     def assign_constant(
         self,
@@ -243,7 +243,7 @@ def assign_constant(
                 child=self.node,
                 assignments=tuple(exprs),
             )
-        ).merge_projections()
+        )
 
     def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue:
         selections = ((ex.free_var(col_id), col_id) for col_id in column_ids)
@@ -252,7 +252,7 @@ def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue:
                 child=self.node,
                 assignments=tuple(selections),
             )
-        ).merge_projections()
+        )
 
     def drop_columns(self, columns: Iterable[str]) -> ArrayValue:
         new_projection = (
@@ -265,7 +265,7 @@ def drop_columns(self, columns: Iterable[str]) -> ArrayValue:
                 child=self.node,
                 assignments=tuple(new_projection),
             )
-        ).merge_projections()
+        )
 
     def aggregate(
         self,

bigframes/core/blocks.py

@@ -2026,27 +2026,17 @@ def to_sql_query(
             idx_labels,
         )
 
-    def cached(self, *, optimize_offsets=False, force: bool = False) -> Block:
-        """Write the block to a session table and create a new block object that references it."""
+    def cached(self, *, optimize_offsets=False, force: bool = False) -> None:
+        """Write the block to a session table."""
         # use a heuristic for whether something needs to be cached
         if (not force) and self.session._is_trivially_executable(self.expr):
-            return self
+            return
         if optimize_offsets:
-            expr = self.session._cache_with_offsets(self.expr)
+            self.session._cache_with_offsets(self.expr)
         else:
-            expr = self.session._cache_with_cluster_cols(
+            self.session._cache_with_cluster_cols(
                 self.expr, cluster_cols=self.index_columns
             )
-        return self.swap_array_expr(expr)
-
-    def swap_array_expr(self, expr: core.ArrayValue) -> Block:
-        # TODO: Validate schema unchanged
-        return Block(
-            expr,
-            index_columns=self.index_columns,
-            column_labels=self.column_labels,
-            index_labels=self.index.names,
-        )
 
     def _is_monotonic(
         self, column_ids: typing.Union[str, Sequence[str]], increasing: bool
@@ -2116,8 +2106,8 @@ def _get_rows_as_json_values(self) -> Block:
         # TODO(shobs): Replace direct SQL manipulation by structured expression
         # manipulation
         ordering_column_name = guid.generate_guid()
-        expr = self.session._cache_with_offsets(self.expr)
-        expr = expr.promote_offsets(ordering_column_name)
+        self.session._cache_with_offsets(self.expr)
+        expr = self.expr.promote_offsets(ordering_column_name)
         expr_sql = self.session._to_sql(expr)
 
         # Names of the columns to serialize for the row.

bigframes/core/indexes/base.py

@@ -95,6 +95,7 @@ def __new__(
         result = typing.cast(Index, object.__new__(klass))  # type: ignore
         result._query_job = None
         result._block = block
+        block.session._register_object(result)
         return result
 
     @classmethod

bigframes/core/tree_properties.py

@@ -15,7 +15,7 @@
 
 import functools
 import itertools
-from typing import Dict
+from typing import Callable, Dict, Optional
 
 import bigframes.core.nodes as nodes
 
@@ -40,46 +40,66 @@ def peekable(node: nodes.BigFrameNode) -> bool:
     return children_peekable and self_peekable
 
 
-def count_complex_nodes(
-    root: nodes.BigFrameNode, min_complexity: float, max_complexity: float
-) -> Dict[nodes.BigFrameNode, int]:
+# Replace modified_cost(node) = cost(apply_cache(node))
+def select_cache_target(
+    root: nodes.BigFrameNode,
+    min_complexity: float,
+    max_complexity: float,
+    cache: dict[nodes.BigFrameNode, nodes.BigFrameNode],
+    heuristic: Callable[[int, int], float],
+) -> Optional[nodes.BigFrameNode]:
+    """Take tree, and return candidate nodes with (# of occurences, post-caching planning complexity).
+
+    heurstic takes two args, node complexity, and node occurence count, in that order
+    """
+
+    @functools.cache
+    def _with_caching(subtree: nodes.BigFrameNode) -> nodes.BigFrameNode:
+        return replace_nodes(subtree, cache)
+
+    def _combine_counts(
+        left: Dict[nodes.BigFrameNode, int], right: Dict[nodes.BigFrameNode, int]
+    ) -> Dict[nodes.BigFrameNode, int]:
+        return {
+            key: left.get(key, 0) + right.get(key, 0)
+            for key in itertools.chain(left.keys(), right.keys())
+        }
+
     @functools.cache
     def _node_counts_inner(
         subtree: nodes.BigFrameNode,
     ) -> Dict[nodes.BigFrameNode, int]:
         """Helper function to count occurences of duplicate nodes in a subtree. Considers only nodes in a complexity range"""
         empty_counts: Dict[nodes.BigFrameNode, int] = {}
-        if subtree.planning_complexity >= min_complexity:
+        subtree_complexity = _with_caching(subtree).planning_complexity
+        if subtree_complexity >= min_complexity:
             child_counts = [_node_counts_inner(child) for child in subtree.child_nodes]
             node_counts = functools.reduce(_combine_counts, child_counts, empty_counts)
-            if subtree.planning_complexity <= max_complexity:
+            if subtree_complexity <= max_complexity:
                 return _combine_counts(node_counts, {subtree: 1})
             else:
                 return node_counts
         return empty_counts
 
-    return _node_counts_inner(root)
+    node_counts = _node_counts_inner(root)
+
+    return max(
+        node_counts.keys(),
+        key=lambda node: heuristic(
+            _with_caching(node).planning_complexity, node_counts[node]
+        ),
+    )
 
 
 def replace_nodes(
     root: nodes.BigFrameNode,
-    to_replace: nodes.BigFrameNode,
-    replacemenet: nodes.BigFrameNode,
+    replacements: dict[nodes.BigFrameNode, nodes.BigFrameNode],
 ):
     @functools.cache
-    def apply_substition(n: nodes.BigFrameNode) -> nodes.BigFrameNode:
-        if n == to_replace:
-            return replacemenet
+    def apply_substition(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
+        if node in replacements.keys():
+            return replacements[node]
         else:
-            return n.transform_children(apply_substition)
-
-    return root.transform_children(apply_substition)
-
+            return node.transform_children(apply_substition)
 
-def _combine_counts(
-    left: Dict[nodes.BigFrameNode, int], right: Dict[nodes.BigFrameNode, int]
-) -> Dict[nodes.BigFrameNode, int]:
-    return {
-        key: left.get(key, 0) + right.get(key, 0)
-        for key in itertools.chain(left.keys(), right.keys())
-    }
+    return apply_substition(root)

bigframes/dataframe.py

@@ -191,6 +191,7 @@ def __init__(
             else:
                 self._block = bigframes.pandas.read_pandas(pd_dataframe)._get_block()
         self._query_job: Optional[bigquery.QueryJob] = None
+        self._block.session._register_object(self)
 
     def __dir__(self):
         return dir(type(self)) + [
@@ -3515,16 +3516,15 @@ def _cached(self, *, force: bool = False) -> DataFrame:
         No-op if the dataframe represents a trivial transformation of an existing materialization.
         Force=True is used for BQML integration where need to copy data rather than use snapshot.
         """
-        self._set_block(self._block.cached(force=force))
+        self._block.cached(force=force)
         return self
 
     def _optimize_query_complexity(self):
         """Reduce query complexity by caching repeated subtrees and recursively materializing maximum-complexity subtrees.
         May generate many queries and take substantial time to execute.
         """
         # TODO: Move all this to session
-        new_expr = self._session._simplify_with_caching(self._block.expr)
-        self._set_block(self._block.swap_array_expr(new_expr))
+        self._session._simplify_with_caching(self._block.expr)
 
     _DataFrameOrSeries = typing.TypeVar("_DataFrameOrSeries")
 

bigframes/series.py

@@ -72,6 +72,7 @@ class Series(bigframes.operations.base.SeriesMethods, vendored_pandas_series.Ser
     def __init__(self, *args, **kwargs):
         self._query_job: Optional[bigquery.QueryJob] = None
         super().__init__(*args, **kwargs)
+        self._block.session._register_object(self)
 
     @property
     def dt(self) -> dt.DatetimeMethods:
@@ -1777,16 +1778,15 @@ def cache(self):
         return self._cached(force=True)
 
     def _cached(self, *, force: bool = True) -> Series:
-        self._set_block(self._block.cached(force=force))
+        self._block.cached(force=force)
         return self
 
     def _optimize_query_complexity(self):
         """Reduce query complexity by caching repeated subtrees and recursively materializing maximum-complexity subtrees.
         May generate many queries and take substantial time to execute.
         """
         # TODO: Move all this to session
-        new_expr = self._block.session._simplify_with_caching(self._block.expr)
-        self._set_block(self._block.swap_array_expr(new_expr))
+        self._block.session._simplify_with_caching(self._block.expr)
 
 
 def _is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]: