more threaded config fixes

Author: TrevorBergeron

packages/bigframes/bigframes/core/blocks.py

@@ -277,9 +277,7 @@ def shape(self) -> typing.Tuple[int, int]:
         row_count = (
             self.session._executor.execute(
                 self.expr.row_count(),
-                execution_spec.ExecutionSpec(
-                    promise_under_10gb=True, ordered=False
-                ).with_current_configuration(),
+                execution_spec.ExecutionSpec(promise_under_10gb=True, ordered=False),
             )
             .batches()
             .to_py_scalar()
@@ -592,7 +590,7 @@ def to_arrow(
             execution_spec.ExecutionSpec(
                 promise_under_10gb=under_10gb,
                 ordered=ordered,
-            ).with_current_configuration(),
+            ),
         )
         pa_table = execute_result.batches().to_arrow_table()
 
@@ -686,9 +684,7 @@ def try_peek(
             )
             result = self.session._executor.execute(
                 self.expr,
-                execution_spec.ExecutionSpec(
-                    promise_under_10gb=under_10gb, peek=n
-                ).with_current_configuration(),
+                execution_spec.ExecutionSpec(promise_under_10gb=under_10gb, peek=n),
             )
             df = result.batches().to_pandas()
             return self._copy_index_to_pandas(df)
@@ -717,7 +713,7 @@ def to_pandas_batches(
             execution_spec.ExecutionSpec(
                 promise_under_10gb=under_10gb,
                 ordered=True,
-            ).with_current_configuration(),
+            ),
         )
         result_batches = execution_result.batches()
 
@@ -781,7 +777,7 @@ def _materialize_local(
             execution_spec.ExecutionSpec(
                 promise_under_10gb=under_10gb,
                 ordered=materialize_options.ordered,
-            ).with_current_configuration(),
+            ),
         )
         result_batches = execute_result.batches()
 
@@ -1625,15 +1621,15 @@ def retrieve_repr_request_results(
             execution_spec.ExecutionSpec(
                 promise_under_10gb=True,
                 ordered=True,
-            ).with_current_configuration(),
+            ),
         )
         row_count = (
             self.session._executor.execute(
                 self.expr.row_count(),
                 execution_spec.ExecutionSpec(
                     promise_under_10gb=True,
                     ordered=False,
-                ).with_current_configuration(),
+                ),
             )
             .batches()
             .to_py_scalar()

packages/bigframes/bigframes/core/indexes/base.py

@@ -294,9 +294,7 @@ def get_loc(self, key) -> typing.Union[int, slice, "bigframes.series.Series"]:
         count_scalar = (
             self._block.session._executor.execute(
                 count_result,
-                ex_spec.ExecutionSpec(
-                    promise_under_10gb=True
-                ).with_current_configuration(),
+                ex_spec.ExecutionSpec(promise_under_10gb=True),
             )
             .batches()
             .to_py_scalar()
@@ -312,9 +310,7 @@ def get_loc(self, key) -> typing.Union[int, slice, "bigframes.series.Series"]:
             position_scalar = (
                 self._block.session._executor.execute(
                     position_result,
-                    ex_spec.ExecutionSpec(
-                        promise_under_10gb=True
-                    ).with_current_configuration(),
+                    ex_spec.ExecutionSpec(promise_under_10gb=True),
                 )
                 .batches()
                 .to_py_scalar()
@@ -353,9 +349,7 @@ def _get_monotonic_slice(
         result_df = (
             self._block.session._executor.execute(
                 combined_result,
-                execution_spec=ex_spec.ExecutionSpec(
-                    promise_under_10gb=True
-                ).with_current_configuration(),
+                execution_spec=ex_spec.ExecutionSpec(promise_under_10gb=True),
             )
             .batches()
             .to_pandas()

packages/bigframes/bigframes/dataframe.py

@@ -4202,7 +4202,7 @@ def to_csv(
                 ex_spec.GcsOutputSpec(
                     uri=path_or_buf, format="csv", export_options=tuple(options.items())
                 )
-            ).with_current_configuration(),
+            ),
         )
         self._set_internal_query_job(result.query_job)
         return None
@@ -4251,7 +4251,7 @@ def to_json(
             export_array.rename_columns(id_overrides),
             ex_spec.ExecutionSpec(
                 ex_spec.GcsOutputSpec(uri=path_or_buf, format="json", export_options=())
-            ).with_current_configuration(),
+            ),
         )
         self._set_internal_query_job(result.query_job)
         return None
@@ -4334,7 +4334,7 @@ def to_gbq(
                     cluster_cols=tuple(clustering_fields),
                     if_exists=if_exists,
                 )
-            ).with_current_configuration(),
+            ),
         )
         assert result.query_job is not None
         self._set_internal_query_job(result.query_job)
@@ -4414,7 +4414,7 @@ def to_parquet(
                     format="parquet",
                     export_options=tuple(export_options.items()),
                 )
