fixes

Author: TrevorBergeron

packages/bigframes/bigframes/core/bq_data.py

@@ -253,6 +253,9 @@ def __post_init__(self):
     # Optimization field, must be correct if set, don't put maybe-stale number here
     n_rows: Optional[int] = None
 
+    def with_ordering(self, ordering: orderings.RowOrdering) -> BigqueryDataSource:
+        return dataclasses.replace(self, ordering=ordering)
+
 
 _WORKER_TIME_INCREMENT = 0.05
 

packages/bigframes/bigframes/core/nodes.py

@@ -846,10 +846,10 @@ def remap_refs(
     ) -> ReadTableNode:
         return self
 
-    def with_order_cols(self):
+    def pull_out_order(self):
         # Maybe the ordering should be required to always be in the scan list, and then we won't need this?
         if self.source.ordering is None:
-            return self, orderings.RowOrdering()
+            return self, RowOrdering()
 
         order_cols = {col.sql for col in self.source.ordering.referenced_columns}
         scan_cols = {col.source_id for col in self.scan_list.items}
@@ -863,10 +863,18 @@ def with_order_cols(self):
         ]
         new_scan_list = ScanList(items=(*self.scan_list.items, *new_scan_cols))
         new_order = self.source.ordering.remap_column_refs(
-            {identifiers.ColumnId(item.source_id): item.id for item in new_scan_cols},
+            {
+                identifiers.ColumnId(item.source_id): item.id
+                for item in new_scan_list.items
+            },
             allow_partial_bindings=True,
         )
-        return dataclasses.replace(self, scan_list=new_scan_list), new_order
+        new_node = dataclasses.replace(
+            self,
+            scan_list=new_scan_list,
+            source=self.source.with_ordering(RowOrdering()),
+        )
+        return new_node, new_order
 
 
 @dataclasses.dataclass(frozen=True, eq=False)

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

@@ -162,7 +162,7 @@ def pull_up_order_inner(
                 )
         elif isinstance(node, bigframes.core.nodes.ReadTableNode):
             if node.source.ordering is not None:
-                return node.with_order_cols()
+                return node.pull_out_order()
             else:
                 # No defined ordering
                 return node, bigframes.core.ordering.RowOrdering()

packages/bigframes/bigframes/session/bq_caching_executor.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import concurrent.futures
+import dataclasses
 import math
 import threading
 from typing import Literal, Mapping, Optional, Sequence, Tuple
@@ -28,8 +29,8 @@
 import bigframes.core
 import bigframes.core.events
 import bigframes.core.guid
-import bigframes.core.ordering
 import bigframes.core.nodes as nodes
+import bigframes.core.ordering
 import bigframes.core.schema as schemata
 import bigframes.core.tree_properties as tree_properties
 import bigframes.dtypes
@@ -43,14 +44,13 @@
 from bigframes.core.compile.sqlglot import sql as sg_sql
 from bigframes.core.compile.sqlglot import sqlglot_ir
 from bigframes.session import (
+    direct_gbq_execution,
     executor,
     loader,
     local_scan_executor,
     read_api_execution,
     semi_executor,
-    direct_gbq_execution,
 )
-import dataclasses
 
 # Max complexity that should be executed as a single query
 QUERY_COMPLEXITY_LIMIT = 1e7
