refactor ir internals to reduce select count

Author: TrevorBergeron

bigframes/core/compile/sqlglot/compiler.py

@@ -126,7 +126,7 @@ def compile_node(
     for current_node in list(node.iter_nodes_topo()):
         if current_node.child_nodes == ():
             # For leaf node, generates a dumpy child to pass the UID generator.
-            child_results = tuple([sqlglot_ir.SQLGlotIR(uid_gen=uid_gen)])
+            child_results = tuple([sqlglot_ir.SQLGlotIR.empty(uid_gen=uid_gen)])
         else:
             # Child nodes should have been compiled in the reverse topological order.
             child_results = tuple(

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import abc
 import dataclasses
 import datetime
 import functools
@@ -38,20 +39,107 @@
     to_wkt = dumps
 
 
+class SelectableFragment(abc.ABC):
+    """
+    Represent a grammar fragment that can be converted to a SELECT or FROM item.
+    """
+
+    def as_select_all(self) -> sge.Select:
+        ...
+
+    def select(self, *items: sge.Expression) -> sge.Select:
+        ...
+
+    def as_from_item(self) -> sge.FromItem:
+        ...
+
+
+class SelectFragment(SelectableFragment):
+    def __init__(self, select_expr: sge.Select):
+        self.select_expr = select_expr
+
+    def as_select_all(self) -> sge.Select:
+        return self.select_expr
+
+    def select(self, *items: sge.Expression) -> sge.Select:
+        return sge.Select().select(*items).from_(self.select_expr.subquery())
+
+    def as_from_item(self) -> sge.FromItem:
+        return self.select_expr.subquery()
+
+
+class TableFragment(SelectableFragment):
+    def __init__(self, table: sge.Table | sge.Unnest):
+        self.table = table
+
+    def as_select_all(self) -> sge.Select:
+        return sge.Select().select(sge.Star()).from_(self.table)
+
+    def select(self, *items: sge.Expression) -> sge.Select:
+        return sge.Select().select(*items).from_(self.table)
+
+    def as_from_item(self) -> sge.FromItem:
+        return self.table
+
+
+class DeferredSelectFragment(SelectableFragment):
+    def __init__(self, select_supplier: typing.Callable[[sge.Select], sge.Select]):
+        self.select_supplier = select_supplier
+
+    def as_select_all(self) -> sge.Select:
+        return self.select_supplier(sge.Select().select(sge.Star()))
+
+    def select(self, *items: sge.Expression) -> sge.Select:
+        return self.select_supplier(sge.Select().select(*items))
+
+    def as_from_item(self) -> sge.FromItem:
+        return self.select_supplier(sge.Select().select(sge.Star())).subquery()
+
+
 @dataclasses.dataclass(frozen=True)
 class SQLGlotIR:
     """Helper class to build SQLGlot Query and generate SQL string."""
 
-    expr: typing.Union[sge.Select, sge.Table] = sg.select()
+    expr: SelectableFragment
     """The SQLGlot expression representing the query."""
 
     uid_gen: guid.SequentialUIDGenerator = guid.SequentialUIDGenerator()
     """Generator for unique identifiers."""
 
+    def __post_init__(self):
+        assert isinstance(self.expr, SelectableFragment)
+
     @property
     def sql(self) -> str:
         """Generate SQL string from the given expression."""
-        return sql.to_sql(self.expr)
+        return sql.to_sql(self.expr.as_select_all())
+
+    @classmethod
+    def empty(
+        cls, uid_gen: guid.SequentialUIDGenerator = guid.SequentialUIDGenerator()
+    ) -> SQLGlotIR:
+        return cls(expr=SelectFragment(sge.select()), uid_gen=uid_gen)
+
+    @classmethod
+    def from_expr(
+        cls,
+        expr: sge.Expression,
+        uid_gen: guid.SequentialUIDGenerator = guid.SequentialUIDGenerator(),
+    ) -> SQLGlotIR:
+        if isinstance(expr, sge.Select):
+            return cls(expr=SelectFragment(expr), uid_gen=uid_gen)
+        elif isinstance(expr, (sge.Table, sge.Unnest)):
+            return cls(expr=TableFragment(expr), uid_gen=uid_gen)
+        else:
+            raise ValueError(f"Unsupported expression type: {type(expr)}")
+
+    @classmethod
+    def from_func(
+        cls,
+        select_handler: typing.Callable[[sge.Select], sge.Select],
+        uid_gen: guid.SequentialUIDGenerator = guid.SequentialUIDGenerator(),
+    ):
+        return cls(expr=DeferredSelectFragment(select_handler), uid_gen=uid_gen)
 
     @classmethod
     def from_pyarrow(
@@ -97,7 +185,7 @@ def from_pyarrow(
                 ),
             ],
         )
-        return cls(expr=sg.select(sge.Star()).from_(expr), uid_gen=uid_gen)
+        return cls.from_expr(expr=expr, uid_gen=uid_gen)
 
     @classmethod
     def from_table(
@@ -143,9 +231,9 @@ def from_table(
             select_expr = select_expr.where(
                 sg.parse_one(sql_predicate, dialect=sql.base.DIALECT), append=False
             )
-            return cls(expr=select_expr, uid_gen=uid_gen)
+            return cls.from_expr(expr=select_expr, uid_gen=uid_gen)
 
-        return cls(expr=table_expr, uid_gen=uid_gen)
+        return cls.from_expr(expr=table_expr, uid_gen=uid_gen)
 
     @classmethod
     def from_cte_ref(
@@ -156,7 +244,7 @@ def from_cte_ref(
         table_expr = sge.Table(
             this=sql.identifier(cte_ref),
         )
-        return cls(expr=table_expr, uid_gen=uid_gen)
+        return cls.from_expr(expr=table_expr, uid_gen=uid_gen)
 
     def select(
         self,
@@ -191,7 +279,7 @@ def select(
         if limit is not None:
             new_expr = new_expr.limit(limit)
 
-        return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
+        return SQLGlotIR.from_expr(expr=new_expr, uid_gen=self.uid_gen)
 
     @classmethod
     def from_unparsed_query(
@@ -209,7 +297,7 @@ def from_unparsed_query(
         )
         select_expr = sge.Select().select(sge.Star()).from_(sge.Table(this=cte_name))
         select_expr = _set_query_ctes(select_expr, [cte])
-        return cls(expr=select_expr, uid_gen=uid_gen)
+        return cls.from_expr(expr=select_expr, uid_gen=uid_gen)
 
     @classmethod
     def from_union(
@@ -241,7 +329,7 @@ def from_union(
         final_select_expr = (
             sge.Select().select(*selections).from_(union_expr.subquery())
         )
-        return cls(expr=final_select_expr, uid_gen=uid_gen)
+        return cls.from_expr(expr=final_select_expr, uid_gen=uid_gen)
 
     def join(
         self,
@@ -262,15 +350,13 @@ def join(
         )
 
         join_type_str = join_type if join_type != "outer" else "full outer"
-        new_expr = (
-            sge.Select()
-            .select(sge.Star())
-            .from_(left_from)
-            .join(right_from, on=join_on, join_type=join_type_str)
+        return SQLGlotIR.from_func(
+            lambda select: select.from_(left_from).join(
+                right_from, on=join_on, join_type=join_type_str
+            ),
+            uid_gen=self.uid_gen,
         )
 
-        return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
-
     def isin_join(
         self,
         right: SQLGlotIR,
@@ -280,7 +366,6 @@ def isin_join(
     ) -> SQLGlotIR:
         """Joins the current query with another SQLGlotIR instance."""
         left_from = self._as_from_item()
-        right_select = right._as_select()
 
         new_column: sge.Expression
         if joins_nulls:
@@ -294,7 +379,7 @@ def isin_join(
                 ]
             )
             right_expr1, right_expr2 = _value_to_non_null_identity(conditions[1])
-            right_select = right_select.select(
+            right_select = right.expr.select(
                 *[
                     sge.Struct(
                         expressions=[
@@ -303,7 +388,6 @@ def isin_join(
                         ]
                     )
                 ],
-                append=False,
             )
 
             new_column = sge.In(
@@ -313,7 +397,7 @@ def isin_join(
         else:
             new_column = sge.In(
                 this=conditions[0].expr,
-                expressions=[right_select.subquery()],
+                expressions=[right._as_subquery()],
             )
 
         new_column = sge.Alias(
@@ -322,7 +406,7 @@ def isin_join(
         )
 
         new_expr = sge.Select().select(sge.Star(), new_column).from_(left_from)
-        return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
+        return SQLGlotIR.from_expr(expr=new_expr, uid_gen=self.uid_gen)
 
     def explode(
         self,
@@ -344,8 +428,8 @@ def sample(self, fraction: float) -> SQLGlotIR:
             expression=sql.literal(fraction, dtypes.FLOAT_DTYPE),
         )
 
-        new_expr = self._as_select().where(condition, append=False)
-        return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
+        new_expr = self.expr.as_select_all().where(condition, append=False)
+        return SQLGlotIR.from_expr(expr=new_expr, uid_gen=self.uid_gen)
 
     def aggregate(
         self,
@@ -368,10 +452,7 @@ def aggregate(
             for id, expr in aggregations
         ]
 
-        new_expr = self._as_select()
-        new_expr = new_expr.group_by(*by_cols).select(
-            *[*by_cols, *aggregations_expr], append=False
-        )
+        new_expr = self.expr.select(*[*by_cols, *aggregations_expr]).group_by(*by_cols)
 
         condition = _and(
             tuple(
@@ -381,7 +462,7 @@ def aggregate(
         )
         if condition is not None:
             new_expr = new_expr.where(condition, append=False)
-        return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
+        return SQLGlotIR.from_expr(expr=new_expr, uid_gen=self.uid_gen)
 
     def with_ctes(
         self,
@@ -395,7 +476,7 @@ def with_ctes(
             for cte_name, cte in ctes
         ]
         select_expr = _set_query_ctes(self._as_select(), sge_ctes)
-        return SQLGlotIR(expr=select_expr, uid_gen=self.uid_gen)
+        return SQLGlotIR.from_expr(expr=select_expr, uid_gen=self.uid_gen)
 
     def resample(
         self,
@@ -431,7 +512,7 @@ def resample(
             .join(unnest_expr, join_type="cross")
         )
 
-        return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
+        return SQLGlotIR.from_expr(expr=new_expr, uid_gen=self.uid_gen)
 
     def _explode_single_column(
         self, column_name: str, offsets_col: typing.Optional[str]
@@ -449,12 +530,9 @@ def _explode_single_column(
         )
         selection = sge.Star(replace=[unnested_column_alias.as_(column)])
 
-        new_expr = self._as_select()
         # Use LEFT JOIN to preserve rows when unnesting empty arrays.
-        new_expr = new_expr.select(selection, append=False).join(
-            unnest_expr, join_type="LEFT"
-        )
-        return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
+        new_expr = self.expr.select(selection).join(unnest_expr, join_type="LEFT")
+        return SQLGlotIR.from_expr(expr=new_expr, uid_gen=self.uid_gen)
 
     def _explode_multiple_columns(
         self,
@@ -492,26 +570,18 @@ def _explode_multiple_columns(
                 for column in columns
             ]
         )
-        new_expr = self._as_select()
         # Use LEFT JOIN to preserve rows when unnesting empty arrays.
-        new_expr = new_expr.select(selection, append=False).join(
-            unnest_expr, join_type="LEFT"
-        )
-        return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
+        new_expr = self.expr.select(selection).join(unnest_expr, join_type="LEFT")
+        return SQLGlotIR.from_expr(expr=new_expr, uid_gen=self.uid_gen)
 
-    def _as_from_item(self) -> typing.Union[sge.Subquery, sge.Table]:
-        if isinstance(self.expr, sge.Select):
-            return self.expr.subquery()
-        else:  # table or cte
-            return self.expr
+    def _as_from_item(self) -> typing.Union[sge.Subquery, sge.Table, sge.Unnest]:
+        return self.expr.as_from_item()
 
     def _as_select(self) -> sge.Select:
-        if isinstance(self.expr, sge.Select):
-            return self.expr
-        else:  # table or cte
-            return sge.Select().select(sge.Star()).from_(self.expr)
+        return self.expr.as_select_all()
 
     def _as_subquery(self) -> sge.Subquery:
+        # Sometimes explicitly need a subquery, e.g. for IN expressions.
         return self._as_select().subquery()
 
 

bigframes/core/rewrite/pruning.py

@@ -67,7 +67,6 @@ def prune_selection_child(
 
     # Important to check this first
     if list(selection.ids) == list(child.ids):
-        # Added all() here - a generator object is natively truthy
         if all(ref.ref.id == ref.id for ref in selection.input_output_pairs):
             # selection is no-op so just remove it entirely
             return child

bigframes/session/bq_caching_executor.py

@@ -490,15 +490,6 @@ def prepare_plan(
 
         return plan
 
-    def simplify_plan(
-        self, plan: nodes.BigFrameNode, use_cache: bool = True
-    ) -> nodes.BigFrameNode:
-        if use_cache:
-            plan = self.replace_cached_subtrees(plan)
-        plan = rewrite.column_pruning(plan)
-        plan = plan.top_down(rewrite.fold_row_counts)
-        return plan
-
     def _cache_with_cluster_cols(
         self, array_value: bigframes.core.ArrayValue, cluster_cols: Sequence[str]
     ):

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_corr/out.sql

@@ -3,5 +3,10 @@ SELECT
 FROM (
   SELECT
     CORR(`int64_col`, `float64_col`) AS `bfcol_2`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  FROM (
+    SELECT
+      `int64_col`,
+      `float64_col`
+    FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  )
 )

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_cov/out.sql

@@ -3,5 +3,10 @@ SELECT
 FROM (
   SELECT
     COVAR_SAMP(`int64_col`, `float64_col`) AS `bfcol_2`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  FROM (
+    SELECT
+      `int64_col`,
+      `float64_col`
+    FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  )
 )