perf: Automatically squash internal projection nodes and use internal schema system.

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -107,7 +107,7 @@ def session(self) -> Session:
     @functools.cached_property
     def schema(self) -> schemata.ArraySchema:
         # TODO: switch to use self.node.schema
-        return self._compiled_schema
+        return self.node.schema
 
     @functools.cached_property
     def _compiled_schema(self) -> schemata.ArraySchema:
@@ -118,18 +118,6 @@ def _compiled_schema(self) -> schemata.ArraySchema:
         )
         return schemata.ArraySchema(items)
 
-    def validate_schema(self):
-        tree_derived = self.node.schema
-        ibis_derived = self._compiled_schema
-        if tree_derived.names != ibis_derived.names:
-            raise ValueError(
-                f"Unexpected names internal {tree_derived.names} vs compiled {ibis_derived.names}"
-            )
-        if tree_derived.dtypes != ibis_derived.dtypes:
-            raise ValueError(
-                f"Unexpected types internal {tree_derived.dtypes} vs compiled {ibis_derived.dtypes}"
-            )
-
     def _try_evaluate_local(self):
         """Use only for unit testing paths - not fully featured. Will throw exception if fails."""
         import ibis
@@ -196,7 +184,7 @@ def project_to_id(self, expression: ex.Expression, output_id: str):
                 child=self.node,
                 assignments=tuple(exprs),
             )
-        )
+        ).rewrite_projection()
 
     def assign(self, source_id: str, destination_id: str) -> ArrayValue:
         if destination_id in self.column_ids:  # Mutate case
@@ -221,7 +209,7 @@ def assign(self, source_id: str, destination_id: str) -> ArrayValue:
                 child=self.node,
                 assignments=tuple(exprs),
             )
-        )
+        ).rewrite_projection()
 
     def assign_constant(
         self,
@@ -251,7 +239,7 @@ def assign_constant(
                 child=self.node,
                 assignments=tuple(exprs),
             )
-        )
+        ).rewrite_projection()
 
     def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue:
         selections = ((ex.free_var(col_id), col_id) for col_id in column_ids)
@@ -260,7 +248,7 @@ def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue:
                 child=self.node,
                 assignments=tuple(selections),
             )
-        )
+        ).rewrite_projection()
 
     def drop_columns(self, columns: Iterable[str]) -> ArrayValue:
         new_projection = (
@@ -273,7 +261,7 @@ def drop_columns(self, columns: Iterable[str]) -> ArrayValue:
                 child=self.node,
                 assignments=tuple(new_projection),
             )
-        )
+        ).rewrite_projection()
 
     def aggregate(
         self,
@@ -404,3 +392,7 @@ def _uniform_sampling(self, fraction: float) -> ArrayValue:
             The row numbers of result is non-deterministic, avoid to use.
         """
         return ArrayValue(nodes.RandomSampleNode(self.node, fraction))
+
+    def rewrite_projection(self) -> ArrayValue:
+        rewritten = bigframes.core.rewrite.maybe_squash_projection(self.node)
+        return ArrayValue(rewritten)

bigframes/core/rewrite.py

@@ -170,6 +170,13 @@ def expand(self) -> nodes.BigFrameNode:
         return nodes.ProjectionNode(child=root, assignments=self.columns)
 
 
+def maybe_squash_projection(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
+    squashed = SquashedSelect.from_node(node)
+    if squashed.root not in node.child_nodes:
+        return squashed.expand()
+    return node
+
+
 def maybe_rewrite_join(join_node: nodes.JoinNode) -> nodes.BigFrameNode:
     left_side = SquashedSelect.from_node(join_node.left_child)
     right_side = SquashedSelect.from_node(join_node.right_child)

bigframes/dataframe.py

@@ -17,7 +17,6 @@
 from __future__ import annotations
 
 import datetime
-import os
 import re
 import sys
 import textwrap
@@ -174,11 +173,6 @@ def __init__(
                 self._block = bigframes.pandas.read_pandas(pd_dataframe)._get_block()
         self._query_job: Optional[bigquery.QueryJob] = None
 
-        # Runs strict validations to ensure internal type predictions and ibis are completely in sync
-        # Do not execute these validations outside of testing suite.
-        if "PYTEST_CURRENT_TEST" in os.environ:
-            self._block.expr.validate_schema()
-
     def __dir__(self):
         return dir(type(self)) + [
             label

tests/system/small/test_dataframe.py

@@ -14,7 +14,6 @@
 
 import io
 import operator
-import sys
 import tempfile
 import typing
 from typing import Tuple
@@ -4034,13 +4033,6 @@ def test_df_dot_operator_series(
     )
 
 
-# TODO(tswast): We may be able to re-enable this test after we break large
-# queries up in https://github.com/googleapis/python-bigquery-dataframes/pull/427
-@pytest.mark.skipif(
-    sys.version_info >= (3, 12),
-    # See: https://github.com/python/cpython/issues/112282
-    reason="setrecursionlimit has no effect on the Python C stack since Python 3.12.",
-)
 def test_recursion_limit(scalars_df_index):
     scalars_df_index = scalars_df_index[["int64_too", "int64_col", "float64_col"]]
     for i in range(400):