more fixes

Author: TrevorBergeron

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

@@ -50,7 +50,9 @@ def bake_order(
 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)
+    import bigframes.core.rewrite.slices
+    node = node.bottom_up(bigframes.core.rewrite.slices.rewrite_slice)
+    return _pull_up_order(node, order_root=True)
 
 
 # Makes ordering explicit in window definitions
@@ -273,7 +275,7 @@ def pull_up_order_inner(
                 offsets_id
             )
             return new_explode, child_order.join(inner_order)
-        raise ValueError(f"Unexpected node: {node}")
+        raise ValueError(f"Unexpected node type {type(node).__name__}")
 
     def pull_order_concat(
         node: bigframes.core.nodes.ConcatNode,

packages/bigframes/bigframes/session/bq_caching_executor.py

@@ -28,7 +28,6 @@
 import bigframes.core
 import bigframes.core.events
 import bigframes.core.guid
-import bigframes.core.identifiers
 import bigframes.core.ordering
 import bigframes.core.nodes as nodes
 import bigframes.core.schema as schemata
@@ -40,7 +39,7 @@
 import bigframes.session.metrics
 import bigframes.session.planner
 import bigframes.session.temporary_storage
-from bigframes.core import bq_data, compile, guid, local_data, rewrite
+from bigframes.core import bq_data, compile, guid, identifiers, 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,23 +175,14 @@ 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.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) and dest_spec.permit_dml:
+        elif 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 (existing_table is not None) and _is_schema_match(
@@ -204,18 +194,18 @@ def _execute_bigquery(
                 results = self._execute_bigquery(array_value, execution_spec)
                 self._export_gbq_with_dml(results, dest_spec)
                 return results
-        if isinstance(dest_spec, ex_spec.SessionTableSpec):
+        elif 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)
+                plan = nodes.PromoteOffsetsNode(plan, identifiers.ColumnId(order_col_id))
                 cluster_cols = [order_col_id]
             elif dest_spec.ordering == "order_key":
-                plan, _ = rewrite.pull_out_order(plan)
+                plan, ordering = rewrite.pull_out_order(plan)
             destination_table = self.storage_manager.create_temp_table(
                 plan.schema.to_bigquery(), cluster_cols
             )
@@ -230,7 +220,16 @@ def _execute_bigquery(
                     permit_dml=False,
                 ),
             )
-            return self._execute_bigquery(arr_value, execution_spec)
+            result = self._execute_bigquery(arr_value, execution_spec)
+            result._data = dataclasses.replace(result._data, ordering=ordering)
+            return result
+        # Force table creation if result might be large (and user explicitly allowed large results)
+        elif isinstance(dest_spec, ex_spec.EphemeralTableSpec) or dest_spec is None:
+            if not execution_spec.promise_under_10gb:
+                execution_spec = dataclasses.replace(
+                    execution_spec, destination_spec=ex_spec.SessionTableSpec()
+                )
+                return self._execute_bigquery(array_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
@@ -315,6 +314,7 @@ def _export_gbq_with_dml(
             job_config=bigquery.QueryJobConfig(),
             metrics=self.metrics,
             publisher=self._publisher,
+            query_with_job=True,
         )
 
     def dry_run(
@@ -407,13 +407,14 @@ def _cache_with_cluster_cols(
         ]
         cluster_cols = cluster_cols[:_MAX_CLUSTER_COLUMNS]
         execution_spec = ex_spec.ExecutionSpec(
-            destination_spec=ex_spec.SessionTableSpec(cluster_cols=tuple(cluster_cols))
+            destination_spec=ex_spec.SessionTableSpec(cluster_cols=tuple(cluster_cols), ordering="order_key")
         )
         result = self.execute(
             array_value,
             execution_spec=execution_spec,
         )
         assert isinstance(result, executor.BQTableExecuteResult)
+        assert result._data.ordering is not None
         self.cache.cache_results_table(array_value.node, result._data)
 
     def _cache_with_offsets(self, array_value: bigframes.core.ArrayValue):
@@ -428,6 +429,7 @@ def _cache_with_offsets(self, array_value: bigframes.core.ArrayValue):
             execution_spec=execution_spec,
         )
         assert isinstance(result, executor.BQTableExecuteResult)
+        assert result._data.ordering is not None
         self.cache.cache_results_table(array_value.node, result._data)
 
     def _cache_with_session_awareness(

packages/bigframes/bigframes/session/direct_gbq_execution.py

@@ -90,8 +90,6 @@ def execute(
         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 if dest_spec.cluster_cols else None

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

@@ -128,19 +128,14 @@ def test_engines_unary_variance_aggregates(
 def test_sql_engines_median_op_aggregates(
     scalars_array_value: array_value.ArrayValue,
     bigquery_client: bigquery.Client,
+    bq_engine,
+    sqlglot_engine,
 ):
     node = apply_agg_to_all_valid(
         scalars_array_value,
         agg_ops.MedianOp(),
     ).node
-    publisher = events.Publisher()
-    left_engine = direct_gbq_execution.DirectGbqExecutor(
-        bigquery_client, publisher=publisher
-    )
-    right_engine = direct_gbq_execution.DirectGbqExecutor(
-        bigquery_client, compiler="sqlglot", publisher=publisher
-    )
-    assert_equivalence_execution(node, left_engine, right_engine)
+    assert_equivalence_execution(node, bq_engine, sqlglot_engine)
 
 
 @pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)

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

@@ -16,7 +16,7 @@
 
 import bigframes.operations as bf_ops
 from bigframes.core import array_value, nodes, ordering
-from bigframes.session import polars_executor
+from bigframes.session import polars_executor, execution_spec
 from bigframes.testing.engine_utils import assert_equivalence_execution
 
 pytest.importorskip("polars")
@@ -96,7 +96,7 @@ def test_polars_engines_skips_unrecognized_order_expr(
         ),
     )
     node = nodes.OrderByNode(node, ORDER_EXPRESSIONS)
-    assert engine.execute(node, ordered=True) is None
+    assert engine.execute(node, execution_spec.ExecutionSpec(ordered=True)) is None
 
 
 def apply_reverse(node: nodes.BigFrameNode) -> nodes.BigFrameNode: