refactor: Simplify compile code paths (#1420)

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -198,12 +198,6 @@ def as_cached(
         )
         return ArrayValue(node)
 
-    def _try_evaluate_local(self):
-        """Use only for unit testing paths - not fully featured. Will throw exception if fails."""
-        import bigframes.core.compile
-
-        return bigframes.core.compile.test_only_try_evaluate(self.node)
-
     def get_column_type(self, key: str) -> bigframes.dtypes.Dtype:
         return self.schema.get_type(key)
 

bigframes/core/blocks.py

@@ -213,14 +213,10 @@ def index(self) -> BlockIndexProperties:
     @functools.cached_property
     def shape(self) -> typing.Tuple[int, int]:
         """Returns dimensions as (length, width) tuple."""
-
-        row_count_expr = self.expr.row_count()
-
-        # Support in-memory engines for hermetic unit tests.
-        if self.expr.session is None:
+        # Support zero-query for hermetic unit tests.
+        if self.expr.session is None and self.expr.node.row_count:
             try:
-                row_count = row_count_expr._try_evaluate_local().squeeze()
-                return (row_count, len(self.value_columns))
+                return self.expr.node.row_count
             except Exception:
                 pass
 

bigframes/core/compile/__init__.py

@@ -13,14 +13,9 @@
 # limitations under the License.
 from __future__ import annotations
 
-from bigframes.core.compile.api import (
-    SQLCompiler,
-    test_only_ibis_inferred_schema,
-    test_only_try_evaluate,
-)
+from bigframes.core.compile.api import SQLCompiler, test_only_ibis_inferred_schema
 
 __all__ = [
     "SQLCompiler",
-    "test_only_try_evaluate",
     "test_only_ibis_inferred_schema",
 ]

bigframes/core/compile/api.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import Mapping, Sequence, Tuple, TYPE_CHECKING
+from typing import Optional, Sequence, Tuple, TYPE_CHECKING
 
 import google.cloud.bigquery as bigquery
 
-import bigframes.core.compile.compiler as compiler
+from bigframes.core import rewrite
+from bigframes.core.compile import compiler
 
 if TYPE_CHECKING:
     import bigframes.core.nodes
@@ -31,31 +32,16 @@ class SQLCompiler:
     def __init__(self, strict: bool = True):
         self._compiler = compiler.Compiler(strict=strict)
 
-    def compile_peek(self, node: bigframes.core.nodes.BigFrameNode, n_rows: int) -> str:
-        """Compile node into sql that selects N arbitrary rows, may not execute deterministically."""
-        return self._compiler.compile_peek_sql(node, n_rows)
-
-    def compile_unordered(
-        self,
-        node: bigframes.core.nodes.BigFrameNode,
-        *,
-        col_id_overrides: Mapping[str, str] = {},
-    ) -> str:
-        """Compile node into sql where rows are unsorted, and no ordering information is preserved."""
-        # TODO: Enable limit pullup, but only if not being used to write to clustered table.
-        output_ids = [col_id_overrides.get(id, id) for id in node.schema.names]
-        return self._compiler.compile_sql(node, ordered=False, output_ids=output_ids)
-
-    def compile_ordered(
+    def compile(
         self,
         node: bigframes.core.nodes.BigFrameNode,
         *,
-        col_id_overrides: Mapping[str, str] = {},
+        ordered: bool = True,
+        limit: Optional[int] = None,
     ) -> str:
         """Compile node into sql where rows are sorted with ORDER BY."""
         # If we are ordering the query anyways, compiling the slice as a limit is probably a good idea.
-        output_ids = [col_id_overrides.get(id, id) for id in node.schema.names]
-        return self._compiler.compile_sql(node, ordered=True, output_ids=output_ids)
+        return self._compiler.compile_sql(node, ordered=ordered, limit=limit)
 
     def compile_raw(
         self,
@@ -67,21 +53,15 @@ def compile_raw(
         return self._compiler.compile_raw(node)
 
 
-def test_only_try_evaluate(node: bigframes.core.nodes.BigFrameNode):
-    """Use only for unit testing paths - not fully featured. Will throw exception if fails."""
-    node = _STRICT_COMPILER._preprocess(node)
-    ibis = _STRICT_COMPILER.compile_node(node)._to_ibis_expr()
-    return ibis.pandas.connect({}).execute(ibis)
-
-
 def test_only_ibis_inferred_schema(node: bigframes.core.nodes.BigFrameNode):
     """Use only for testing paths to ensure ibis inferred schema does not diverge from bigframes inferred schema."""
     import bigframes.core.schema
 
-    node = _STRICT_COMPILER._preprocess(node)
-    compiled = _STRICT_COMPILER.compile_node(node)
+    node = _STRICT_COMPILER._replace_unsupported_ops(node)
+    node, _ = rewrite.pull_up_order(node, order_root=False)
+    ir = _STRICT_COMPILER.compile_node(node)
     items = tuple(
-        bigframes.core.schema.SchemaItem(name, compiled.get_column_type(ibis_id))
-        for name, ibis_id in zip(node.schema.names, compiled.column_ids)
+        bigframes.core.schema.SchemaItem(name, ir.get_column_type(ibis_id))
+        for name, ibis_id in zip(node.schema.names, ir.column_ids)
     )
     return bigframes.core.schema.ArraySchema(items)

bigframes/core/compile/compiled.py

@@ -72,7 +72,11 @@ def to_sql(
     ) -> str:
         ibis_table = self._to_ibis_expr()
         # This set of output transforms maybe should be its own output node??
-        if order_by or limit:
+        if (
+            order_by
+            or limit
+            or (selections and (tuple(selections) != tuple(self.column_ids)))
+        ):
             sql = ibis_bigquery.Backend().compile(ibis_table)
             sql = (
                 bigframes.core.compile.googlesql.Select()

bigframes/core/compile/compiler.py

@@ -33,8 +33,6 @@
 import bigframes.core.compile.ibis_types
 import bigframes.core.compile.scalar_op_compiler as compile_scalar
 import bigframes.core.compile.schema_translator
-import bigframes.core.expression as ex
-import bigframes.core.identifiers as ids
 import bigframes.core.nodes as nodes
 import bigframes.core.ordering as bf_ordering
 import bigframes.core.rewrite as rewrites
@@ -52,65 +50,54 @@ class Compiler:
     scalar_op_compiler = compile_scalar.ScalarOpCompiler()
 
     def compile_sql(
-        self, node: nodes.BigFrameNode, ordered: bool, output_ids: typing.Sequence[str]
+        self,
+        node: nodes.BigFrameNode,
+        ordered: bool,
+        limit: typing.Optional[int] = None,
     ) -> str:
-        # TODO: get rid of output_ids arg
-        assert len(output_ids) == len(list(node.fields))
-        node = set_output_names(node, output_ids)
-        node = nodes.top_down(node, rewrites.rewrite_timedelta_expressions)
+        # later steps might add ids, so snapshot before those steps.
+        output_ids = node.schema.names
         if ordered:
-            node, limit = rewrites.pullup_limit_from_slice(node)
-            node = nodes.bottom_up(node, rewrites.rewrite_slice)
-            # TODO: Extract out CTEs
-            node, ordering = rewrites.pull_up_order(
-                node, order_root=True, ordered_joins=self.strict
-            )
-            node = rewrites.column_pruning(node)
-            ir = self.compile_node(node)
-            return ir.to_sql(
-                order_by=ordering.all_ordering_columns,
-                limit=limit,
-                selections=output_ids,
-            )
-        else:
-            node = nodes.bottom_up(node, rewrites.rewrite_slice)
-            node, _ = rewrites.pull_up_order(
-                node, order_root=False, ordered_joins=self.strict
-            )
-            node = rewrites.column_pruning(node)
-            ir = self.compile_node(node)
-            return ir.to_sql(selections=output_ids)
+            # Need to do this before replacing unsupported ops, as that will rewrite slice ops
+            node, pulled_up_limit = rewrites.pullup_limit_from_slice(node)
+            if (pulled_up_limit is not None) and (
+                (limit is None) or limit > pulled_up_limit
+            ):
+                limit = pulled_up_limit
 
-    def compile_peek_sql(self, node: nodes.BigFrameNode, n_rows: int) -> str:
-        ids = [id.sql for id in node.ids]
-        node = nodes.bottom_up(node, rewrites.rewrite_slice)
-        node = nodes.top_down(node, rewrites.rewrite_timedelta_expressions)
-        node, _ = rewrites.pull_up_order(
-            node, order_root=False, ordered_joins=self.strict
+        node = self._replace_unsupported_ops(node)
+        # prune before pulling up order to avoid unnnecessary row_number() ops
+        node = rewrites.column_pruning(node)
+        node, ordering = rewrites.pull_up_order(
+            node, order_root=ordered, ordered_joins=self.strict
         )
+        # final pruning to cleanup up any leftovers unused values
         node = rewrites.column_pruning(node)
-        return self.compile_node(node).to_sql(limit=n_rows, selections=ids)
+        return self.compile_node(node).to_sql(
+            order_by=ordering.all_ordering_columns if ordered else (),
+            limit=limit,
+            selections=output_ids,
+        )
 
     def compile_raw(
         self,
-        node: bigframes.core.nodes.BigFrameNode,
+        node: nodes.BigFrameNode,
     ) -> typing.Tuple[
         str, typing.Sequence[google.cloud.bigquery.SchemaField], bf_ordering.RowOrdering
     ]:
-        node = nodes.bottom_up(node, rewrites.rewrite_slice)
-        node = nodes.top_down(node, rewrites.rewrite_timedelta_expressions)
-        node, ordering = rewrites.pull_up_order(node, ordered_joins=self.strict)
+        node = self._replace_unsupported_ops(node)
+        node = rewrites.column_pruning(node)
+        node, ordering = rewrites.pull_up_order(
+            node, order_root=True, ordered_joins=self.strict
+        )
         node = rewrites.column_pruning(node)
-        ir = self.compile_node(node)
-        sql = ir.to_sql()
+        sql = self.compile_node(node).to_sql()
         return sql, node.schema.to_bigquery(), ordering
 
-    def _preprocess(self, node: nodes.BigFrameNode):
+    def _replace_unsupported_ops(self, node: nodes.BigFrameNode):
+        # TODO: Run all replacement rules as single bottom-up pass
         node = nodes.bottom_up(node, rewrites.rewrite_slice)
-        node = nodes.top_down(node, rewrites.rewrite_timedelta_expressions)
-        node, _ = rewrites.pull_up_order(
-            node, order_root=False, ordered_joins=self.strict
-        )
+        node = nodes.bottom_up(node, rewrites.rewrite_timedelta_expressions)
         return node
 
     # TODO: Remove cache when schema no longer requires compilation to derive schema (and therefor only compiles for execution)
@@ -305,16 +292,3 @@ def compile_explode(self, node: nodes.ExplodeNode):
     @_compile_node.register
     def compile_random_sample(self, node: nodes.RandomSampleNode):
         return self.compile_node(node.child)._uniform_sampling(node.fraction)
-
-
-def set_output_names(
-    node: bigframes.core.nodes.BigFrameNode, output_ids: typing.Sequence[str]
-):
-    # TODO: Create specialized output operators that will handle final names
-    return nodes.SelectionNode(
-        node,
-        tuple(
-            bigframes.core.nodes.AliasedRef(ex.DerefOp(old_id), ids.ColumnId(out_id))
-            for old_id, out_id in zip(node.ids, output_ids)
-        ),
-    )

bigframes/core/nodes.py

@@ -1555,6 +1555,14 @@ def remap_refs(
         return dataclasses.replace(self, column_ids=new_ids)  # type: ignore
 
 
+# Introduced during planing/compilation
+@dataclasses.dataclass(frozen=True, eq=False)
+class ResultNode(UnaryNode):
+    output_names: tuple[str, ...]
+    order_by: Tuple[OrderingExpression, ...] = ()
+    limit: Optional[int] = None
+
+
 # Tree operators
 def top_down(
     root: BigFrameNode,

bigframes/session/executor.py

@@ -40,6 +40,7 @@
 import pyarrow
 
 import bigframes.core
+from bigframes.core import expression
 import bigframes.core.compile
 import bigframes.core.guid
 import bigframes.core.identifiers
@@ -231,11 +232,9 @@ def to_sql(
             if enable_cache
             else array_value.node
         )
-        if ordered:
-            return self.compiler.compile_ordered(
-                node, col_id_overrides=col_id_overrides
-            )
-        return self.compiler.compile_unordered(node, col_id_overrides=col_id_overrides)
+        if col_id_overrides:
+            node = override_ids(node, col_id_overrides)
+        return self.compiler.compile(node, ordered=ordered)
 
     def execute(
         self,
@@ -377,7 +376,7 @@ def peek(
             msg = "Peeking this value cannot be done efficiently."
             warnings.warn(msg)
 
-        sql = self.compiler.compile_peek(plan, n_rows)
+        sql = self.compiler.compile(plan, ordered=False, limit=n_rows)
 
         # TODO(swast): plumb through the api_name of the user-facing api that
         # caused this query.
@@ -416,7 +415,7 @@ def head(
             assert tree_properties.can_fast_head(plan)
 
         head_plan = generate_head_plan(plan, n_rows)
-        sql = self.compiler.compile_ordered(head_plan)
+        sql = self.compiler.compile(head_plan)
 
         # TODO(swast): plumb through the api_name of the user-facing api that
         # caused this query.
@@ -439,7 +438,7 @@ def get_row_count(self, array_value: bigframes.core.ArrayValue) -> int:
             row_count_plan = self.replace_cached_subtrees(
                 generate_row_count_plan(array_value.node)
             )
-            sql = self.compiler.compile_unordered(row_count_plan)
+            sql = self.compiler.compile(row_count_plan, ordered=False)
             iter, _ = self._run_execute_query(sql)
             return next(iter)[0]
 
@@ -549,8 +548,8 @@ def _cache_with_offsets(self, array_value: bigframes.core.ArrayValue):
         """Executes the query and uses the resulting table to rewrite future executions."""
         offset_column = bigframes.core.guid.generate_guid("bigframes_offsets")
         w_offsets, offset_column = array_value.promote_offsets()
-        sql = self.compiler.compile_unordered(
-            self.replace_cached_subtrees(w_offsets.node)
+        sql = self.compiler.compile(
+            self.replace_cached_subtrees(w_offsets.node), ordered=False
         )
 
         tmp_table = self._sql_as_cached_temp_table(
@@ -666,3 +665,18 @@ def generate_head_plan(node: nodes.BigFrameNode, n: int):
 
 def generate_row_count_plan(node: nodes.BigFrameNode):
     return nodes.RowCountNode(node)
+
+
+def override_ids(
+    node: nodes.BigFrameNode, col_id_overrides: Mapping[str, str]
+) -> nodes.SelectionNode:
+    output_ids = [col_id_overrides.get(id, id) for id in node.schema.names]
+    return nodes.SelectionNode(
+        node,
+        tuple(
+            nodes.AliasedRef(
+                expression.DerefOp(old_id), bigframes.core.identifiers.ColumnId(out_id)
+            )
+            for old_id, out_id in zip(node.ids, output_ids)
+        ),
+    )