From 0beb85539eb1cdc93e138d963c8899e529ba83d0 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Tue, 26 Mar 2024 05:16:50 +0000 Subject: [PATCH 1/6] perf: Automatically squash internal projection nodes and use internal schema system. --- bigframes/core/__init__.py | 29 +++++++------------- bigframes/core/compile/compiled.py | 16 +++++++---- bigframes/core/rewrite.py | 43 ++++++++++++++++++++++++++---- bigframes/dataframe.py | 6 ----- bigframes/series.py | 6 ----- 5 files changed, 59 insertions(+), 41 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 6fd6fc23c2..2b446b96bc 100644 --- a/bigframes/core/__init__.py +++ b/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) diff --git a/bigframes/core/compile/compiled.py b/bigframes/core/compile/compiled.py index af2d69275a..6d9864dc23 100644 --- a/bigframes/core/compile/compiled.py +++ b/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, ) diff --git a/bigframes/core/rewrite.py b/bigframes/core/rewrite.py index 61fe28b7b5..046b6737c5 100644 --- a/bigframes/core/rewrite.py +++ b/bigframes/core/rewrite.py @@ -35,6 +35,7 @@ class SquashedSelect: columns: Tuple[Tuple[scalar_exprs.Expression, str], ...] predicate: Optional[scalar_exprs.Expression] ordering: Tuple[order.OrderingExpression, ...] + reverse_root: bool = False @classmethod def from_node(cls, node: nodes.BigFrameNode) -> SquashedSelect: @@ -63,7 +64,9 @@ def project( new_columns = tuple( (expr.bind_all_variables(self.column_lookup), id) for expr, id in projection ) - return SquashedSelect(self.root, new_columns, self.predicate, self.ordering) + return SquashedSelect( + self.root, new_columns, self.predicate, self.ordering, self.reverse_root + ) def filter(self, predicate: scalar_exprs.Expression) -> SquashedSelect: if self.predicate is None: @@ -72,18 +75,24 @@ def filter(self, predicate: scalar_exprs.Expression) -> SquashedSelect: new_predicate = ops.and_op.as_expr( self.predicate, predicate.bind_all_variables(self.column_lookup) ) - return SquashedSelect(self.root, self.columns, new_predicate, self.ordering) + return SquashedSelect( + self.root, self.columns, new_predicate, self.ordering, self.reverse_root + ) def reverse(self) -> SquashedSelect: new_ordering = tuple(expr.with_reverse() for expr in self.ordering) - return SquashedSelect(self.root, self.columns, self.predicate, new_ordering) + return SquashedSelect( + self.root, self.columns, self.predicate, new_ordering, not self.reverse_root + ) def order_with(self, by: Tuple[order.OrderingExpression, ...]): adjusted_orderings = [ order_part.bind_variables(self.column_lookup) for order_part in by ] new_ordering = (*adjusted_orderings, *self.ordering) - return SquashedSelect(self.root, self.columns, self.predicate, new_ordering) + return SquashedSelect( + self.root, self.columns, self.predicate, new_ordering, self.reverse_root + ) def maybe_join( self, right: SquashedSelect, join_def: join_defs.JoinDefinition @@ -126,8 +135,10 @@ def maybe_join( new_columns = remap_names(join_def, lselection, rselection) # Reconstruct ordering + reverse_root = self.reverse_root if join_type == "right": new_ordering = right.ordering + reverse_root = right.reverse_root elif join_type == "outer": if lmask is not None: prefix = order.OrderingExpression(lmask, order.OrderingDirection.DESC) @@ -158,11 +169,15 @@ def maybe_join( new_ordering = self.ordering else: raise ValueError(f"Unexpected join type {join_type}") - return SquashedSelect(self.root, new_columns, new_predicate, new_ordering) + return SquashedSelect( + self.root, new_columns, new_predicate, new_ordering, reverse_root + ) def expand(self) -> nodes.BigFrameNode: # Safest to apply predicates first, as it may filter out inputs that cannot be handled by other expressions root = self.root + if self.reverse_root: + root = nodes.ReversedNode(child=root) if self.predicate: root = nodes.FilterNode(child=root, predicate=self.predicate) if self.ordering: @@ -170,6 +185,24 @@ def expand(self) -> nodes.BigFrameNode: return nodes.ProjectionNode(child=root, assignments=self.columns) +def is_squashable(node: nodes.BigFrameNode) -> bool: + squashable_classes = ( + nodes.ProjectionNode, + nodes.FilterNode, + nodes.ReversedNode, + nodes.OrderByNode, + ) + return isinstance(node, squashable_classes) and isinstance( + node.child, squashable_classes + ) + + +def maybe_squash_projection(node: nodes.BigFrameNode) -> nodes.BigFrameNode: + if is_squashable(node): + return SquashedSelect.from_node(node).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) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 07dae2c53b..e3bebf2273 100644 --- a/bigframes/dataframe.py +++ b/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 diff --git a/bigframes/series.py b/bigframes/series.py index e7b358c2fe..fc45c027e2 100644 --- a/bigframes/series.py +++ b/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) From 374b9d1ac224cc10cfe9ed149828519211207773 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Wed, 24 Apr 2024 18:05:21 +0000 Subject: [PATCH 2/6] squash less aggressively --- bigframes/core/__init__.py | 1 + bigframes/core/rewrite.py | 29 +++++++++++------------------ 2 files changed, 12 insertions(+), 18 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 6de32b7f76..abef0abcc3 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -468,5 +468,6 @@ def _uniform_sampling(self, fraction: float) -> ArrayValue: return ArrayValue(nodes.RandomSampleNode(self.node, fraction)) def rewrite_projection(self) -> ArrayValue: + # Relatively conservative approach, logic can also handle filter/reordering nodes, but will leave those out. rewritten = bigframes.core.rewrite.maybe_squash_projection(self.node) return ArrayValue(rewritten) diff --git a/bigframes/core/rewrite.py b/bigframes/core/rewrite.py index 046b6737c5..6018c02cdd 100644 --- a/bigframes/core/rewrite.py +++ b/bigframes/core/rewrite.py @@ -38,14 +38,16 @@ class SquashedSelect: reverse_root: bool = False @classmethod - def from_node(cls, node: nodes.BigFrameNode) -> SquashedSelect: + def from_node( + cls, node: nodes.BigFrameNode, projections_only: bool = False + ) -> SquashedSelect: if isinstance(node, nodes.ProjectionNode): return cls.from_node(node.child).project(node.assignments) - elif isinstance(node, nodes.FilterNode): + elif not projections_only and isinstance(node, nodes.FilterNode): return cls.from_node(node.child).filter(node.predicate) - elif isinstance(node, nodes.ReversedNode): + elif not projections_only and isinstance(node, nodes.ReversedNode): return cls.from_node(node.child).reverse() - elif isinstance(node, nodes.OrderByNode): + elif not projections_only and isinstance(node, nodes.OrderByNode): return cls.from_node(node.child).order_with(node.by) else: selection = tuple( @@ -185,21 +187,12 @@ def expand(self) -> nodes.BigFrameNode: return nodes.ProjectionNode(child=root, assignments=self.columns) -def is_squashable(node: nodes.BigFrameNode) -> bool: - squashable_classes = ( - nodes.ProjectionNode, - nodes.FilterNode, - nodes.ReversedNode, - nodes.OrderByNode, - ) - return isinstance(node, squashable_classes) and isinstance( - node.child, squashable_classes - ) - - def maybe_squash_projection(node: nodes.BigFrameNode) -> nodes.BigFrameNode: - if is_squashable(node): - return SquashedSelect.from_node(node).expand() + if isinstance(node, nodes.ProjectionNode) and isinstance( + node.child, nodes.ProjectionNode + ): + # Conservative approach, only squash consecutive projections, even though could also squash filters, reorderings + return SquashedSelect.from_node(node, projections_only=True).expand() return node From 542ae49c9f6eaf1a56dc5ff446406f931837112b Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Thu, 25 Apr 2024 01:17:17 +0000 Subject: [PATCH 3/6] improve order baking --- bigframes/core/compile/compiled.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/bigframes/core/compile/compiled.py b/bigframes/core/compile/compiled.py index 35041337e0..c89646645b 100644 --- a/bigframes/core/compile/compiled.py +++ b/bigframes/core/compile/compiled.py @@ -1066,8 +1066,12 @@ def _bake_ordering(self) -> OrderedIR: ) new_exprs.append(new_expr) elif isinstance(expr.scalar_expression, ex.UnboundVariableExpression): + order_col = expr.scalar_expression.id new_exprs.append(expr) - new_baked_cols.append(self._ibis_bindings[expr.scalar_expression.id]) + if order_col not in self.columns: + new_baked_cols.append( + self._ibis_bindings[expr.scalar_expression.id] + ) new_ordering = ExpressionOrdering( tuple(new_exprs), From 8c899fabd5bd8d9f58450d6384a06b3699584bcf Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Thu, 25 Apr 2024 01:22:45 +0000 Subject: [PATCH 4/6] move rewrite to ArrayValue init --- bigframes/core/__init__.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index abef0abcc3..72f0c342ba 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -75,7 +75,12 @@ def from_ibis( hidden_ordering_columns=tuple(hidden_ordering_columns), ordering=ordering, ) - return cls(node) + # Relatively conservative approach, logic can also handle filter/reordering nodes, but will leave those out. + # Squash rewrites here to make sure tree doesn't get too deep, wasting memory and exceeding recursion limits. + # Avoid deep rewrites here though, as can invalidate caching strategies, instead do comprehensive rewrites at + # compile time. + rewritten = bigframes.core.rewrite.maybe_squash_projection(node) + return cls(rewritten) @classmethod def from_pyarrow(cls, arrow_table: pa.Table, session: Session): @@ -183,7 +188,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 @@ -208,7 +213,7 @@ def assign(self, source_id: str, destination_id: str) -> ArrayValue: child=self.node, assignments=tuple(exprs), ) - ).rewrite_projection() + ) def assign_constant( self, @@ -242,7 +247,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) @@ -251,7 +256,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 = ( @@ -264,7 +269,7 @@ def drop_columns(self, columns: Iterable[str]) -> ArrayValue: child=self.node, assignments=tuple(new_projection), ) - ).rewrite_projection() + ) def aggregate( self, @@ -466,8 +471,3 @@ 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: - # Relatively conservative approach, logic can also handle filter/reordering nodes, but will leave those out. - rewritten = bigframes.core.rewrite.maybe_squash_projection(self.node) - return ArrayValue(rewritten) From b49a8161d4934aa4cb3b8603e22f0b9bf5a60ea7 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Thu, 25 Apr 2024 21:37:34 +0000 Subject: [PATCH 5/6] fix condition in bake_ordering --- bigframes/core/__init__.py | 4 ++-- bigframes/core/compile/compiled.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 72f0c342ba..84d3110260 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -79,8 +79,8 @@ def from_ibis( # Squash rewrites here to make sure tree doesn't get too deep, wasting memory and exceeding recursion limits. # Avoid deep rewrites here though, as can invalidate caching strategies, instead do comprehensive rewrites at # compile time. - rewritten = bigframes.core.rewrite.maybe_squash_projection(node) - return cls(rewritten) + # node = bigframes.core.rewrite.maybe_squash_projection(node) + return cls(node) @classmethod def from_pyarrow(cls, arrow_table: pa.Table, session: Session): diff --git a/bigframes/core/compile/compiled.py b/bigframes/core/compile/compiled.py index c89646645b..9240118ba2 100644 --- a/bigframes/core/compile/compiled.py +++ b/bigframes/core/compile/compiled.py @@ -1053,8 +1053,8 @@ def _hide_column(self, column_id) -> OrderedIR: def _bake_ordering(self) -> OrderedIR: """Bakes ordering expression into the selection, maybe creating hidden columns.""" ordering_expressions = self._ordering.all_ordering_columns - new_exprs = [] - new_baked_cols = [] + new_exprs: list[OrderingExpression] = [] + new_baked_cols: list[ibis_types.Value] = [] for expr in ordering_expressions: if isinstance(expr.scalar_expression, ex.OpExpression): baked_column = self._compile_expression(expr.scalar_expression).name( @@ -1068,7 +1068,7 @@ def _bake_ordering(self) -> OrderedIR: elif isinstance(expr.scalar_expression, ex.UnboundVariableExpression): order_col = expr.scalar_expression.id new_exprs.append(expr) - if order_col not in self.columns: + if order_col not in self.column_ids: new_baked_cols.append( self._ibis_bindings[expr.scalar_expression.id] ) From 63f0465e9ff466bd369efcce35349ad237329d4e Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Thu, 25 Apr 2024 21:44:44 +0000 Subject: [PATCH 6/6] fix squash logic --- bigframes/core/__init__.py | 19 +++++++++---------- bigframes/core/rewrite.py | 4 +++- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 84d3110260..185ce7cd4f 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -75,11 +75,6 @@ def from_ibis( hidden_ordering_columns=tuple(hidden_ordering_columns), ordering=ordering, ) - # Relatively conservative approach, logic can also handle filter/reordering nodes, but will leave those out. - # Squash rewrites here to make sure tree doesn't get too deep, wasting memory and exceeding recursion limits. - # Avoid deep rewrites here though, as can invalidate caching strategies, instead do comprehensive rewrites at - # compile time. - # node = bigframes.core.rewrite.maybe_squash_projection(node) return cls(node) @classmethod @@ -188,7 +183,7 @@ def project_to_id(self, expression: ex.Expression, output_id: str): child=self.node, assignments=tuple(exprs), ) - ) + ).merge_projections() def assign(self, source_id: str, destination_id: str) -> ArrayValue: if destination_id in self.column_ids: # Mutate case @@ -213,7 +208,7 @@ def assign(self, source_id: str, destination_id: str) -> ArrayValue: child=self.node, assignments=tuple(exprs), ) - ) + ).merge_projections() def assign_constant( self, @@ -247,7 +242,7 @@ def assign_constant( child=self.node, assignments=tuple(exprs), ) - ) + ).merge_projections() def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue: selections = ((ex.free_var(col_id), col_id) for col_id in column_ids) @@ -256,7 +251,7 @@ def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue: child=self.node, assignments=tuple(selections), ) - ) + ).merge_projections() def drop_columns(self, columns: Iterable[str]) -> ArrayValue: new_projection = ( @@ -269,7 +264,7 @@ def drop_columns(self, columns: Iterable[str]) -> ArrayValue: child=self.node, assignments=tuple(new_projection), ) - ) + ).merge_projections() def aggregate( self, @@ -471,3 +466,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 merge_projections(self) -> ArrayValue: + new_node = bigframes.core.rewrite.maybe_squash_projection(self.node) + return ArrayValue(new_node) diff --git a/bigframes/core/rewrite.py b/bigframes/core/rewrite.py index 6018c02cdd..e3a07c04b4 100644 --- a/bigframes/core/rewrite.py +++ b/bigframes/core/rewrite.py @@ -42,7 +42,9 @@ def from_node( cls, node: nodes.BigFrameNode, projections_only: bool = False ) -> SquashedSelect: if isinstance(node, nodes.ProjectionNode): - return cls.from_node(node.child).project(node.assignments) + return cls.from_node(node.child, projections_only=projections_only).project( + node.assignments + ) elif not projections_only and isinstance(node, nodes.FilterNode): return cls.from_node(node.child).filter(node.predicate) elif not projections_only and isinstance(node, nodes.ReversedNode):