feat: support automatic per-cell execution history filtering and isolated callbacks

packages/bigframes/bigframes/core/blocks.py

@@ -696,6 +696,7 @@ def to_pandas_batches(
         page_size: Optional[int] = None,
         max_results: Optional[int] = None,
         allow_large_results: Optional[bool] = None,
+        cell_execution_count: Optional[int] = None,
     ) -> PandasBatches:
         """Download results one message at a time.
 
@@ -713,6 +714,7 @@ def to_pandas_batches(
             execution_spec.ExecutionSpec(
                 promise_under_10gb=under_10gb,
                 ordered=True,
+                cell_execution_count=cell_execution_count,
             ),
         )
         result_batches = execution_result.batches()

packages/bigframes/bigframes/core/events.py

@@ -20,7 +20,7 @@
 import datetime
 import threading
 import uuid
-from typing import Any, Callable, Literal, Set
+from typing import Any, Callable, Literal, Optional, Set
 
 import google.cloud.bigquery._job_helpers
 import google.cloud.bigquery.job.query
@@ -127,8 +127,21 @@ class Event:
 
 @dataclasses.dataclass(frozen=True)
 class EventEnvelope:
+    """An envelope that wraps an execution event with metadata and display options.
+
+    Attributes:
+        event:
+            The actual execution event details (e.g., ExecutionStarted, BigQuerySentEvent).
+        progress_bar:
+            Specifies the style of progress bar to display during execution.
+        cell_execution_count:
+            The 1-indexed execution count of the notebook cell that triggered the event.
+            Used to group and filter execution history on a per-cell basis.
+    """
+
     event: Event
     progress_bar: ProgressBarType = _DEFAULT
+    cell_execution_count: Optional[int] = None
 
 
 @dataclasses.dataclass(frozen=True)

packages/bigframes/bigframes/core/global_session.py

@@ -19,7 +19,7 @@
 import threading
 import traceback
 import warnings
-from typing import TYPE_CHECKING, Callable, Optional, TypeVar
+from typing import TYPE_CHECKING, Callable, Iterable, Optional, TypeVar
 
 import google.auth.exceptions
 
@@ -124,12 +124,20 @@ def with_default_session(func_: Callable[..., _T], *args, **kwargs) -> _T:
     return func_(get_global_session(), *args, **kwargs)
 
 
-def execution_history() -> "bigframes.session._ExecutionHistory":
-    import pandas  # noqa: F401
-
+def execution_history(
+    *,
+    events: Optional[Iterable[bigframes.core.events.Event]] = None,
+    job_ids: Optional[Iterable[str]] = None,
+    all_cells: bool = True,
+) -> "bigframes.session._ExecutionHistory":
     import bigframes.session
 
-    return with_default_session(bigframes.session.Session.execution_history)
+    return with_default_session(
+        bigframes.session.Session.execution_history,
+        events=events,
+        job_ids=job_ids,
+        all_cells=all_cells,
+    )
 
 
 class _GlobalSessionContext:

packages/bigframes/bigframes/core/utils.py

@@ -249,3 +249,16 @@ def timedelta_to_micros(
         ) * 1_000_000 + timedelta.microseconds
 
     raise TypeError(f"Unrecognized input type: {type(timedelta)}")
+
+
+def get_ipython_execution_count() -> typing.Optional[int]:
+    """Returns the current IPython cell execution count if running in a notebook, else None."""
+    try:
+        import IPython
+
+        ipy = IPython.get_ipython()
+        if ipy is not None and hasattr(ipy, "execution_count"):
+            return ipy.execution_count
+    except (ImportError, NameError):
+        pass
+    return None

packages/bigframes/bigframes/dataframe.py

@@ -1755,6 +1755,7 @@ def to_pandas_batches(
         max_results: Optional[int] = None,
         *,
         allow_large_results: Optional[bool] = None,
+        cell_execution_count: Optional[int] = None,
     ) -> blocks.PandasBatches:
         """Stream DataFrame results to an iterable of pandas DataFrame.
 
