refactor query metadata

Author: TrevorBergeron

bigframes/session/bq_caching_executor.py

@@ -320,7 +320,7 @@ def _export_gbq(
 
         # TODO(swast): plumb through the api_name of the user-facing api that
         # caused this query.
-        _, query_job = self._run_execute_query(
+        iterator, _ = self._run_execute_query(
             sql=sql,
             job_config=job_config,
         )
@@ -338,7 +338,7 @@ def _export_gbq(
 
         return executor.EmptyExecuteResult(
             bf_schema=array_value.schema,
-            query_job=query_job,
+            execution_metadata=executor.ExecutionMetadata.from_iterator(iterator),
         )
 
     def dry_run(
@@ -681,18 +681,20 @@ def _execute_plan_gbq(
             assert compiled.row_order is not None
             self.cache.cache_results_table(og_plan, result_bq_data)
 
+        execution_metadata = executor.ExecutionMetadata.from_iterator(iterator)
         if result_bq_data is not None:
             return executor.BQTableExecuteResult(
                 data=result_bq_data,
                 project_id=self.bqclient.project,
                 storage_client=self.bqstoragereadclient,
-                query_job=query_job,
+                execution_metadata=execution_metadata,
                 selected_fields=tuple((col, col) for col in og_schema.names),
             )
         else:
             return executor.LocalExecuteResult(
                 data=iterator.to_arrow().select(og_schema.names),
                 bf_schema=plan.schema,
+                execution_metadata=execution_metadata,
             )
 
 

bigframes/session/direct_gbq_execution.py

@@ -68,6 +68,7 @@ def execute(
         return executor.LocalExecuteResult(
             data=iterator.to_arrow(),
             bf_schema=plan.schema,
+            execution_metadata=executor.ExecutionMetadata.from_iterator(iterator),
         )
 
     def _run_execute_query(

bigframes/session/executor.py

@@ -21,6 +21,7 @@
 from typing import Iterator, Literal, Optional, Sequence, Union
 
 from google.cloud import bigquery, bigquery_storage_v1
+import google.cloud.bigquery.table as bq_table
 import pandas as pd
 import pyarrow
 import pyarrow as pa
@@ -40,7 +41,9 @@
 
 
 class ResultsIterator(Iterator[pa.RecordBatch]):
-    """Interface for mutable objects with state represented by a block value object."""
+    """
+    Iterator for query results, with some extra metadata attached.
+    """
 
     def __init__(
         self,
@@ -141,12 +144,7 @@ def to_py_scalar(self):
 class ExecuteResult(abc.ABC):
     @property
     @abc.abstractmethod
-    def query_job(self) -> Optional[bigquery.QueryJob]:
-        ...
-
-    @property
-    @abc.abstractmethod
-    def total_bytes_processed(self) -> Optional[int]:
+    def execution_metadata(self) -> ExecutionMetadata:
         ...
 
     @property
@@ -158,18 +156,40 @@ def schema(self) -> bigframes.core.schema.ArraySchema:
     def batches(self) -> ResultsIterator:
         ...
 
-
-class LocalExecuteResult(ExecuteResult):
-    def __init__(self, data: pa.Table, bf_schema: bigframes.core.schema.ArraySchema):
-        self._data = local_data.ManagedArrowTable.from_pyarrow(data, bf_schema)
-
     @property
     def query_job(self) -> Optional[bigquery.QueryJob]:
-        return None
+        return self.execution_metadata.query_job
 
     @property
     def total_bytes_processed(self) -> Optional[int]:
-        return None
+        return self.execution_metadata.bytes_processed
+
+
+@dataclasses.dataclass(frozen=True)
+class ExecutionMetadata:
+    query_job: Optional[bigquery.QueryJob] = None
+    bytes_processed: Optional[int] = None
+
+    @classmethod
+    def from_iterator(cls, iterator: bq_table.RowIterator) -> ExecutionMetadata:
+        return cls(
+            query_job=iterator.query, bytes_processed=iterator.total_bytes_processed
+        )
+
+
+class LocalExecuteResult(ExecuteResult):
+    def __init__(
+        self,
+        data: pa.Table,
+        bf_schema: bigframes.core.schema.ArraySchema,
+        execution_metadata: ExecutionMetadata = ExecutionMetadata(),
+    ):
+        self._data = local_data.ManagedArrowTable.from_pyarrow(data, bf_schema)
+        self._execution_metadata = execution_metadata
+
+    @property
+    def execution_metadata(self) -> ExecutionMetadata:
+        return self._execution_metadata
 
     @property
     def schema(self) -> bigframes.core.schema.ArraySchema:
@@ -188,20 +208,14 @@ class EmptyExecuteResult(ExecuteResult):
     def __init__(
         self,
         bf_schema: bigframes.core.schema.ArraySchema,
-        query_job: Optional[bigquery.QueryJob] = None,
+        execution_metadata: ExecutionMetadata = ExecutionMetadata(),
     ):
         self._schema = bf_schema
-        self._query_job = query_job
+        self._execution_metadata = execution_metadata
 
     @property
-    def query_job(self) -> Optional[bigquery.QueryJob]:
-        return self._query_job
-
-    @property
-    def total_bytes_processed(self) -> Optional[int]:
-        if self.query_job:
-            return self.query_job.total_bytes_processed
-        return None
+    def execution_metadata(self) -> ExecutionMetadata:
+        return self._execution_metadata
 
     @property
     def schema(self) -> bigframes.core.schema.ArraySchema:
@@ -218,28 +232,22 @@ def __init__(
         storage_client: bigquery_storage_v1.BigQueryReadClient,
         project_id: str,
         *,
-        query_job: Optional[bigquery.QueryJob] = None,
+        execution_metadata: ExecutionMetadata = ExecutionMetadata(),
         limit: Optional[int] = None,
         selected_fields: Optional[Sequence[tuple[str, str]]] = None,
     ):
         self._data = data
         self._project_id = project_id
-        self._query_job = query_job
+        self._execution_metadata = execution_metadata
         self._storage_client = storage_client
         self._limit = limit
         self._selected_fields = selected_fields or [
             (name, name) for name in data.schema.names
         ]
 
     @property
-    def query_job(self) -> Optional[bigquery.QueryJob]:
-        return self._query_job
-
-    @property
-    def total_bytes_processed(self) -> Optional[int]:
-        if self.query_job:
-            return self.query_job.total_bytes_processed
-        return None
+    def execution_metadata(self) -> ExecutionMetadata:
+        return self._execution_metadata
 
     @property
     @functools.cache