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

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -106,8 +106,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 +117,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 +183,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 +208,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 +238,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 +247,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 +260,7 @@ def drop_columns(self, columns: Iterable[str]) -> ArrayValue:
                 child=self.node,
                 assignments=tuple(new_projection),
             )
-        )
+        ).rewrite_projection()
 
     def aggregate(
         self,
@@ -404,3 +391,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/compile/compiled.py

@@ -1228,18 +1228,24 @@ def _bake_ordering(self) -> OrderedIR:
                 )
                 new_baked_cols.append(baked_column)
                 new_expr = OrderingExpression(
-                    ex.free_var(baked_column.name), expr.direction, expr.na_last
+                    ex.free_var(baked_column.get_name()), expr.direction, expr.na_last
                 )
                 new_exprs.append(new_expr)
-            else:
+            elif isinstance(expr.scalar_expression, ex.UnboundVariableExpression):
                 new_exprs.append(expr)
+                new_baked_cols.append(self._ibis_bindings[expr.scalar_expression.id])
 
-        ordering = self._ordering.with_ordering_columns(new_exprs)
+        new_ordering = ExpressionOrdering(
+            tuple(new_exprs),
+            self._ordering.integer_encoding,
+            self._ordering.string_encoding,
+            self._ordering.total_ordering_columns,
+        )
         return OrderedIR(
             self._table,
             columns=self.columns,
-            hidden_ordering_columns=[*self._hidden_ordering_columns, *new_baked_cols],
-            ordering=ordering,
+            hidden_ordering_columns=tuple(new_baked_cols),
+            ordering=new_ordering,
             predicates=self._predicates,
         )
 

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

bigframes/series.py

@@ -19,7 +19,6 @@
 import functools
 import itertools
 import numbers
-import os
 import textwrap
 import typing
 from typing import Any, Literal, Mapping, Optional, Tuple, Union
@@ -72,11 +71,6 @@ def __init__(self, *args, **kwargs):
         self._query_job: Optional[bigquery.QueryJob] = None
         super().__init__(*args, **kwargs)
 
-        # 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()
-
     @property
     def dt(self) -> dt.DatetimeMethods:
         return dt.DatetimeMethods(self._block)

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):