@@ -1807,6 +1808,7 @@ def to_pandas_batches(
             page_size=page_size,
             max_results=max_results,
             allow_large_results=allow_large_results,
+            cell_execution_count=cell_execution_count,
         )
 
     def _to_pandas_batches(
@@ -1815,11 +1817,13 @@ def _to_pandas_batches(
         max_results: Optional[int] = None,
         *,
         allow_large_results: Optional[bool] = None,
+        cell_execution_count: Optional[int] = None,
     ) -> blocks.PandasBatches:
         return self._block.to_pandas_batches(
             page_size=page_size,
             max_results=max_results,
             allow_large_results=allow_large_results,
+            cell_execution_count=cell_execution_count,
         )
 
     def _compute_dry_run(self) -> google.cloud.bigquery.job.QueryJob:

packages/bigframes/bigframes/display/anywidget.py

@@ -92,6 +92,10 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
 
         self._dataframe = dataframe
 
+        from bigframes.core.utils import get_ipython_execution_count
+
+        self._cell_execution_count = get_ipython_execution_count()
+
         super().__init__()
 
         # Initialize attributes that might be needed by observers first
@@ -286,7 +290,10 @@ def _reset_batch_cache(self) -> None:
     def _reset_batches_for_new_page_size(self) -> None:
         """Reset the batch iterator when page size changes."""
         with bigframes.option_context("display.progress_bar", None):
-            self._batches = self._dataframe.to_pandas_batches(page_size=self.page_size)
+            self._batches = self._dataframe.to_pandas_batches(
+                page_size=self.page_size,
+                cell_execution_count=self._cell_execution_count,
+            )
 
         self._reset_batch_cache()
 
@@ -318,7 +325,8 @@ def _set_table_html(self) -> None:
             current_sort_state = _SortState(tuple(sort_columns), tuple(sort_ascending))
             if self._last_sort_state != current_sort_state:
                 self._batches = df_to_display.to_pandas_batches(
-                    page_size=self.page_size
+                    page_size=self.page_size,
+                    cell_execution_count=self._cell_execution_count,
                 )
                 self._reset_batch_cache()
                 self._last_sort_state = current_sort_state

packages/bigframes/bigframes/pandas/io/api.py

@@ -300,23 +300,26 @@ def _try_read_gbq_colab_sessionless_dry_run(
 def _read_gbq_colab(  # type: ignore[overload-overlap]
     query_or_table: str,
     *,
-    pyformat_args: Optional[Dict[str, Any]] = ...,
-    dry_run: Literal[False] = ...,
+    callback: Optional[Callable[[bigframes.core.events.EventEnvelope], None]] = None,
+    pyformat_args: Optional[Dict[str, Any]] = None,
+    dry_run: Literal[False] = False,
 ) -> bigframes.dataframe.DataFrame: ...
 
 
 @overload
 def _read_gbq_colab(
     query_or_table: str,
     *,
-    pyformat_args: Optional[Dict[str, Any]] = ...,
-    dry_run: Literal[True] = ...,
+    callback: Optional[Callable[[bigframes.core.events.EventEnvelope], None]] = None,
+    pyformat_args: Optional[Dict[str, Any]] = None,
+    dry_run: Literal[True],
 ) -> pandas.Series: ...
 
 
 def _read_gbq_colab(
     query_or_table: str,
     *,
+    callback: Optional[Callable[[bigframes.core.events.EventEnvelope], None]] = None,
     pyformat_args: Optional[Dict[str, Any]] = None,
     dry_run: bool = False,
 ) -> bigframes.dataframe.DataFrame | pandas.Series:
@@ -328,6 +331,8 @@ def _read_gbq_colab(
     Args:
         query_or_table (str):
             SQL query or table ID (table ID not yet supported).
+        callback (Optional[Callable[[bigframes.core.events.EventEnvelope], None]]):
+            Callback to receive query execution events.
         pyformat_args (Optional[Dict[str, Any]]):
             Parameters to format into the query string.
         dry_run (bool):
@@ -379,6 +384,7 @@ def _read_gbq_colab(
     return global_session.with_default_session(
         bigframes.session.Session._read_gbq_colab,
         query_or_table,
+        callback=callback,
         pyformat_args=pyformat_args,
         dry_run=dry_run,
     )

packages/bigframes/bigframes/series.py

@@ -760,6 +760,7 @@ def to_pandas_batches(
         max_results: Optional[int] = None,
         *,
         allow_large_results: Optional[bool] = None,
+        cell_execution_count: Optional[int] = None,
     ) -> Iterable[pandas.Series]:
         """Stream Series results to an iterable of pandas Series.
 
@@ -812,6 +813,7 @@ def to_pandas_batches(
             page_size=page_size,
             max_results=max_results,
             allow_large_results=allow_large_results,
+            cell_execution_count=cell_execution_count,
         )
         return map(lambda df: cast(pandas.Series, df.squeeze(1)), batches)