refactor: Reduce compile interface to returning just SQL text (#786)

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -27,7 +27,7 @@
 import pyarrow as pa
 import pyarrow.feather as pa_feather
 
-import bigframes.core.compile as compiling
+import bigframes.core.compile
 import bigframes.core.expression as ex
 import bigframes.core.guid
 import bigframes.core.join_def as join_def
@@ -142,12 +142,7 @@ def schema(self) -> schemata.ArraySchema:
 
     @functools.cached_property
     def _compiled_schema(self) -> schemata.ArraySchema:
-        compiled = self._compile_unordered()
-        items = tuple(
-            schemata.SchemaItem(id, compiled.get_column_type(id))
-            for id in compiled.column_ids
-        )
-        return schemata.ArraySchema(items)
+        return bigframes.core.compile.test_only_ibis_inferred_schema(self.node)
 
     def as_cached(
         self: ArrayValue,
@@ -169,21 +164,11 @@ def as_cached(
 
     def _try_evaluate_local(self):
         """Use only for unit testing paths - not fully featured. Will throw exception if fails."""
-        import ibis
-
-        return ibis.pandas.connect({}).execute(
-            self._compile_ordered()._to_ibis_expr(ordering_mode="unordered")
-        )
+        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)
 
-    def _compile_ordered(self) -> compiling.OrderedIR:
-        return compiling.compile_ordered_ir(self.node)
-
-    def _compile_unordered(self) -> compiling.UnorderedIR:
-        return compiling.compile_unordered_ir(self.node)
-
     def row_count(self) -> ArrayValue:
         """Get number of rows in ArrayValue as a single-entry ArrayValue."""
         return ArrayValue(nodes.RowCountNode(child=self.node))

bigframes/core/compile/__init__.py

@@ -11,13 +11,22 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
 
-from bigframes.core.compile.compiled import OrderedIR, UnorderedIR
-from bigframes.core.compile.compiler import compile_ordered_ir, compile_unordered_ir
+from bigframes.core.compile.api import (
+    compile_ordered,
+    compile_peek,
+    compile_raw,
+    compile_unordered,
+    test_only_ibis_inferred_schema,
+    test_only_try_evaluate,
+)
 
 __all__ = [
-    "compile_ordered_ir",
-    "compile_unordered_ir",
-    "OrderedIR",
-    "UnorderedIR",
+    "compile_peek",
+    "compile_unordered",
+    "compile_ordered",
+    "compile_raw",
+    "test_only_try_evaluate",
+    "test_only_ibis_inferred_schema",
 ]

bigframes/core/compile/api.py

@@ -0,0 +1,72 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from typing import Mapping, Tuple, TYPE_CHECKING
+
+import bigframes.core.compile.compiler as compiler
+
+if TYPE_CHECKING:
+    import bigframes.core.nodes
+    import bigframes.core.ordering
+    import bigframes.core.schema
+
+
+def compile_peek(node: bigframes.core.nodes.BigFrameNode, n_rows: int) -> str:
+    """Compile node into sql that selects N arbitrary rows, may not execute deterministically."""
+    return compiler.compile_unordered_ir(node).peek_sql(n_rows)
+
+
+def compile_unordered(
+    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."""
+    return compiler.compile_unordered_ir(node).to_sql(col_id_overrides=col_id_overrides)
+
+
+def compile_ordered(
+    node: bigframes.core.nodes.BigFrameNode, *, col_id_overrides: Mapping[str, str] = {}
+) -> str:
+    """Compile node into sql where rows are sorted with ORDER BY."""
+    return compiler.compile_ordered_ir(node).to_sql(
+        col_id_overrides=col_id_overrides, sorted=True
+    )
+
+
+def compile_raw(
+    node: bigframes.core.nodes.BigFrameNode,
+) -> Tuple[str, bigframes.core.ordering.ExpressionOrdering]:
+    """Compile node into sql that exposes all columns, including hidden ordering-only columns."""
+    ir = compiler.compile_ordered_ir(node)
+    sql = ir.raw_sql()
+    ordering_info = ir._ordering
+    return sql, ordering_info
+
+
+def test_only_try_evaluate(node: bigframes.core.nodes.BigFrameNode):
+    """Use only for unit testing paths - not fully featured. Will throw exception if fails."""
+    ibis = compiler.compile_ordered_ir(node)._to_ibis_expr(ordering_mode="unordered")
+    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
+
+    compiled = compiler.compile_unordered_ir(node)
+    items = tuple(
+        bigframes.core.schema.SchemaItem(id, compiled.get_column_type(id))
+        for id in compiled.column_ids
+    )
+    return bigframes.core.schema.ArraySchema(items)

bigframes/core/compile/compiled.py

@@ -923,6 +923,15 @@ def to_sql(
             )
         return typing.cast(str, sql)
 
+    def raw_sql(self) -> str:
+        """Return sql with all hidden columns. Used to cache with ordering information."""
+        return ibis_bigquery.Backend().compile(
+            self._to_ibis_expr(
+                ordering_mode="unordered",
+                expose_hidden_cols=True,
+            )
+        )
+
     def _to_ibis_expr(
         self,
         *,

bigframes/core/nodes.py

@@ -375,6 +375,19 @@ class CachedTableNode(BigFrameNode):
 
     ordering: typing.Optional[orderings.ExpressionOrdering] = field()
 
+    def __post_init__(self):
+        # enforce invariants
+        physical_names = set(map(lambda i: i.name, self.physical_schema))
+        logical_names = self.original_node.schema.names
+        if not set(logical_names).issubset(physical_names):
+            raise ValueError(
+                f"Requested schema {logical_names} cannot be derived from table schema {self.physical_schema}"
+            )
+        if not set(self.hidden_columns).issubset(physical_names):
+            raise ValueError(
+                f"Requested hidden columns {self.hidden_columns} cannot be derived from table schema {self.physical_schema}"
+            )
+
     @property
     def session(self):
         return self.original_node.session

bigframes/session/__init__.py

@@ -64,7 +64,6 @@
 import google.cloud.storage as storage  # type: ignore
 import ibis
 import ibis.backends.bigquery as ibis_bigquery
-import ibis.expr.types as ibis_types
 import jellyfish
 import numpy as np
 import pandas
@@ -248,6 +247,8 @@ def __init__(
         # the ibis client has been created
         original_default_query_job_config = self.bqclient.default_query_job_config
 
+        # Only used to fetch remote function metadata.
+        # TODO: Remove in favor of raw bq client
         self.ibis_client = typing.cast(
             ibis_bigquery.Backend,
             ibis.bigquery.connect(
@@ -1497,14 +1498,14 @@ def _create_empty_temp_table(
         )
         return bigquery.TableReference.from_string(table)
 
-    def _ibis_to_temp_table(
+    def _sql_to_temp_table(
         self,
-        table: ibis_types.Table,
+        sql: str,
         cluster_cols: Iterable[str],
         api_name: str,
     ) -> bigquery.TableReference:
         destination, _ = self._query_to_destination(
-            self.ibis_client.compile(table),
+            sql,
             index_cols=list(cluster_cols),
             api_name=api_name,
         )
@@ -1847,17 +1848,15 @@ def _cache_with_cluster_cols(
         # TODO: May want to support some partial ordering info even for non-strict ordering mode
         keep_order_info = self._strictly_ordered
 
-        compiled_value = self._compile_ordered(array_value)
-
-        ibis_expr = compiled_value._to_ibis_expr(
-            ordering_mode="unordered", expose_hidden_cols=keep_order_info
+        sql, ordering_info = bigframes.core.compile.compile_raw(
+            self._with_cached_executions(array_value.node)
         )
-        tmp_table = self._ibis_to_temp_table(
-            ibis_expr, cluster_cols=cluster_cols, api_name="cached"
+        tmp_table = self._sql_to_temp_table(
+            sql, cluster_cols=cluster_cols, api_name="cached"
         )
         cached_replacement = array_value.as_cached(
             cache_table=self.bqclient.get_table(tmp_table),
-            ordering=compiled_value._ordering if keep_order_info else None,
+            ordering=ordering_info if keep_order_info else None,
         ).node
         self._cached_executions[array_value.node] = cached_replacement
 
@@ -1869,13 +1868,14 @@ def _cache_with_offsets(self, array_value: core.ArrayValue):
             raise ValueError(
                 "Caching with offsets only supported in strictly ordered mode."
             )
-        compiled_value = self._compile_ordered(array_value)
-
-        ibis_expr = compiled_value._to_ibis_expr(
-            ordering_mode="offset_col", order_col_name="bigframes_offsets"
+        sql = bigframes.core.compile.compile_unordered(
+            self._with_cached_executions(
+                array_value.promote_offsets("bigframes_offsets").node
+            )
         )
-        tmp_table = self._ibis_to_temp_table(
-            ibis_expr, cluster_cols=["bigframes_offsets"], api_name="cached"
+
+        tmp_table = self._sql_to_temp_table(
+            sql, cluster_cols=["bigframes_offsets"], api_name="cached"
         )
         cached_replacement = array_value.as_cached(
             cache_table=self.bqclient.get_table(tmp_table),
@@ -1962,7 +1962,9 @@ def _peek(
         """A 'peek' efficiently accesses a small number of rows in the dataframe."""
         if not tree_properties.peekable(self._with_cached_executions(array_value.node)):
             warnings.warn("Peeking this value cannot be done efficiently.")
-        sql = self._compile_unordered(array_value).peek_sql(n_rows)
+        sql = bigframes.core.compile.compile_peek(
+            self._with_cached_executions(array_value.node), n_rows
+        )
 
         # TODO(swast): plumb through the api_name of the user-facing api that
         # caused this query.
@@ -1979,26 +1981,13 @@ def _to_sql(
     ) -> str:
         if offset_column:
             array_value = array_value.promote_offsets(offset_column)
+        node_w_cached = self._with_cached_executions(array_value.node)
         if sorted:
-            return self._compile_ordered(array_value).to_sql(
-                col_id_overrides=col_id_overrides, sorted=True
+            return bigframes.core.compile.compile_ordered(
+                node_w_cached, col_id_overrides=col_id_overrides
             )
-        return self._compile_unordered(array_value).to_sql(
-            col_id_overrides=col_id_overrides
-        )
-
-    def _compile_ordered(
-        self, array_value: core.ArrayValue
-    ) -> bigframes.core.compile.OrderedIR:
-        return bigframes.core.compile.compile_ordered_ir(
-            self._with_cached_executions(array_value.node)
-        )
-
-    def _compile_unordered(
-        self, array_value: core.ArrayValue
-    ) -> bigframes.core.compile.UnorderedIR:
-        return bigframes.core.compile.compile_unordered_ir(
-            self._with_cached_executions(array_value.node)
+        return bigframes.core.compile.compile_unordered(
+            node_w_cached, col_id_overrides=col_id_overrides
         )
 
     def _get_table_size(self, destination_table):