Implement execution history

Author: shuoweil

packages/bigframes/bigframes/__init__.py

@@ -45,6 +45,7 @@
 from bigframes._config.bigquery_options import BigQueryOptions  # noqa: E402
 from bigframes.core.global_session import (  # noqa: E402
     close_session,
+    execution_history,
     get_global_session,
 )
 from bigframes.session import Session, connect  # noqa: E402
@@ -69,6 +70,7 @@ def load_ipython_extension(ipython):
     "BigQueryOptions",
     "get_global_session",
     "close_session",
+    "execution_history",
     "enums",
     "exceptions",
     "connect",

packages/bigframes/bigframes/core/global_session.py

@@ -26,6 +26,7 @@
 import bigframes.exceptions as bfe
 
 if TYPE_CHECKING:
+    import pandas
     import bigframes.session
 
 _global_session: Optional[bigframes.session.Session] = None
@@ -124,6 +125,14 @@ def with_default_session(func_: Callable[..., _T], *args, **kwargs) -> _T:
     return func_(get_global_session(), *args, **kwargs)
 
 
+def execution_history() -> "pandas.DataFrame":
+    import pandas  # noqa: F401
+
+    import bigframes.session
+
+    return with_default_session(bigframes.session.Session.execution_history)
+
+
 class _GlobalSessionContext:
     """
     Context manager for testing that sets global session.

packages/bigframes/bigframes/session/__init__.py

@@ -109,6 +109,39 @@
 logger = logging.getLogger(__name__)
 
 
+class _ExecutionHistory(pandas.DataFrame):
+    @property
+    def _constructor(self):
+        return _ExecutionHistory
+
+    def _repr_html_(self) -> str | None:
+        try:
+            import bigframes.formatting_helpers as formatter
+
+            if self.empty:
+                return "<div>No executions found.</div>"
+
+            cols = ["job_id", "status", "total_bytes_processed", "job_url"]
+            df_display = self[cols].copy()
+            df_display["total_bytes_processed"] = df_display[
+                "total_bytes_processed"
+            ].apply(formatter.get_formatted_bytes)
+
+            def format_url(url):
+                return f'<a target="_blank" href="{url}">Open Job</a>' if url else ""
+
+            df_display["job_url"] = df_display["job_url"].apply(format_url)
+
+            # Rename job_id to query_id to match user expectations
+            df_display = df_display.rename(columns={"job_id": "query_id"})
+
+            compact_html = df_display.to_html(escape=False, index=False)
+
+            return compact_html
+        except Exception:
+            return super()._repr_html_()  # type: ignore
+
+
 @log_adapter.class_logger
 class Session(
     third_party_pandas_gbq.GBQIOMixin,
@@ -233,6 +266,7 @@ def __init__(
         )
 
         self._metrics = metrics.ExecutionMetrics()
+        self._publisher.subscribe(self._metrics.on_event)
         self._function_session = bff_session.FunctionSession()
         self._anon_dataset_manager = anonymous_dataset.AnonymousDatasetManager(
             self._clients_provider.bqclient,
@@ -371,6 +405,10 @@ def slot_millis_sum(self):
         """The sum of all slot time used by bigquery jobs in this session."""
         return self._metrics.slot_millis
 
+    def execution_history(self) -> pandas.DataFrame:
+        """Returns a list of underlying BigQuery executions initiated by BigFrames in the current session."""
+        return _ExecutionHistory([job.__dict__ for job in self._metrics.jobs])
+
     @property
     def _allows_ambiguity(self) -> bool:
         return self._allow_ambiguity
@@ -432,7 +470,8 @@ def read_gbq(  # type: ignore[overload-overlap]
         col_order: Iterable[str] = ...,
         dry_run: Literal[False] = ...,
         allow_large_results: Optional[bool] = ...,
-    ) -> dataframe.DataFrame: ...
+    ) -> dataframe.DataFrame:
+        ...
 
     @overload
     def read_gbq(
@@ -448,7 +487,8 @@ def read_gbq(
         col_order: Iterable[str] = ...,
         dry_run: Literal[True] = ...,
         allow_large_results: Optional[bool] = ...,
-    ) -> pandas.Series: ...
+    ) -> pandas.Series:
+        ...
 
     def read_gbq(
         self,
@@ -520,7 +560,8 @@ def _read_gbq_colab(
         *,
         pyformat_args: Optional[Dict[str, Any]] = None,
         dry_run: Literal[False] = ...,
-    ) -> dataframe.DataFrame: ...
+    ) -> dataframe.DataFrame:
+        ...
 
     @overload
     def _read_gbq_colab(
@@ -529,7 +570,8 @@ def _read_gbq_colab(
         *,
         pyformat_args: Optional[Dict[str, Any]] = None,
         dry_run: Literal[True] = ...,
-    ) -> pandas.Series: ...
+    ) -> pandas.Series:
+        ...
 
     @log_adapter.log_name_override("read_gbq_colab")
     def _read_gbq_colab(
@@ -590,7 +632,8 @@ def read_gbq_query(  # type: ignore[overload-overlap]
         filters: third_party_pandas_gbq.FiltersType = ...,
         dry_run: Literal[False] = ...,
         allow_large_results: Optional[bool] = ...,
-    ) -> dataframe.DataFrame: ...
+    ) -> dataframe.DataFrame:
+        ...
 
     @overload
     def read_gbq_query(
@@ -606,7 +649,8 @@ def read_gbq_query(
         filters: third_party_pandas_gbq.FiltersType = ...,
         dry_run: Literal[True] = ...,
         allow_large_results: Optional[bool] = ...,
-    ) -> pandas.Series: ...
+    ) -> pandas.Series:
+        ...
 
     def read_gbq_query(
         self,
@@ -753,7 +797,8 @@ def read_gbq_table(  # type: ignore[overload-overlap]
         use_cache: bool = ...,
         col_order: Iterable[str] = ...,
         dry_run: Literal[False] = ...,
-    ) -> dataframe.DataFrame: ...
+    ) -> dataframe.DataFrame:
+        ...
 
     @overload
     def read_gbq_table(
@@ -767,7 +812,8 @@ def read_gbq_table(
         use_cache: bool = ...,
         col_order: Iterable[str] = ...,
         dry_run: Literal[True] = ...,
-    ) -> pandas.Series: ...
+    ) -> pandas.Series:
+        ...
 
     def read_gbq_table(
         self,
@@ -918,23 +964,26 @@ def read_pandas(
         pandas_dataframe: pandas.Index,
         *,
         write_engine: constants.WriteEngineType = "default",
-    ) -> bigframes.core.indexes.Index: ...
+    ) -> bigframes.core.indexes.Index:
+        ...
 
     @typing.overload
     def read_pandas(
         self,
         pandas_dataframe: pandas.Series,
         *,
         write_engine: constants.WriteEngineType = "default",
-    ) -> bigframes.series.Series: ...
+    ) -> bigframes.series.Series:
+        ...
 
     @typing.overload
     def read_pandas(
         self,
         pandas_dataframe: pandas.DataFrame,
         *,
         write_engine: constants.WriteEngineType = "default",
-    ) -> dataframe.DataFrame: ...
+    ) -> dataframe.DataFrame:
+        ...
 
     def read_pandas(
         self,

packages/bigframes/bigframes/session/loader.py

@@ -49,6 +49,8 @@
 import google.cloud.bigquery
 import google.cloud.bigquery as bigquery
 import google.cloud.bigquery.table
+from google.cloud.bigquery.job.load import LoadJob
+from google.cloud.bigquery.job.query import QueryJob
 import pandas
 import pyarrow as pa
 from google.cloud import bigquery_storage_v1
@@ -605,6 +607,9 @@ def _start_generic_job(self, job: formatting_helpers.GenericJob):
         else:
             job.result()
 
+        if self._metrics is not None and isinstance(job, (QueryJob, LoadJob)):
+            self._metrics.count_job_stats(query_job=job)
+
     @overload
     def read_gbq_table(  # type: ignore[overload-overlap]
         self,
@@ -626,7 +631,8 @@ def read_gbq_table(  # type: ignore[overload-overlap]
         n_rows: Optional[int] = None,
         index_col_in_columns: bool = False,
         publish_execution: bool = True,
-    ) -> dataframe.DataFrame: ...
+    ) -> dataframe.DataFrame:
+        ...
 
     @overload
     def read_gbq_table(
@@ -649,7 +655,8 @@ def read_gbq_table(
         n_rows: Optional[int] = None,
         index_col_in_columns: bool = False,
         publish_execution: bool = True,
-    ) -> pandas.Series: ...
+    ) -> pandas.Series:
+        ...
 
     def read_gbq_table(
         self,
@@ -1133,7 +1140,8 @@ def read_gbq_query(  # type: ignore[overload-overlap]
         dry_run: Literal[False] = ...,
         force_total_order: Optional[bool] = ...,
         allow_large_results: bool,
-    ) -> dataframe.DataFrame: ...
+    ) -> dataframe.DataFrame:
+        ...
 
     @overload
     def read_gbq_query(
@@ -1149,7 +1157,8 @@ def read_gbq_query(
         dry_run: Literal[True] = ...,
         force_total_order: Optional[bool] = ...,
         allow_large_results: bool,
-    ) -> pandas.Series: ...
+    ) -> pandas.Series:
+        ...
 
     def read_gbq_query(
         self,