-            ).with_current_configuration(),
+            ),
         )
         self._set_internal_query_job(result.query_job)
         return None

packages/bigframes/bigframes/formatting_helpers.py

@@ -26,6 +26,8 @@
 import google.cloud.bigquery as bigquery
 import humanize
 
+import bigframes._config
+
 if TYPE_CHECKING:
     import bigframes.core.events
 

packages/bigframes/bigframes/session/bq_caching_executor.py

@@ -41,6 +41,7 @@
 import bigframes.session.metrics
 import bigframes.session.planner
 import bigframes.session.temporary_storage
+from bigframes._config import ComputeOptions
 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
@@ -183,6 +184,8 @@ def execute(
         array_value: bigframes.core.ArrayValue,
         execution_spec: ex_spec.ExecutionSpec,
     ) -> executor.ExecuteResult:
+        # Need to grab thread local before starting async execution.
+        execution_spec = execution_spec.with_compute_options(bigframes.options.compute)
         return _run_sync(
             self._execute_async(
                 array_value,
@@ -319,7 +322,9 @@ async def _execute_gbq_query_only(
         array_value: bigframes.core.ArrayValue,
         execution_spec: ex_spec.ExecutionSpec,
     ) -> executor.ExecuteResult:
-        gbq_plan = await self._prepare_plan_bq_execution(array_value.node)
+        gbq_plan = await self._prepare_plan_bq_execution(
+            array_value.node, execution_spec.bigquery_config
+        )
         result = await self._gbq_executor.execute(gbq_plan, execution_spec)
         if result is None:
             raise ValueError(
@@ -395,10 +400,13 @@ def dry_run(
     def cached(
         self, array_value: bigframes.core.ArrayValue, *, config: executor.CacheConfig
     ) -> None:
+        # Get compute options before passing to async method, can be thread-local
+        bq_compute_options = ex_spec.BqComputeOptions.from_compute_options(
+            bigframes.options.compute
+        )
         return _run_sync(
             self._cached_async(
-                array_value,
-                config=config,
+                array_value, config=config, compute_options=bq_compute_options
             )
         )
 
@@ -407,6 +415,7 @@ async def _cached_async(
         array_value: bigframes.core.ArrayValue,
         *,
         config: executor.CacheConfig,
+        compute_options: ex_spec.BqComputeOptions,
     ) -> None:
         """Write the block to a session table."""
         # First, see if we can reuse the existing cache
@@ -429,25 +438,23 @@ async def _cached_async(
 
         if config.optimize_for == "auto":
             await self._cache_with_session_awareness(
-                array_value, enable_multi_query_execution=enable_multi_query_execution
+                array_value, compute_options=compute_options
             )
         elif config.optimize_for == "head":
-            await self._cache_with_offsets(
-                array_value, enable_multi_query_execution=enable_multi_query_execution
-            )
+            await self._cache_with_offsets(array_value, compute_options=compute_options)
         else:
             assert isinstance(config.optimize_for, executor.HierarchicalKey)
             await self._cache_with_cluster_cols(
                 array_value,
                 cluster_cols=config.optimize_for.columns,
-                enable_multi_query_execution=enable_multi_query_execution,
+                compute_options=compute_options,
             )
 
     async def _execute_to_cached_table(
         self,
         plan: nodes.BigFrameNode,
         cache_spec: ex_spec.CacheSpec,
-        enable_multi_query_execution: bool,
+        compute_options: ex_spec.BqComputeOptions,
     ) -> executor.ExecuteResult:
         # "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
@@ -474,7 +481,7 @@ async def _execute_to_cached_table(
                 cluster_cols=cluster_cols,
                 if_exists="replace",
             ),
-            enable_multi_query_execution=enable_multi_query_execution,
+            bigquery_config=compute_options,
         )
         # We don't use _execute_gbq_table_export, as this result is internal, not exported.
         result = await self._execute_gbq_query_only(
@@ -506,13 +513,13 @@ def _prepare_plan_simplify(self, plan: nodes.BigFrameNode) -> nodes.BigFrameNode
         return plan
 
     async def _prepare_plan_bq_execution(
-        self, plan: nodes.BigFrameNode, enable_multi_query_execution: bool
+        self,
+        plan: nodes.BigFrameNode,
+        compute_options: ex_spec.BqComputeOptions,
     ) -> nodes.BigFrameNode:
         """Prepare the plan for BigQuery execution by caching subtrees and uploading large local sources."""
-        if enable_multi_query_execution:
-            await self._simplify_with_caching(
-                plan, enable_multi_query_execution=enable_multi_query_execution
-            )
+        if compute_options.enable_multi_query_execution:
+            await self._simplify_with_caching(plan, compute_options=compute_options)
         plan = self._prepare_plan_simplify(plan)
         plan = await self._substitute_large_local_sources(plan)
         return plan
@@ -521,7 +528,7 @@ async def _cache_with_cluster_cols(
         self,
         array_value: bigframes.core.ArrayValue,
         cluster_cols: Sequence[str],
-        enable_multi_query_execution: bool,
+        compute_options: ex_spec.BqComputeOptions,
     ):
         """Executes the query and uses the resulting table to rewrite future executions."""
         cluster_cols = [
@@ -533,7 +540,7 @@ async def _cache_with_cluster_cols(
         result = await self._execute_to_cached_table(
             array_value.node,
             ex_spec.CacheSpec(cluster_cols=tuple(cluster_cols), ordering="order_key"),
-            enable_multi_query_execution=enable_multi_query_execution,
+            compute_options=compute_options,
         )
         assert isinstance(result, executor.BQTableExecuteResult)
         assert result._data.ordering is not None
@@ -542,13 +549,13 @@ async def _cache_with_cluster_cols(
     async def _cache_with_offsets(
         self,
         array_value: bigframes.core.ArrayValue,
-        enable_multi_query_execution: bool,
+        compute_options: ex_spec.BqComputeOptions,
     ):
         """Executes the query and uses the resulting table to rewrite future executions."""
         result = await self._execute_to_cached_table(
             array_value.node,
             ex_spec.CacheSpec(ordering="offsets_col"),
-            enable_multi_query_execution=enable_multi_query_execution,
+            compute_options=compute_options,
         )
         assert isinstance(result, executor.BQTableExecuteResult)
         assert result._data.ordering is not None
@@ -557,7 +564,7 @@ async def _cache_with_offsets(
     async def _cache_with_session_awareness(
         self,
         array_value: bigframes.core.ArrayValue,
-        enable_multi_query_execution: bool,
+        compute_options: ex_spec.BqComputeOptions,
     ) -> None:
         session_forest = [obj._block._expr.node for obj in array_value.session.objects]
         # These node types are cheap to re-compute
@@ -569,22 +576,22 @@ async def _cache_with_session_awareness(
             await self._cache_with_cluster_cols(
                 bigframes.core.ArrayValue(target),
                 cluster_cols_sql_names,
-                enable_multi_query_execution=enable_multi_query_execution,
+                compute_options=compute_options,
             )
         elif not target.order_ambiguous:
             await self._cache_with_offsets(
                 bigframes.core.ArrayValue(target),
-                enable_multi_query_execution=enable_multi_query_execution,
+                compute_options=compute_options,
             )
         else:
             await self._cache_with_cluster_cols(
                 bigframes.core.ArrayValue(target),
                 [],
-                enable_multi_query_execution=enable_multi_query_execution,
+                compute_options=compute_options,
             )
 
     async def _simplify_with_caching(
-        self, plan: nodes.BigFrameNode, enable_multi_query_execution: bool
+        self, plan: nodes.BigFrameNode, compute_options: ex_spec.BqComputeOptions
     ):
         """Attempts to handle the complexity by caching duplicated subtrees and breaking the query into pieces."""
         # Apply existing caching first
@@ -596,13 +603,13 @@ async def _simplify_with_caching(
                 return
 
             did_cache = await self._cache_most_complex_subtree(
-                plan, enable_multi_query_execution=enable_multi_query_execution
+                plan, compute_options=compute_options
             )
             if not did_cache:
                 return
 
     async def _cache_most_complex_subtree(
-        self, node: nodes.BigFrameNode, enable_multi_query_execution: bool
+        self, node: nodes.BigFrameNode, compute_options: ex_spec.BqComputeOptions
     ) -> bool:
         # TODO: If query fails, retry with lower complexity limit
         selection = tree_properties.select_cache_target(
@@ -622,7 +629,7 @@ async def _cache_most_complex_subtree(
         await self._cache_with_cluster_cols(
             bigframes.core.ArrayValue(selection),
             [],
-            enable_multi_query_execution=enable_multi_query_execution,
+            compute_options=compute_options,
         )
         return True
 

packages/bigframes/bigframes/session/direct_gbq_execution.py

@@ -95,10 +95,13 @@ async def execute(
         job_config.labels["bigframes-dtypes"] = compiled.encoded_type_refs
         if self._labels:
             job_config.labels.update(self._labels)
-        if spec.labels:
-            job_config.labels.update(spec.labels)
-        if spec.maximum_bytes_billed is not None:
-            job_config.maximum_bytes_billed = spec.maximum_bytes_billed
+        if spec.bigquery_config is not None:
+            if spec.bigquery_config.extra_query_labels:
+                job_config.labels.update(spec.bigquery_config.extra_query_labels)
+            if spec.bigquery_config.maximum_bytes_billed is not None:
+                job_config.maximum_bytes_billed = (
+                    spec.bigquery_config.maximum_bytes_billed
+                )
 
         iterator, query_job = await asyncio.to_thread(
             self._run_execute_query,