perf: multi-query execution of complex dataframes

Author: TrevorBergeron

bigframes/core/blocks.py

@@ -1808,6 +1808,10 @@ def cached(self, *, optimize_offsets=False, force: bool = False) -> Block:
             expr = 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,

bigframes/core/nodes.py

@@ -14,11 +14,12 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass, field, fields
+import abc
+from dataclasses import dataclass, field, fields, replace
 import functools
 import itertools
 import typing
-from typing import Tuple
+from typing import Callable, Tuple
 
 import pandas
 
@@ -100,6 +101,19 @@ def roots(self) -> typing.Set[BigFrameNode]:
         )
         return set(roots)
 
+    @functools.cached_property
+    @abc.abstractmethod
+    def complexity(self) -> int:
+        """A crude measure of the query complexity. Not necessarily predictive of the complexity of the actual computation."""
+        ...
+
+    @abc.abstractmethod
+    def transform_children(
+        self, t: Callable[[BigFrameNode], BigFrameNode]
+    ) -> BigFrameNode:
+        """Apply a function to each child node."""
+        ...
+
 
 @dataclass(frozen=True)
 class UnaryNode(BigFrameNode):
@@ -109,6 +123,15 @@ class UnaryNode(BigFrameNode):
     def child_nodes(self) -> typing.Sequence[BigFrameNode]:
         return (self.child,)
 
+    def transform_children(
+        self, t: Callable[[BigFrameNode], BigFrameNode]
+    ) -> BigFrameNode:
+        return replace(self, child=t(self.child))
+
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity + 1
+
 
 @dataclass(frozen=True)
 class JoinNode(BigFrameNode):
@@ -138,6 +161,18 @@ def peekable(self) -> bool:
         single_root = len(self.roots) == 1
         return children_peekable and single_root
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        child_complexity = sum(child.complexity for child in self.child_nodes)
+        return child_complexity * 2
+
+    def transform_children(
+        self, t: Callable[[BigFrameNode], BigFrameNode]
+    ) -> BigFrameNode:
+        return replace(
+            self, left_child=t(self.left_child), right_child=t(self.right_child)
+        )
+
 
 @dataclass(frozen=True)
 class ConcatNode(BigFrameNode):
@@ -150,6 +185,16 @@ def child_nodes(self) -> typing.Sequence[BigFrameNode]:
     def __hash__(self):
         return self._node_hash
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        child_complexity = sum(child.complexity for child in self.child_nodes)
+        return child_complexity * 2
+
+    def transform_children(
+        self, t: Callable[[BigFrameNode], BigFrameNode]
+    ) -> BigFrameNode:
+        return replace(self, children=tuple(t(child) for child in self.children))
+
 
 # Input Nodex
 @dataclass(frozen=True)
@@ -167,6 +212,16 @@ def peekable(self) -> bool:
     def roots(self) -> typing.Set[BigFrameNode]:
         return {self}
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        # TODO: Set to number of columns once this is more readily available
+        return 500
+
+    def transform_children(
+        self, t: Callable[[BigFrameNode], BigFrameNode]
+    ) -> BigFrameNode:
+        return self
+
 
 # TODO: Refactor to take raw gbq object reference
 @dataclass(frozen=True)
@@ -192,6 +247,15 @@ def peekable(self) -> bool:
     def roots(self) -> typing.Set[BigFrameNode]:
         return {self}
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return len(self.columns) + 5
+
+    def transform_children(
+        self, t: Callable[[BigFrameNode], BigFrameNode]
+    ) -> BigFrameNode:
+        return self
+
 
 # Unary nodes
 @dataclass(frozen=True)
@@ -209,6 +273,10 @@ def peekable(self) -> bool:
     def non_local(self) -> bool:
         return False
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity * 2
+
 
 @dataclass(frozen=True)
 class FilterNode(UnaryNode):
@@ -221,6 +289,10 @@ def row_preserving(self) -> bool:
     def __hash__(self):
         return self._node_hash
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity + 1
+
 
 @dataclass(frozen=True)
 class OrderByNode(UnaryNode):
@@ -229,6 +301,10 @@ class OrderByNode(UnaryNode):
     def __hash__(self):
         return self._node_hash
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity + 1
+
 
 @dataclass(frozen=True)
 class ReversedNode(UnaryNode):
