perf: multi-query execution of complex dataframes

Author: TrevorBergeron

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)
+
+    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/session/__init__.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import datetime
+import functools
 import itertools
 import logging
 import os
@@ -69,6 +70,7 @@
 import bigframes.core.blocks as blocks
 import bigframes.core.compile
 import bigframes.core.guid as guid
+import bigframes.core.nodes as nodes
 from bigframes.core.ordering import IntegerEncoding, OrderingColumnReference
 import bigframes.core.ordering as orderings
 import bigframes.core.traversal as traversals
@@ -110,6 +112,16 @@
 # TODO(tbergeron): Convert to bytes-based limit
 MAX_INLINE_DF_SIZE = 5000
 
+# Chosen to prevent very large expression trees from being directly executed as a single query
+# Beyond this limit, expression should be cached or decomposed into multiple queries for execution.
+COMPLEXITY_SOFT_LIMIT = 10**5
+# Beyond the hard limite, even decomposing the query is unlikely to succeed
+COMPLEXITY_HARD_LIMIT = 10**25
+# Number of times to factor out and cache a repeated subtree
+MAX_SUBTREE_FACTORINGS = 10
+# Limits how much bigframes will attempt to decompose complex queries into smaller queries
+MAX_DECOMPOSITION_STEPS = 5
+
 logger = logging.getLogger(__name__)
 
 
@@ -1671,6 +1683,87 @@ def _cache_with_offsets(self, array_value: core.ArrayValue) -> core.ArrayValue:
             ordering=orderings.ExpressionOrdering.from_offset_col("bigframes_offsets"),
         )
 
+    def _cache_subtrees(self, node: nodes.BigFrameNode) -> nodes.BigFrameNode:
+        """If the computaiton is complex, execute subtrees of the compuation."""
+        node = node
+        # Recurse into subtrees below complexity limit, then cache each of those
+        if node.complexity > COMPLEXITY_SOFT_LIMIT:
+            reduced = node.transform_children(lambda n: self._cache_subtrees(n))
+            return reduced
+        else:
+            # TODO: Add clustering columns based on access patterns
+            return self._cache_with_cluster_cols(core.ArrayValue(node), []).node
+
+    def _cache_repeated_subtrees(
+        self, node: nodes.BigFrameNode, max_iterations: int
+    ) -> nodes.BigFrameNode:
+        """If the computaiton is complex, execute subtrees of the compuation."""
+        root = node
+        # Identify node that maximizes complexity * (repetitions - 1)
+        # Recurse into subtrees below complexity limit, then cache each of those
+
+        for _ in range(max_iterations):
+            node_counts: Dict[nodes.BigFrameNode, int] = dict()
+
+            candidates = list(root.child_nodes)
+            while candidates:
+                candidate = candidates.pop(0)
+                # Don't bother with low complexity nodes
+                if candidate.complexity > COMPLEXITY_SOFT_LIMIT / 50:
+                    if candidate.complexity < COMPLEXITY_SOFT_LIMIT:
+                        node_count = node_counts.get(candidate, 0) + 1
+                        node_counts[candidate] = node_count
+                    candidates.extend(candidate.child_nodes)
+
+            # Only consider nodes the occur at least twice
+            valid_candidates = filter(lambda x: x[1] >= 2, node_counts.items())
+            # Heuristic: Complexity * (copies of subtree - 1), in other words, how much is complexity
+            # reduced by deduplicating
+            best_candidate = max(
+                valid_candidates, key=lambda i: i[0].complexity * i[1] - 1, default=None
+            )
+
+            if best_candidate is None:
+                # No good subtrees to cache, just return original tree
+                return root
+
+            node_to_replace = best_candidate[0]
+
+            # TODO: Add clustering columns based on access patterns
+            cached_node = self._cache_with_cluster_cols(
+                core.ArrayValue(best_candidate[0]), []
+            ).node
+
+            @functools.cache
+            def apply_substition(n: nodes.BigFrameNode) -> nodes.BigFrameNode:
+                if n == node_to_replace:
+                    return cached_node
+                else:
+                    return n.transform_children(apply_substition)
+
+            root = root.transform_children(apply_substition)
+        return root
+
+    def _simplify_with_caching(self, array_value: core.ArrayValue) -> core.ArrayValue:
+        """Attempts to handle the complexity by caching duplicated subtrees and breaking the query into pieces."""
+        node = array_value.node
+        if node.complexity < COMPLEXITY_SOFT_LIMIT:
+            return array_value
+        node = self._cache_repeated_subtrees(
+            node, max_iterations=MAX_SUBTREE_FACTORINGS
+        )
+        if node.complexity > COMPLEXITY_HARD_LIMIT:
+            raise ValueError(
+                f"This dataframe is too complex to convert to SQL queries. Internal complexity measure: {node.complexity}"
+            )
+        # TODO: add config flag to enable/disable recursive execution
+        for _ in range(MAX_DECOMPOSITION_STEPS):
+            if node.complexity > COMPLEXITY_SOFT_LIMIT:
+                node = self._cache_subtrees(node)
+            else:
+                return core.ArrayValue(node)
+        return core.ArrayValue(node)
+
     def _is_trivially_executable(self, array_value: core.ArrayValue):
         """
         Can the block be evaluated very cheaply?
@@ -1689,6 +1782,8 @@ def _execute(
         dry_run=False,
         col_id_overrides: Mapping[str, str] = {},
     ) -> tuple[bigquery.table.RowIterator, bigquery.QueryJob]:
+        array_value = self._simplify_with_caching(array_value)
+
         sql = self._to_sql(
             array_value, sorted=sorted, col_id_overrides=col_id_overrides
         )  # type:ignore