perf: multi-query execution of complex dataframes

Author: TrevorBergeron

bigframes/core/nodes.py

@@ -14,11 +14,11 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass, field, fields
+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 +100,16 @@ def roots(self) -> typing.Set[BigFrameNode]:
         )
         return set(roots)
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        """A crude measure of the query complexity. Not necessarily predictive of the complexity of the actual computation."""
+        return 0
+
+    def transform_children(
+        self, t: Callable[[BigFrameNode], BigFrameNode]
+    ) -> BigFrameNode:
+        return self
+
 
 @dataclass(frozen=True)
 class UnaryNode(BigFrameNode):
@@ -109,6 +119,11 @@ 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))
+
 
 @dataclass(frozen=True)
 class JoinNode(BigFrameNode):
@@ -138,6 +153,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 +177,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 +204,11 @@ 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
+
 
 # TODO: Refactor to take raw gbq object reference
 @dataclass(frozen=True)
@@ -192,6 +234,10 @@ def peekable(self) -> bool:
     def roots(self) -> typing.Set[BigFrameNode]:
         return {self}
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return len(self.columns)
+
 
 # Unary nodes
 @dataclass(frozen=True)
@@ -209,6 +255,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 +271,10 @@ def row_preserving(self) -> bool:
     def __hash__(self):
         return self._node_hash
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity
+
 
 @dataclass(frozen=True)
 class OrderByNode(UnaryNode):
@@ -229,6 +283,10 @@ class OrderByNode(UnaryNode):
     def __hash__(self):
         return self._node_hash
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity
+
 
 @dataclass(frozen=True)
 class ReversedNode(UnaryNode):
@@ -238,6 +296,10 @@ class ReversedNode(UnaryNode):
     def __hash__(self):
         return self._node_hash
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity
+
 
 @dataclass(frozen=True)
 class ProjectionNode(UnaryNode):
@@ -246,6 +308,10 @@ class ProjectionNode(UnaryNode):
     def __hash__(self):
         return self._node_hash
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity
+
 
 # TODO: Merge RowCount into Aggregate Node?
 # Row count can be compute from table metadata sometimes, so it is a bit special.
@@ -259,6 +325,10 @@ def row_preserving(self) -> bool:
     def non_local(self) -> bool:
         return True
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity
+
 
 @dataclass(frozen=True)
 class AggregateNode(UnaryNode):
@@ -281,6 +351,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 +376,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 +421,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 +438,9 @@ def deterministic(self) -> bool:
     def row_preserving(self) -> bool:
         return False
 
+    @functools.cached_property
+    def complexity(self) -> int:
+        return self.child.complexity
+
     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

@@ -70,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
@@ -112,6 +113,15 @@
 # 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.
+# May need tuning
+COMPLEXITY_SOFT_LIMIT = 50000
+# Beyond the complexity hard limit, will not even attempt to execute, as would require breaking into too many smaller queries.
+COMPLEXITY_HARD_LIMIT = COMPLEXITY_SOFT_LIMIT * 10
+# Limits how much bigframes will attempt to decompose complex queries into smaller queries
+MAX_DECOMPOSITION_STEPS = 1
+
 logger = logging.getLogger(__name__)
 
 
@@ -1610,6 +1620,82 @@ 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
+
+            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=10)
+        if node.complexity > COMPLEXITY_HARD_LIMIT:
+            raise ValueError("This dataframe is too complex to convert to SQL queries.")
+        # 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?
@@ -1628,6 +1714,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

tests/system/small/test_dataframe.py

@@ -27,6 +27,7 @@
 import bigframes
 import bigframes._config.display_options as display_options
 import bigframes.dataframe as dataframe
+import bigframes.pandas
 import bigframes.series as series
 from tests.system.utils import (
     assert_pandas_df_equal,
@@ -3882,6 +3883,13 @@ def test_recursion_limit(scalars_df_index):
     scalars_df_index.to_pandas()
 
 
+def test_query_complexity(scalars_df_index):
+    # Recursively union the data
+    for _ in range(3):
+        scalars_df_index = bigframes.pandas.concat(10 * [scalars_df_index]).head(5)
+    scalars_df_index.to_pandas()
+
+
 def test_to_pandas_downsampling_option_override(session):
     df = session.read_gbq("bigframes-dev.bigframes_tests_sys.batting")
     download_size = 1