refactor: Unify gbq execution as semi-executor

TrevorBergeron · TrevorBergeron · commit 076cb98af458 · 2026-04-22T22:31:37.000Z
diff --git a/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py b/packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py
@@ -177,7 +177,7 @@ def from_table(
         project_id: str,
         dataset_id: str,
         table_id: str,
-        uid_gen: guid.SequentialUIDGenerator,
+        uid_gen: guid.SequentialUIDGenerator | None = None,
         columns: typing.Sequence[str] = (),
         sql_predicate: typing.Optional[str] = None,
         system_time: typing.Optional[datetime.datetime] = None,
@@ -202,6 +202,8 @@ def from_table(
             if system_time
             else None
         )
+        if uid_gen is None:
+            uid_gen = guid.SequentialUIDGenerator()
         table_alias = next(uid_gen.get_uid_stream("bft_"))
         table_expr = sge.Table(
             this=sql.identifier(table_id),
diff --git a/packages/bigframes/bigframes/session/bq_caching_executor.py b/packages/bigframes/bigframes/session/bq_caching_executor.py
diff --git a/packages/bigframes/bigframes/session/direct_gbq_execution.py b/packages/bigframes/bigframes/session/direct_gbq_execution.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import Literal, Optional, Tuple
+from typing import Callable, Literal, Optional, Tuple
 
 import google.cloud.bigquery.job as bq_job
 import google.cloud.bigquery.table as bq_table
@@ -23,76 +23,180 @@
 import bigframes.core.compile.sqlglot.compiler as sqlglot_compiler
 import bigframes.core.events
 import bigframes.session._io.bigquery as bq_io
-from bigframes.core import compile, nodes
-from bigframes.session import executor, semi_executor
+from bigframes.core import bq_data, compile, nodes
+from bigframes.session import executor, semi_executor, execution_spec
+from bigframes import exceptions as bfe
+import bigframes.core.schema as schemata
+
+import google.api_core.exceptions
+
+
+_WRITE_DISPOSITIONS = {
+    "fail": bigquery.WriteDisposition.WRITE_EMPTY,
+    "replace": bigquery.WriteDisposition.WRITE_TRUNCATE,
+    "append": bigquery.WriteDisposition.WRITE_APPEND,
+}
 
 
-# used only in testing right now, BigQueryCachingExecutor is the fully featured engine
-# simplified, doesnt not do large >10 gb result queries, error handling, respect global config
-# or record metrics. Also avoids caching, and most pre-compile rewrites, to better serve as a
-# reference for validating more complex executors.
 class DirectGbqExecutor(semi_executor.SemiExecutor):
     def __init__(
         self,
         bqclient: bigquery.Client,
-        compiler: Literal["ibis", "sqlglot"] = "ibis",
+        compiler: Literal["ibis", "sqlglot"]
+        | Callable[[compile.CompileRequest], executor.CompileResult] = "ibis",
         *,
         publisher: bigframes.core.events.Publisher,
     ):
         self.bqclient = bqclient
-        self._compile_fn = (
-            ibis_compiler.compile_sql
-            if compiler == "ibis"
-            else sqlglot_compiler.compile_sql
-        )
+        if isinstance(compiler, str):
+            self._compile_fn = (
+                ibis_compiler.compile_sql
+                if compiler == "ibis"
+                else sqlglot_compiler.compile_sql
+            )
+        else:
+            self._compile_fn = compiler
         self._publisher = publisher
 
     def execute(
         self,
         plan: nodes.BigFrameNode,
-        ordered: bool,
-        peek: Optional[int] = None,
+        spec: execution_spec.ExecutionSpec,
     ) -> executor.ExecuteResult:
         """Just execute whatever plan as is, without further caching or decomposition."""
-        # TODO(swast): plumb through the api_name of the user-facing api that
-        # caused this query.
 
-        compiled = self._compile_fn(
-            compile.CompileRequest(plan, sort_rows=ordered, peek_count=peek)
+        og_schema = plan.schema
+
+        compiled = compile.compiler().compile_sql(
+            compile.CompileRequest(
+                plan,
+                sort_rows=spec.ordered,
+                peek_count=spec.peek,
+            )
         )
+        # might have more columns than og schema, for hidden ordering columns
+        compiled_schema = compiled.sql_schema
+
+        job_config = bigquery.QueryJobConfig()
+        if isinstance(spec.destination_spec, execution_spec.TableOutputSpec):
+            job_config.destination = spec.destination_spec.table
+            job_config.write_disposition = _WRITE_DISPOSITIONS[spec.destination_spec.if_exists]
+            job_config.clustering_fields = spec.destination_spec.cluster_cols
+        elif isinstance(spec.destination_spec, execution_spec.TempTableSpec) and spec.destination_spec.lifetime == "ephemeral":
+            pass
+        elif spec.destination_spec is not None:
+            raise ValueError(f"Direct GBQ Executor does not support destination: {spec.destination_spec}")
 
+        job_config.labels["bigframes-dtypes"] = compiled.encoded_type_refs
+        can_skip_job = spec.destination_spec is None and spec.promise_under_10gb
         iterator, query_job = self._run_execute_query(
             sql=compiled.sql,
+            job_config=job_config,
+            query_with_job=(not can_skip_job),
             session=plan.session,
         )
 
-        # just immediately downlaod everything for simplicity
-        return executor.LocalExecuteResult(
-            data=iterator.to_arrow(),
-            bf_schema=plan.schema,
-            execution_metadata=executor.ExecutionMetadata.from_iterator_and_job(
-                iterator, query_job
-            ),
+        result_bq_data = None
+        if query_job and query_job.destination:
+            # we might add extra sql columns in compilation, esp if caching w ordering, infer a bigframes type for them
+            result_bf_schema = _result_schema(og_schema, list(compiled.sql_schema))
+            dst = query_job.destination
+            result_bq_data = bq_data.BigqueryDataSource(
+                table=bq_data.GbqNativeTable.from_ref_and_schema(
+                    dst,
+                    tuple(compiled_schema),
+                    cluster_cols=spec.destination_spec.cluster_cols,
+                    location=iterator.location or self.storage_manager.location,
+                    table_type="TABLE",
+                ),
+                schema=result_bf_schema,
+                ordering=compiled.row_order,
+                n_rows=iterator.total_rows,
+            )
+
+        execution_metadata = executor.ExecutionMetadata.from_iterator_and_job(
+            iterator, query_job
         )
+        result_mostly_cached = (
+            hasattr(iterator, "_is_almost_completely_cached")
+            and iterator._is_almost_completely_cached()
+        )
+        if result_bq_data is not None and not result_mostly_cached:
+            return executor.BQTableExecuteResult(
+                data=result_bq_data,
+                project_id=self.bqclient.project,
+                storage_client=self.bqstoragereadclient,
+                execution_metadata=execution_metadata,
+                selected_fields=tuple((col, col) for col in og_schema.names),
+            )
+        else:
+            return executor.LocalExecuteResult(
+                data=iterator.to_arrow().select(og_schema.names),
+                bf_schema=plan.schema,
+                execution_metadata=execution_metadata,
+            )
 
     def _run_execute_query(
         self,
         sql: str,
         job_config: Optional[bq_job.QueryJobConfig] = None,
+        query_with_job: bool = True,
         session=None,
     ) -> Tuple[bq_table.RowIterator, Optional[bigquery.QueryJob]]:
         """
         Starts BigQuery query job and waits for results.
         """
-        return bq_io.start_query_with_client(
-            self.bqclient,
-            sql,
-            job_config=job_config or bq_job.QueryJobConfig(),
-            project=None,
-            location=None,
-            timeout=None,
-            metrics=None,
-            query_with_job=False,
-            publisher=self._publisher,
-            session=session,
-        )
+        job_config = bq_job.QueryJobConfig() if job_config is None else job_config
+        if bigframes.options.compute.maximum_bytes_billed is not None:
+            job_config.maximum_bytes_billed = (
+                bigframes.options.compute.maximum_bytes_billed
+            )
+
+        if self._labels:
+            job_config.labels.update(self._labels)
+
+        try:
+            # Trick the type checker into thinking we got a literal.
+            if query_with_job:
+                return bq_io.start_query_with_client(
+                    self.bqclient,
+                    sql,
+                    job_config=job_config,
+                    metrics=self.metrics,
+                    project=None,
+                    location=None,
+                    timeout=None,
+                    query_with_job=True,
+                    publisher=self._publisher,
+                    session=session,
+                )
+            else:
+                return bq_io.start_query_with_client(
+                    self.bqclient,
+                    sql,
+                    job_config=job_config,
+                    metrics=self.metrics,
+                    project=None,
+                    location=None,
+                    timeout=None,
+                    query_with_job=False,
+                    publisher=self._publisher,
+                    session=session,
+                )
+
+        except google.api_core.exceptions.BadRequest as e:
+            # Unfortunately, this error type does not have a separate error code or exception type
+            if "Resources exceeded during query execution" in e.message:
+                new_message = "Computation is too complex to execute as a single query. Try using DataFrame.cache() on intermediate results, or setting bigframes.options.compute.enable_multi_query_execution."
+                raise bfe.QueryComplexityError(new_message) from e
+            else:
+                raise
+
+def _result_schema(
+    logical_schema: schemata.ArraySchema, sql_schema: list[bigquery.SchemaField]
+) -> schemata.ArraySchema:
+    inferred_schema = bigframes.dtypes.bf_type_from_type_kind(sql_schema)
+    inferred_schema.update(logical_schema._mapping)
+    return schemata.ArraySchema(
+        tuple(schemata.SchemaItem(col, dtype) for col, dtype in inferred_schema.items())
+    )
diff --git a/packages/bigframes/bigframes/session/execution_spec.py b/packages/bigframes/bigframes/session/execution_spec.py
@@ -22,24 +22,43 @@
 
 @dataclasses.dataclass(frozen=True)
 class ExecutionSpec:
-    destination_spec: Union[TableOutputSpec, GcsOutputSpec, CacheSpec, None] = None
+    # destination for the result of the operation. Executor may also incidentally create other temporary tables for its own purposes.
+    destination_spec: Union[TableOutputSpec, GcsOutputSpec, TempTableSpec, None] = None
+    # If set, the result will be truncated to the given number of rows. Which N rows is
+    # implementation dependent and not stable.
     peek: Optional[int] = None
-    ordered: bool = (
-        False  # ordered and promise_under_10gb must both be together for bq execution
-    )
+    # Controls whether output iterator is ordered. Cannot be true if destination is not
+    # guaranteed to be ordered. 
+    ordered: bool = False
     # This is an optimization flag for gbq execution, it doesn't change semantics, but if promise is falsely made, errors may occur
     promise_under_10gb: bool = False
 
 
-# This one is temporary, in future, caching will not be done through immediate execution, but will label nodes
-# that will be cached only when a super-tree is executed
+# Used internally by execution
 @dataclasses.dataclass(frozen=True)
-class CacheSpec:
-    cluster_cols: tuple[str, ...]
+class TempTableSpec:
+    """
+    Specifies that the result of an operation should be a session temp table.
+    The table will be automatically deleted after the session ends.
+    """
+    cluster_cols: tuple[str, ...]   # if empty, will cluster using order key if ordering_key is set
+    lifetime: Literal["session", "ephemeral"] = "session"
+    # Controls ordering and whether extra columns are materialized to preserve ordering
+    # Any extra columns will be appended to the end of the schema.
+    # None: ordering may be discarded entirely (ordering metadata will still be provided if ordering is derivable from materialized columns)
+    # order_rows: the result iterator itself will be ordered. For gbq execution, result cannot exceed 10GB.
+    # order_key: the result set ordered by a key, may materialize extra columns.
+    # offsets_col: order the result set by an offsets column, materializes one extra column.
+    ordering: Literal["order_rows", "offsets_col", "order_key"] | None = None
 
 
 @dataclasses.dataclass(frozen=True)
 class TableOutputSpec:
+    """
+    Specifies that the result of an operation should be exported to a specific named table.
+
+    The executor is not responsible for managing lifecycle of the table.
+    """
     table: bigquery.TableReference
     cluster_cols: tuple[str, ...]
     if_exists: Literal["fail", "replace", "append"] = "fail"
diff --git a/packages/bigframes/bigframes/session/local_scan_executor.py b/packages/bigframes/bigframes/session/local_scan_executor.py
@@ -16,7 +16,7 @@
 from typing import Optional
 
 from bigframes.core import bigframe_node, rewrite
-from bigframes.session import executor, semi_executor
+from bigframes.session import executor, semi_executor, execution_spec
 
 
 class LocalScanExecutor(semi_executor.SemiExecutor):
@@ -27,15 +27,17 @@ class LocalScanExecutor(semi_executor.SemiExecutor):
     def execute(
         self,
         plan: bigframe_node.BigFrameNode,
-        ordered: bool,
-        peek: Optional[int] = None,
+        execution_spec: execution_spec.ExecutionSpec,
     ) -> Optional[executor.ExecuteResult]:
+        if execution_spec.destination_spec is not None:
+            return None
+
         reduced_result = rewrite.try_reduce_to_local_scan(plan)
         if not reduced_result:
             return None
 
         node, limit = reduced_result
-
+        peek = execution_spec.peek
         if limit is not None:
             if peek is None or limit < peek:
                 peek = limit
diff --git a/packages/bigframes/bigframes/session/polars_executor.py b/packages/bigframes/bigframes/session/polars_executor.py
@@ -34,7 +34,7 @@
     numeric_ops,
     string_ops,
 )
-from bigframes.session import executor, semi_executor
+from bigframes.session import executor, semi_executor, execution_spec
 
 if TYPE_CHECKING:
     import polars as pl
@@ -140,20 +140,20 @@ def __init__(self):
     def execute(
         self,
         plan: bigframe_node.BigFrameNode,
-        ordered: bool,
-        peek: Optional[int] = None,
+        execution_spec: execution_spec.ExecutionSpec,
     ) -> Optional[executor.ExecuteResult]:
         if not self._can_execute(plan):
             return None
-        # Note: Ignoring ordered flag, as just executing totally ordered is fine.
+        if execution_spec.destination_spec is not None:
+            return None
         try:
             lazy_frame: pl.LazyFrame = self._compiler.compile(
                 array_value.ArrayValue(plan).node
             )
         except Exception:
             return None
-        if peek is not None:
-            lazy_frame = lazy_frame.limit(peek)
+        if execution_spec.peek is not None:
+            lazy_frame = lazy_frame.limit(execution_spec.peek)
         pa_table = lazy_frame.collect().to_arrow()
         return executor.LocalExecuteResult(
             data=pa_table,
diff --git a/packages/bigframes/bigframes/session/read_api_execution.py b/packages/bigframes/bigframes/session/read_api_execution.py
@@ -18,7 +18,7 @@
 from google.cloud import bigquery_storage_v1
 
 from bigframes.core import bigframe_node, bq_data, nodes, rewrite
-from bigframes.session import executor, semi_executor
+from bigframes.session import executor, semi_executor, execution_spec
 
 
 class ReadApiSemiExecutor(semi_executor.SemiExecutor):
@@ -37,14 +37,16 @@ def __init__(
     def execute(
         self,
         plan: bigframe_node.BigFrameNode,
-        ordered: bool,
-        peek: Optional[int] = None,
+        execution_spec: execution_spec.ExecutionSpec,
     ) -> Optional[executor.ExecuteResult]:
-        adapt_result = self._try_adapt_plan(plan, ordered)
+        if execution_spec.destination_spec is not None:
+            return None
+
+        adapt_result = self._try_adapt_plan(plan, execution_spec.ordered)
         if not adapt_result:
             return None
         node, limit = adapt_result
-        if node.explicitly_ordered and ordered:
+        if node.explicitly_ordered and execution_spec.ordered:
             return None
 
         if not isinstance(node.source.table, bq_data.GbqNativeTable):
@@ -53,6 +55,7 @@ def execute(
         if not node.source.table.is_physically_stored:
             return None
 
+        peek = execution_spec.peek
         if limit is not None:
             if peek is None or limit < peek:
                 peek = limit
diff --git a/packages/bigframes/bigframes/session/semi_executor.py b/packages/bigframes/bigframes/session/semi_executor.py
diff --git a/packages/bigframes/tests/system/small/test_session.py b/packages/bigframes/tests/system/small/test_session.py
