fix issues

Author: TrevorBergeron

packages/bigframes/bigframes/core/rewrite/__init__.py

@@ -19,7 +19,7 @@
 from bigframes.core.rewrite.implicit_align import try_row_join
 from bigframes.core.rewrite.legacy_align import legacy_join_as_projection
 from bigframes.core.rewrite.nullity import simplify_join
-from bigframes.core.rewrite.order import bake_order, defer_order
+from bigframes.core.rewrite.order import bake_order, defer_order, pull_out_order
 from bigframes.core.rewrite.pruning import column_pruning
 from bigframes.core.rewrite.scan_reduction import (
     try_reduce_to_local_scan,
@@ -50,6 +50,7 @@
     "rewrite_range_rolling",
     "try_reduce_to_table_scan",
     "bake_order",
+    "pull_out_order",
     "try_reduce_to_local_scan",
     "fold_row_counts",
     "pull_out_window_order",

packages/bigframes/bigframes/core/rewrite/order.py

@@ -47,6 +47,12 @@ def bake_order(
     return node
 
 
+def pull_out_order(
+    node: bigframes.core.nodes.BigFrameNode,
+) -> Tuple[bigframes.core.nodes.BigFrameNode, bigframes.core.ordering.RowOrdering]:
+    return _pull_up_order(node, order_root=False)
+
+
 # Makes ordering explicit in window definitions
 def _pull_up_order(
     root: bigframes.core.nodes.BigFrameNode,

packages/bigframes/bigframes/session/bq_caching_executor.py

@@ -40,7 +40,7 @@
 import bigframes.session.metrics
 import bigframes.session.planner
 import bigframes.session.temporary_storage
-from bigframes.core import compile, guid, local_data, rewrite
+from bigframes.core import bq_data, compile, guid, local_data, rewrite
 from bigframes.core.compile.sqlglot import sql as sg_sql
 from bigframes.core.compile.sqlglot import sqlglot_ir
 from bigframes.session import (
@@ -176,59 +176,61 @@ def _execute_bigquery(
     ) -> executor.ExecuteResult:
         dest_spec = execution_spec.destination_spec
         # Recursive handlers for different cases, maybe extract to explicit interface.
+        if (
+            isinstance(dest_spec, ex_spec.EphemeralTableSpec)
+            and not execution_spec.promise_under_10gb
+        ):
+            # Results over 10GB need to explicitly allocate a table.
+            execution_spec = dataclasses.replace(
+                execution_spec, destination_spec=ex_spec.SessionTableSpec()
+            )
+            return self._execute_bigquery(array_value, execution_spec)
         if isinstance(dest_spec, ex_spec.GcsOutputSpec):
             execution_spec = dataclasses.replace(
-                execution_spec,
-                destination_spec=ex_spec.TempTableSpec(
-                    cluster_cols=dest_spec.cluster_cols, lifetime="ephemeral"
-                ),
+                execution_spec, destination_spec=ex_spec.EphemeralTableSpec()
             )
             results = self._execute_bigquery(array_value, execution_spec)
             self._export_result_gcs(results, dest_spec)
             return results
-        if isinstance(dest_spec, ex_spec.TableOutputSpec):
+        if isinstance(dest_spec, ex_spec.TableOutputSpec) and dest_spec.permit_dml:
             # Special DML path - maybe this should be configurable, dml vs query destination has tradeoffs
             existing_table = self._maybe_find_existing_table(dest_spec)
-            if execution_spec.ordered:
-                raise ValueError("Ordering not supported with table outputs")
             if (existing_table is not None) and _is_schema_match(
                 existing_table.schema, array_value.schema
             ):
                 execution_spec = dataclasses.replace(
-                    execution_spec,
-                    destination_spec=ex_spec.TempTableSpec(
-                        cluster_cols=execution_spec.destination_spec.cluster_cols,
-                        lifetime="ephemeral",
-                    ),
+                    execution_spec, destination_spec=ex_spec.EphemeralTableSpec()
                 )
                 results = self._execute_bigquery(array_value, execution_spec)
                 self._export_gbq_with_dml(results, dest_spec)
                 return results
-        if isinstance(dest_spec, ex_spec.TempTableSpec):
+        if isinstance(dest_spec, ex_spec.SessionTableSpec):
             # "ephemeral" temp tables created in the course of exeuction, don't need to be allocated
             # materialized ordering only really makes sense for internal temp tables used by caching
             cluster_cols = dest_spec.cluster_cols
+            # Rewrite plan to materialize ordering as extra columns
             plan = array_value.node
             if dest_spec.ordering == "offsets_col":
                 order_col_id = guid.generate_guid()
                 plan = nodes.PromoteOffsetsNode(plan, order_col_id)
                 cluster_cols = [order_col_id]
             elif dest_spec.ordering == "order_key":
-                plan = nodes.defer_order(plan, output_hidden_row_keys=True)
-            if dest_spec.lifetime == "session":
-                destination_table = self.storage_manager.create_temp_table(
-                    plan.schema, cluster_cols
-                )
-                arr_value = bigframes.core.ArrayValue(plan)
-                execution_spec = dataclasses.replace(
-                    execution_spec,
-                    destination_spec=ex_spec.TableOutputSpec(
-                        table=destination_table,
-                        cluster_cols=dest_spec.cluster_cols,
-                        if_exists="replace",
-                    ),
-                )
-                return self._execute_bigquery(arr_value, execution_spec)
+                plan, _ = rewrite.pull_out_order(plan)
+            destination_table = self.storage_manager.create_temp_table(
+                plan.schema.to_bigquery(), cluster_cols
+            )
+            arr_value = bigframes.core.ArrayValue(plan)
+            execution_spec = dataclasses.replace(
+                execution_spec,
+                destination_spec=ex_spec.TableOutputSpec(
+                    table=destination_table,
+                    cluster_cols=dest_spec.cluster_cols,
+                    if_exists="replace",
+                    # Avoid loops, also dml is mostly used to avoid quotas on user-owned tables
+                    permit_dml=False,
+                ),
+            )
+            return self._execute_bigquery(arr_value, execution_spec)
 
         # At this point, dst should be unspecified, a specific bq table, or an ephemeral temp table
         # Also, ordering mode will either be none or row-sorted
@@ -405,32 +407,28 @@ def _cache_with_cluster_cols(
         ]
         cluster_cols = cluster_cols[:_MAX_CLUSTER_COLUMNS]
         execution_spec = ex_spec.ExecutionSpec(
-            destination_spec=ex_spec.TempTableSpec(
-                cluster_cols=tuple(cluster_cols),
-                lifetime="session",
-                ordering="order_key",
-            )
+            destination_spec=ex_spec.SessionTableSpec(cluster_cols=tuple(cluster_cols))
         )
-        result_bq_data = self.execute(
+        result = self.execute(
             array_value,
             execution_spec=execution_spec,
         )
-        assert isinstance(result_bq_data, bigframes.core.BigqueryDataSource)
-        self.cache.cache_results_table(array_value.node, result_bq_data)
+        assert isinstance(result, executor.BQTableExecuteResult)
+        self.cache.cache_results_table(array_value.node, result._data)
 
     def _cache_with_offsets(self, array_value: bigframes.core.ArrayValue):
         """Executes the query and uses the resulting table to rewrite future executions."""
         execution_spec = ex_spec.ExecutionSpec(
-            destination_spec=ex_spec.TempTableSpec(
-                cluster_cols=(), lifetime="session", ordering="offsets_col"
+            destination_spec=ex_spec.SessionTableSpec(
+                cluster_cols=(), ordering="offsets_col"
             )
         )
-        result_bq_data = self.execute(
+        result = self.execute(
             array_value,
             execution_spec=execution_spec,
         )
-        assert isinstance(result_bq_data, bigframes.core.BigqueryDataSource)
-        self.cache.cache_results_table(array_value.node, result_bq_data)
+        assert isinstance(result, executor.BQTableExecuteResult)
+        self.cache.cache_results_table(array_value.node, result._data)
 
     def _cache_with_session_awareness(
         self,

packages/bigframes/bigframes/session/direct_gbq_execution.py

@@ -87,25 +87,25 @@ def execute(
 
         job_config = bigquery.QueryJobConfig()
         dest_spec = spec.destination_spec
-        cluster_cols = ()
+        cluster_cols = None
+        can_skip_job = True
         if isinstance(dest_spec, execution_spec.TableOutputSpec):
+            if spec.ordered:
+                raise ValueError("Ordering not supported with destination table")
             job_config.destination = dest_spec.table
             job_config.write_disposition = _WRITE_DISPOSITIONS[dest_spec.if_exists]
-            cluster_cols = dest_spec.cluster_cols
-            job_config.clustering_fields = dest_spec.cluster_cols
-        elif (
-            isinstance(dest_spec, execution_spec.TempTableSpec)
-            and dest_spec.lifetime == "ephemeral"
-        ):
-            cluster_cols = dest_spec.cluster_cols
-            job_config.clustering_fields = dest_spec.cluster_cols
+            cluster_cols = dest_spec.cluster_cols if dest_spec.cluster_cols else None
+            job_config.clustering_fields = cluster_cols
+            can_skip_job = False
+        elif isinstance(dest_spec, execution_spec.EphemeralTableSpec):
+            # Need destination table, but jobless execution might not create a destination table
+            can_skip_job = False
         elif dest_spec is not None:
             raise ValueError(
                 f"Direct GBQ Executor does not support destination: {dest_spec}"
             )
 
         job_config.labels["bigframes-dtypes"] = compiled.encoded_type_refs
-        can_skip_job = dest_spec is None and spec.promise_under_10gb
         iterator, query_job = self._run_execute_query(
             sql=compiled.sql,
             job_config=job_config,
@@ -121,7 +121,7 @@ def execute(
                 table=bq_data.GbqNativeTable.from_ref_and_schema(
                     dst,
                     tuple(compiled_schema),
-                    cluster_cols=cluster_cols,
+                    cluster_cols=cluster_cols or (),
                     location=iterator.location or self.bqclient.location,
                     table_type="TABLE",
                 ),
@@ -137,7 +137,10 @@ def execute(
             hasattr(iterator, "_is_almost_completely_cached")
             and iterator._is_almost_completely_cached()
         )
-        if result_bq_data is not None and not result_mostly_cached:
+
+        if (isinstance(dest_spec, execution_spec.EphemeralTableSpec)) or (
+            result_bq_data is not None and not result_mostly_cached
+        ):
             return executor.BQTableExecuteResult(
                 data=result_bq_data,
                 project_id=self.bqclient.project,

packages/bigframes/bigframes/session/execution_spec.py

@@ -23,7 +23,9 @@
 @dataclasses.dataclass(frozen=True)
 class ExecutionSpec:
     # 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
+    destination_spec: Union[
+        TableOutputSpec, GcsOutputSpec, EphemeralTableSpec, SessionTableSpec, 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
@@ -36,16 +38,28 @@ class ExecutionSpec:
 
 # Used internally by execution
 @dataclasses.dataclass(frozen=True)
-class TempTableSpec:
+class EphemeralTableSpec:
+    """
+    Specifies that the result of an operation should be a temporary table of some sort.
+
+    No guarantees on lifetime, may be a session temp table, or a bq-created temp table with <24hr life.
+
+    Used internally when results need temporary staging, because they are large (>10GB), or needed in subsequent operations.
+    """
+
+    pass
+
+
+@dataclasses.dataclass(frozen=True)
+class SessionTableSpec:
     """
     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"
+    ] = ()  # if empty, will cluster using order key if ordering_key is set
     # 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)
@@ -66,6 +80,8 @@ class TableOutputSpec:
     table: bigquery.TableReference
     cluster_cols: tuple[str, ...]
     if_exists: Literal["fail", "replace", "append"] = "fail"
+    # Allow DML to be used to populate table
+    permit_dml: bool = True
 
 
 @dataclasses.dataclass(frozen=True)

packages/bigframes/tests/system/small/engines/conftest.py

@@ -44,31 +44,37 @@ def fake_session() -> Generator[bigframes.Session, None, None]:
     with bigframes.core.global_session._GlobalSessionContext(session):
         yield session
 
+@pytest.fixture(scope="session")
+def sqlglot_engine(bigquery_client: bigquery.Client, bigquery_storage_read_client: google.cloud.bigquery_storage_v1.BigQueryReadClient):
+    return direct_gbq_execution.DirectGbqExecutor(
+        bigquery_client,
+        bqstoragereadclient=bigquery_storage_read_client,
+        compiler="sqlglot",
+        publisher=events.Publisher(),
+    )
+
+@pytest.fixture(scope="session")
+def bq_engine(bigquery_client: bigquery.Client, bigquery_storage_read_client: google.cloud.bigquery_storage_v1.BigQueryReadClient):
+    return direct_gbq_execution.DirectGbqExecutor(
+        bigquery_client,
+        bqstoragereadclient=bigquery_storage_read_client,
+        publisher=events.Publisher(),
+    )
 
 @pytest.fixture(scope="session", params=["pyarrow", "polars", "bq", "bq-sqlglot"])
 def engine(
     request,
-    bigquery_client: bigquery.Client,
-    bigquery_storage_read_client: google.cloud.bigquery_storage_v1.BigQueryReadClient,
+    sqlglot_engine,
+    bq_engine,
 ) -> semi_executor.SemiExecutor:
     if request.param == "pyarrow":
         return local_scan_executor.LocalScanExecutor()
     if request.param == "polars":
         return polars_executor.PolarsExecutor()
-    publisher = events.Publisher()
     if request.param == "bq":
-        return direct_gbq_execution.DirectGbqExecutor(
-            bigquery_client,
-            bqstoragereadclient=bigquery_storage_read_client,
-            publisher=publisher,
-        )
+        return bq_engine
     if request.param == "bq-sqlglot":
-        return direct_gbq_execution.DirectGbqExecutor(
-            bigquery_client,
-            bqstoragereadclient=bigquery_storage_read_client,
-            compiler="sqlglot",
-            publisher=publisher,
-        )
+        return sqlglot_engine
     raise ValueError(f"Unrecognized param: {request.param}")
 
 

packages/bigframes/tests/system/small/engines/test_windowing.py

@@ -46,8 +46,9 @@ def test_engines_with_offsets(
 @pytest.mark.parametrize("agg_op", [agg_ops.sum_op, agg_ops.count_op])
 def test_engines_with_rows_window(
     scalars_array_value: array_value.ArrayValue,
-    bigquery_client: bigquery.Client,
     agg_op,
+    bq_engine,
+    sqlglot_engine,
 ):
     window = window_spec.WindowSpec(
         bounds=window_spec.RowsWindowBounds.from_window_size(3, "left"),
@@ -62,12 +63,4 @@ def test_engines_with_rows_window(
         ),
         window_spec=window,
     )
-
-    publisher = events.Publisher()
-    bq_executor = direct_gbq_execution.DirectGbqExecutor(
-        bigquery_client, publisher=publisher
-    )
-    bq_sqlgot_executor = direct_gbq_execution.DirectGbqExecutor(
-        bigquery_client, compiler="sqlglot", publisher=publisher
-    )
-    assert_equivalence_execution(window_node, bq_executor, bq_sqlgot_executor)
+    assert_equivalence_execution(window_node, bq_engine, sqlglot_engine)

packages/bigframes/tests/system/small/test_session.py

@@ -118,7 +118,7 @@ def test_read_gbq_tokyo(
     exec_result = session_tokyo._executor.execute(
         df._block.expr,
         bigframes.session.execution_spec.ExecutionSpec(
-            bigframes.session.execution_spec.TempTableSpec(()), promise_under_10gb=False
+            bigframes.session.execution_spec.EphemeralTableSpec(), promise_under_10gb=False
         ),
     )
     assert exec_result.query_job is not None
@@ -951,7 +951,7 @@ def test_read_pandas_tokyo(
     result = session_tokyo._executor.execute(
         df._block.expr,
         bigframes.session.execution_spec.ExecutionSpec(
-            bigframes.session.execution_spec.TempTableSpec(()), promise_under_10gb=False
+            bigframes.session.execution_spec.EphemeralTableSpec(), promise_under_10gb=False
         ),
     )
     assert result.query_job is not None