fix bugs

Author: TrevorBergeron

packages/bigframes/bigframes/session/bq_caching_executor.py

@@ -87,10 +87,8 @@ def __init__(
         self.cache: execution_cache.ExecutionCache = execution_cache.ExecutionCache()
         self.metrics = metrics
         self.loader = loader
-        self.bqstoragereadclient = bqstoragereadclient
         self._enable_polars_execution = enable_polars_execution
         self._publisher = publisher
-        self._labels = labels
 
         # TODO(tswast): Send events from semi-executors, too.
         self._semi_executors: Sequence[semi_executor.SemiExecutor] = (
@@ -109,7 +107,12 @@ def __init__(
             )
         self._upload_lock = threading.Lock()
         self._gbq_executor = direct_gbq_execution.DirectGbqExecutor(
-            bqclient, compile.compiler, metrics=self.metrics, publisher=self._publisher
+            bqclient,
+            compiler=compile.compiler,
+            bqstoragereadclient=bqstoragereadclient,
+            metrics=self.metrics,
+            publisher=self._publisher,
+            labels=labels,
         )
 
     def to_sql(
@@ -132,7 +135,6 @@ def to_sql(
         )
         return compiled.sql
 
-
     def execute(
         self,
         array_value: bigframes.core.ArrayValue,
@@ -167,7 +169,6 @@ def _try_execute_semi_executors(
                 return maybe_result
         return None
 
-
     def _execute_bigquery(
         self,
         array_value: bigframes.core.ArrayValue,
@@ -177,7 +178,10 @@ def _execute_bigquery(
         # Recursive handlers for different cases, maybe extract to explicit interface.
         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.TempTableSpec(
+                    cluster_cols=dest_spec.cluster_cols, lifetime="ephemeral"
+                ),
             )
             results = self._execute_bigquery(array_value, execution_spec)
             self._export_result_gcs(results, dest_spec)
@@ -191,7 +195,11 @@ def _execute_bigquery(
                 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.TempTableSpec(
+                        cluster_cols=execution_spec.destination_spec.cluster_cols,
+                        lifetime="ephemeral",
+                    ),
                 )
                 results = self._execute_bigquery(array_value, execution_spec)
                 self._export_gbq_with_dml(results, dest_spec)
@@ -213,10 +221,14 @@ def _execute_bigquery(
                 )
                 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")
+                    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)
-                
 
         # 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
@@ -393,7 +405,11 @@ 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.TempTableSpec(
+                cluster_cols=tuple(cluster_cols),
+                lifetime="session",
+                ordering="order_key",
+            )
         )
         result_bq_data = self.execute(
             array_value,
@@ -405,7 +421,9 @@ def _cache_with_cluster_cols(
     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.TempTableSpec(
+                cluster_cols=(), lifetime="session", ordering="offsets_col"
+            )
         )
         result_bq_data = self.execute(
             array_value,

packages/bigframes/bigframes/session/direct_gbq_execution.py

@@ -26,8 +26,10 @@
 import bigframes.session._io.bigquery as bq_io
 from bigframes.core import bq_data, compile, nodes
 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
 
@@ -43,11 +45,13 @@ class DirectGbqExecutor(semi_executor.SemiExecutor):
     def __init__(
         self,
         bqclient: bigquery.Client,
-        compiler: Literal["ibis", "sqlglot"]
-        | Callable[[compile.CompileRequest], executor.CompileResult] = "ibis",
+        bqstoragereadclient: google.cloud.bigquery_storage_v1.BigQueryReadClient,
         *,
+        compiler: Literal["ibis", "sqlglot"]
+        | Callable[[CompileRequest], CompileResult] = "ibis",
         metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
         publisher: Optional[bigframes.core.events.Publisher] = None,
+        labels: Mapping[str, str] = {},
     ):
         self.bqclient = bqclient
         if isinstance(compiler, str):
@@ -58,8 +62,10 @@ def __init__(
             )
         else:
             self._compile_fn = compiler
+        self._bqstoragereadclient = bqstoragereadclient
         self._publisher = publisher
         self._metrics = metrics
+        self._labels = labels
 
     def execute(
         self,
@@ -69,36 +75,43 @@ def execute(
         """Just execute whatever plan as is, without further caching or decomposition."""
 
         og_schema = plan.schema
-        compile_request = compile.CompileRequest(
-                plan,
-                sort_rows=spec.ordered,
-                peek_count=spec.peek,
-            )
+        compile_request = CompileRequest(
+            plan,
+            sort_rows=spec.ordered,
+            peek_count=spec.peek,
+        )
 
         compiled = self._compile_fn(compile_request)
         # might have more columns than og schema, for hidden ordering columns
         compiled_schema = compiled.sql_schema
 
         job_config = bigquery.QueryJobConfig()
-        if isinstance(spec.destination_spec, execution_spec.TableOutputSpec):
-            job_config.destination = spec.destination_spec.table
-            job_config.write_disposition = _WRITE_DISPOSITIONS[spec.destination_spec.if_exists]
-            job_config.clustering_fields = spec.destination_spec.cluster_cols
-        elif isinstance(spec.destination_spec, execution_spec.TempTableSpec) and spec.destination_spec.lifetime == "ephemeral":
-            pass
-        elif spec.destination_spec is not None:
-            raise ValueError(f"Direct GBQ Executor does not support destination: {spec.destination_spec}")
+        dest_spec = spec.destination_spec
+        cluster_cols = ()
+        if isinstance(dest_spec, execution_spec.TableOutputSpec):
+            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
+        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 = spec.destination_spec is None and spec.promise_under_10gb
+        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,
             query_with_job=(not can_skip_job),
             session=plan.session,
         )
-
-        cluster_cols = spec.destination_spec.cluster_cols if spec.desination_spec else ()
         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
@@ -128,7 +141,7 @@ def execute(
             return executor.BQTableExecuteResult(
                 data=result_bq_data,
                 project_id=self.bqclient.project,
-                storage_client=self.bqstoragereadclient,
+                storage_client=self._bqstoragereadclient,
                 execution_metadata=execution_metadata,
                 selected_fields=tuple((col, col) for col in og_schema.names),
             )
@@ -159,34 +172,15 @@ def _run_execute_query(
             job_config.labels.update(self._labels)
 
         try:
-            # Trick the type checker into thinking we got a literal.
-            if query_with_job:
-                return bq_io.start_query_with_client(
-                    self.bqclient,
-                    sql,
-                    job_config=job_config,
-                    metrics=self._metrics,
-                    project=None,
-                    location=None,
-                    timeout=None,
-                    query_with_job=True,
-                    publisher=self._publisher,
-                    session=session,
-                )
-            else:
-                return bq_io.start_query_with_client(
-                    self.bqclient,
-                    sql,
-                    job_config=job_config,
-                    metrics=self._metrics,
-                    project=None,
-                    location=None,
-                    timeout=None,
-                    query_with_job=False,
-                    publisher=self._publisher,
-                    session=session,
-                )
-
+            return bq_io.start_query_with_client(
+                self.bqclient,
+                sql,
+                job_config=job_config,
+                metrics=self._metrics,
+                query_with_job=query_with_job,
+                publisher=self._publisher,
+                session=session,
+            )
         except google.api_core.exceptions.BadRequest as e:
             # Unfortunately, this error type does not have a separate error code or exception type
             if "Resources exceeded during query execution" in e.message:
@@ -195,6 +189,7 @@ def _run_execute_query(
             else:
                 raise
 
+
 def _result_schema(
     logical_schema: schemata.ArraySchema, sql_schema: list[bigquery.SchemaField]
 ) -> schemata.ArraySchema:

packages/bigframes/bigframes/session/execution_spec.py

@@ -28,7 +28,7 @@ class ExecutionSpec:
     # implementation dependent and not stable.
     peek: Optional[int] = None
     # Controls whether output iterator is ordered. Cannot be true if destination is not
-    # guaranteed to be ordered. 
+    # guaranteed to be ordered.
     ordered: bool = False
     # This is an optimization flag for gbq execution, it doesn't change semantics, but if promise is falsely made, errors may occur
     promise_under_10gb: bool = False
@@ -41,7 +41,10 @@ class TempTableSpec:
     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
+
+    cluster_cols: tuple[
+        str, ...
+    ]  # if empty, will cluster using order key if ordering_key is set
     lifetime: Literal["session", "ephemeral"] = "session"
     # Controls ordering and whether extra columns are materialized to preserve ordering
     # Any extra columns will be appended to the end of the schema.
@@ -59,6 +62,7 @@ class TableOutputSpec:
 
     The executor is not responsible for managing lifecycle of the table.
     """
+
     table: bigquery.TableReference
     cluster_cols: tuple[str, ...]
     if_exists: Literal["fail", "replace", "append"] = "fail"