@@ -189,9 +189,9 @@ def _execute_bigquery(
             self._export_result_gcs(results, dest_spec)
             return results
         elif isinstance(dest_spec, ex_spec.TableOutputSpec):
-            return self._execute_gbq_table_output(array_value, execution_spec)
+            return self._execute_gbq_table_export(array_value, execution_spec)
         # 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:
+        elif isinstance(dest_spec, ex_spec.EphemeralTableSpec) or (dest_spec is None):
             if not execution_spec.promise_under_10gb:
                 table = self.storage_manager.create_temp_table(
                     array_value.schema.to_bigquery()
@@ -202,12 +202,12 @@ def _execute_bigquery(
                         table=table, if_exists="append"
                     ),
                 )
-                # We don't use _execute_gbq_table_output, because we want to skip slower DML path.
+                # We don't use _execute_gbq_table_export, as this result is internal, not exported.
                 return self._execute_gbq_query_only(array_value, execution_spec)
         # At this point, dst should be unspecified, a specific bq table, or an ephemeral temp table that fits in <10gb
         return self._execute_gbq_query_only(array_value, execution_spec)
 
-    def _execute_gbq_table_output(
+    def _execute_gbq_table_export(
         self,
         array_value: bigframes.core.ArrayValue,
         execution_spec: ex_spec.ExecutionSpec,
@@ -224,9 +224,20 @@ def _execute_gbq_table_output(
             )
             results = self._execute_bigquery(array_value, execution_spec)
             self._export_gbq_with_dml(results, dest_spec)
-            return results
-        # If not compatible with DML path, just run query with destination unchanged
-        return self._execute_gbq_query_only(array_value, execution_spec)
+            result = results
+        else:
+            result = self._execute_gbq_query_only(array_value, execution_spec)
+
+        has_special_dtype_col = any(
+            t in (bigframes.dtypes.TIMEDELTA_DTYPE, bigframes.dtypes.OBJ_REF_DTYPE)
+            for t in array_value.schema.dtypes
+        )
+        if dest_spec.if_exists != "append" and has_special_dtype_col:
+            table = self.bqclient.get_table(dest_spec.table)
+            table.schema = array_value.schema.to_bigquery()
+            self.bqclient.update_table(table, ["schema"])
+
+        return result
 
     def _execute_gbq_query_only(
         self,
@@ -347,7 +358,9 @@ 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)
+            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(
@@ -361,7 +374,7 @@ def _execute_to_cached_table(
                 if_exists="replace",
             )
         )
-        # We don't use _execute_gbq_table_output, because we want to skip slower DML path.
+        # We don't use _execute_gbq_table_export, as this result is internal, not exported.
         result = self._execute_gbq_query_only(arr_value, execution_spec)
         result._data = dataclasses.replace(result._data, ordering=ordering)
         return result

packages/bigframes/bigframes/session/direct_gbq_execution.py

@@ -15,25 +15,23 @@
 
 from typing import Callable, Literal, Mapping, Optional, Tuple
 
+import google.api_core.exceptions
 import google.cloud.bigquery.job as bq_job
 import google.cloud.bigquery.table as bq_table
+import google.cloud.bigquery_storage_v1
 from google.cloud import bigquery
 
+import bigframes.core.compile
 import bigframes.core.compile.ibis_compiler.ibis_compiler as ibis_compiler
 import bigframes.core.compile.sqlglot.compiler as sqlglot_compiler
 import bigframes.core.events
-import bigframes.session.metrics
+import bigframes.core.schema as schemata
 import bigframes.session._io.bigquery as bq_io
+import bigframes.session.metrics
+from bigframes import exceptions as bfe
 from bigframes.core import bq_data, compile, nodes
-import bigframes.core.compile
-from bigframes.session import executor, semi_executor, execution_spec
 from bigframes.core.compile.configs import CompileRequest, CompileResult
-from bigframes import exceptions as bfe
-import bigframes.core.schema as schemata
-import google.cloud.bigquery_storage_v1
-
-import google.api_core.exceptions
-
+from bigframes.session import execution_spec, executor, semi_executor
 
 _WRITE_DISPOSITIONS = {
     "fail": bigquery.WriteDisposition.WRITE_EMPTY,
@@ -66,20 +64,14 @@ def execute(
         spec: execution_spec.ExecutionSpec,
     ) -> executor.ExecuteResult:
         """Just execute whatever plan as is, without further caching or decomposition."""
-
-        og_schema = plan.schema
-        compile_request = CompileRequest(
-            plan,
-            sort_rows=spec.ordered,
-            peek_count=spec.peek,
-        )
-
         compiled = compile.compile_sql(
-            compile_request, compiler_name=self._compiler_name
+            CompileRequest(
+                plan,
+                sort_rows=spec.ordered,
+                peek_count=spec.peek,
+            ),
+            compiler_name=self._compiler_name,
         )
-        # might have more columns than og schema, for hidden ordering columns
-        compiled_schema = compiled.sql_schema
-
         job_config = bigquery.QueryJobConfig()
         dest_spec = spec.destination_spec
         cluster_cols = None
@@ -110,18 +102,16 @@ def execute(
         )
         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),
+                    tuple(compiled.sql_schema),
                     cluster_cols=cluster_cols or (),
                     location=iterator.location or self.bqclient.location,
                     table_type="TABLE",
                 ),
-                schema=result_bf_schema,
+                schema=plan.schema,
                 ordering=compiled.row_order,
                 n_rows=iterator.total_rows,
             )
@@ -143,26 +133,25 @@ def execute(
                 project_id=self.bqclient.project,
                 storage_client=self._bqstoragereadclient,
                 execution_metadata=execution_metadata,
-                selected_fields=tuple((col, col) for col in og_schema.names),
+                selected_fields=tuple((col, col) for col in plan.schema.names),
             )
         else:
             return executor.LocalExecuteResult(
-                data=iterator.to_arrow().select(og_schema.names),
+                data=iterator.to_arrow().select(plan.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,
+        job_config: bq_job.QueryJobConfig,
+        query_with_job: bool,
+        session,
     ) -> Tuple[bq_table.RowIterator, Optional[bigquery.QueryJob]]:
         """
         Starts BigQuery query job and waits for results.
         """
-        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
@@ -188,13 +177,3 @@ def _run_execute_query(
                 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())
-    )

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

@@ -114,10 +114,11 @@ def test_read_gbq_tokyo(
     df.sort_index(inplace=True)
     expected = scalars_pandas_df_index
 
-    # don't promise under 10gb, so table creation, and job creation are guaranteed
     exec_result = session_tokyo._executor.execute(
         df._block.expr,
-        bigframes.session.execution_spec.ExecutionSpec(promise_under_10gb=False),
+        bigframes.session.execution_spec.ExecutionSpec(
+            destination_spec=bigframes.session.execution_spec.EphemeralTableSpec()
+        ),
     )
     assert exec_result.query_job is not None
     assert exec_result.query_job.location == tokyo_location
@@ -948,7 +949,9 @@ def test_read_pandas_tokyo(
 
     result = session_tokyo._executor.execute(
         df._block.expr,
-        bigframes.session.execution_spec.ExecutionSpec(promise_under_10gb=False),
+        bigframes.session.execution_spec.ExecutionSpec(
+            destination_spec=bigframes.session.execution_spec.EphemeralTableSpec()
+        ),
     )
     assert result.query_job is not None
     assert result.query_job.location == tokyo_location