fixes

Author: TrevorBergeron

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

@@ -51,6 +51,7 @@ def pull_out_order(
     node: bigframes.core.nodes.BigFrameNode,
 ) -> Tuple[bigframes.core.nodes.BigFrameNode, bigframes.core.ordering.RowOrdering]:
     import bigframes.core.rewrite.slices
+
     node = node.bottom_up(bigframes.core.rewrite.slices.rewrite_slice)
     return _pull_up_order(node, order_root=True)
 

packages/bigframes/bigframes/session/bq_caching_executor.py

@@ -347,19 +347,19 @@ def _execute_to_cached_table(
             order_col_id = guid.generate_guid()
             plan = nodes.PromoteOffsetsNode(plan, identifiers.ColumnId(order_col_id))
             cluster_cols = [order_col_id]
+            ordering = bigframes.core.ordering.TotalOrdering.from_offset_col(order_col_id)
         elif cache_spec.ordering == "order_key":
             plan, ordering = 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,
+        execution_spec = ex_spec.ExecutionSpec(
             destination_spec=ex_spec.TableOutputSpec(
                 table=destination_table,
-                cluster_cols=cache_spec.cluster_cols,
+                cluster_cols=cluster_cols,
                 if_exists="replace",
-            ),
+            )
         )
         # We don't use _execute_gbq_table_output, because we want to skip slower DML path.
         result = self._execute_gbq_query_only(arr_value, execution_spec)
@@ -427,7 +427,7 @@ def _cache_with_offsets(self, array_value: bigframes.core.ArrayValue):
         """Executes the query and uses the resulting table to rewrite future executions."""
         result = self._execute_to_cached_table(
             array_value.node,
-            ex_spec.CacheSpec(cluster_cols=tuple(cluster_cols), ordering="order_key"),
+            ex_spec.CacheSpec(ordering="offsets_col"),
         )
         assert isinstance(result, executor.BQTableExecuteResult)
         assert result._data.ordering is not None

packages/bigframes/bigframes/session/direct_gbq_execution.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import Callable, Literal, Optional, Tuple
+from typing import Callable, Literal, Mapping, Optional, Tuple
 
 import google.cloud.bigquery.job as bq_job
 import google.cloud.bigquery.table as bq_table
@@ -135,8 +135,9 @@ def execute(
         )
 
         if (isinstance(dest_spec, execution_spec.EphemeralTableSpec)) or (
-            result_bq_data is not None and not result_mostly_cached
+            (result_bq_data is not None) and not result_mostly_cached
         ):
+            assert result_bq_data is not None, "expected result table but none exists"
             return executor.BQTableExecuteResult(
                 data=result_bq_data,
                 project_id=self.bqclient.project,

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

@@ -14,12 +14,12 @@
 import pathlib
 from typing import Generator
 
+import google.cloud.bigquery_storage_v1
 import pandas as pd
 import pytest
 from google.cloud import bigquery
 
 import bigframes
-import google.cloud.bigquery_storage_v1
 from bigframes.core import ArrayValue, events, local_data
 from bigframes.session import (
     direct_gbq_execution,
@@ -44,23 +44,32 @@ 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):
+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):
+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,