@@ -238,6 +314,10 @@ class ReversedNode(UnaryNode):
     def __hash__(self):
         return self._node_hash
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity + 1
+
 
 @dataclass(frozen=True)
 class ProjectionNode(UnaryNode):
@@ -246,6 +326,10 @@ class ProjectionNode(UnaryNode):
     def __hash__(self):
         return self._node_hash
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity + 1
+
 
 # TODO: Merge RowCount into Aggregate Node?
 # Row count can be compute from table metadata sometimes, so it is a bit special.
@@ -259,6 +343,10 @@ def row_preserving(self) -> bool:
     def non_local(self) -> bool:
         return True
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity + 1
+
 
 @dataclass(frozen=True)
 class AggregateNode(UnaryNode):
@@ -281,6 +369,10 @@ def peekable(self) -> bool:
     def non_local(self) -> bool:
         return True
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity * 2
+
 
 @dataclass(frozen=True)
 class WindowOpNode(UnaryNode):
@@ -302,12 +394,22 @@ def peekable(self) -> bool:
     def non_local(self) -> bool:
         return True
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        if self.skip_reproject_unsafe:
+            return self.child.complexity
+        return self.child.complexity * 2
+
 
 @dataclass(frozen=True)
 class ReprojectOpNode(UnaryNode):
     def __hash__(self):
         return self._node_hash
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity * 2
+
 
 @dataclass(frozen=True)
 class UnpivotNode(UnaryNode):
@@ -337,6 +439,10 @@ def non_local(self) -> bool:
     def peekable(self) -> bool:
         return False
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity * 2
+
 
 @dataclass(frozen=True)
 class RandomSampleNode(UnaryNode):
@@ -350,5 +456,9 @@ def deterministic(self) -> bool:
     def row_preserving(self) -> bool:
         return False
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity + 1
+
     def __hash__(self):
         return self._node_hash

bigframes/core/traversal.py

@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
+
 import bigframes.core.nodes as nodes
 
 
+@functools.cache
 def is_trivially_executable(node: nodes.BigFrameNode) -> bool:
     if local_only(node):
         return True

bigframes/dataframe.py

@@ -1060,6 +1060,7 @@ def to_pandas(
                 downsampled rows and all columns of this DataFrame.
         """
         # TODO(orrbradford): Optimize this in future. Potentially some cases where we can return the stored query job
+        self._optimize_query_complexity()
         df, query_job = self._block.to_pandas(
             max_download_size=max_download_size,
             sampling_method=sampling_method,
@@ -1071,6 +1072,7 @@ def to_pandas(
 
     def to_pandas_batches(self) -> Iterable[pandas.DataFrame]:
         """Stream DataFrame results to an iterable of pandas DataFrame"""
+        self._optimize_query_complexity()
         return self._block.to_pandas_batches()
 
     def _compute_dry_run(self) -> bigquery.QueryJob:
@@ -2973,6 +2975,7 @@ def _run_io_query(
         """Executes a query job presenting this dataframe and returns the destination
         table."""
         session = self._block.expr.session
+        self._optimize_query_complexity()
         export_array, id_overrides = self._prepare_export(
             index=index, ordering_id=ordering_id
         )
@@ -3109,6 +3112,14 @@ def _cached(self, *, force: bool = False) -> DataFrame:
         self._set_block(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))
+
     _DataFrameOrSeries = typing.TypeVar("_DataFrameOrSeries")
 
     def dot(self, other: _DataFrameOrSeries) -> _DataFrameOrSeries:

bigframes/series.py

@@ -160,6 +160,7 @@ def __len__(self):
         return self.shape[0]
 
     def __iter__(self) -> typing.Iterator:
+        self._optimize_query_complexity()
         return itertools.chain.from_iterable(
             map(lambda x: x.squeeze(axis=1), self._block.to_pandas_batches())
         )
@@ -312,6 +313,7 @@ def to_pandas(
             pandas.Series: A pandas Series with all rows of this Series if the data_sampling_threshold_mb
                 is not exceeded; otherwise, a pandas Series with downsampled rows of the DataFrame.
         """
+        self._optimize_query_complexity()
         df, query_job = self._block.to_pandas(
             max_download_size=max_download_size,
             sampling_method=sampling_method,
@@ -1573,6 +1575,14 @@ def _cached(self, *, force: bool = True) -> Series:
         self._set_block(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))
+
 
 def _is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]:
     return pandas.api.types.is_list_like(obj)