feat: propagate cell execution count to widget queries and rename history filter

Author: shuoweil

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/global_session.py

@@ -128,7 +128,7 @@ def execution_history(
     *,
     events: Optional[Iterable[bigframes.core.events.Event]] = None,
     job_ids: Optional[Iterable[str]] = None,
-    filter_by_cell: bool = True,
+    current_cell_only: bool = True,
 ) -> "bigframes.session._ExecutionHistory":
     import pandas  # noqa: F401
 
@@ -138,7 +138,7 @@ def execution_history(
         bigframes.session.Session.execution_history,
         events=events,
         job_ids=job_ids,
-        filter_by_cell=filter_by_cell,
+        current_cell_only=current_cell_only,
     )
 
 

packages/bigframes/bigframes/dataframe.py

@@ -1739,6 +1739,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.
 
@@ -1791,6 +1792,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(
@@ -1799,11 +1801,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,16 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
 
         self._dataframe = dataframe
 
+        self._cell_execution_count = None
+        try:
+            import IPython
+
+            ipy = IPython.get_ipython()
+            if ipy is not None and hasattr(ipy, "execution_count"):
+                self._cell_execution_count = ipy.execution_count
+        except (ImportError, NameError):
+            pass
+
         super().__init__()
 
         # Initialize attributes that might be needed by observers first
@@ -286,7 +296,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 +331,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/series.py

@@ -756,6 +756,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.
 
@@ -808,6 +809,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)
 

packages/bigframes/bigframes/session/__init__.py

@@ -448,7 +448,7 @@ def execution_history(
         *,
         events: Optional[Iterable[bigframes.core.events.Event]] = None,
         job_ids: Optional[Iterable[str]] = None,
-        filter_by_cell: bool = True,
+        current_cell_only: bool = True,
     ) -> _ExecutionHistory:
         """Returns the history of executions initiated by BigFrames in the current session.
 
@@ -459,7 +459,7 @@ def execution_history(
                 Filter execution history to only include jobs associated with the given events.
             job_ids (Iterable[str], optional):
                 Filter execution history to only include jobs matching the given job IDs.
-            filter_by_cell (bool, optional):
+            current_cell_only (bool, optional):
                 If True and running in Colab/Jupyter, automatically filter history to only include
                 jobs executed within the current cell. Defaults to True.
         """
@@ -503,7 +503,7 @@ def execution_history(
                 )
             ]
 
-        elif filter_by_cell:
+        elif current_cell_only:
             try:
                 import IPython
 

packages/bigframes/tests/unit/display/test_anywidget.py

@@ -179,3 +179,25 @@ def test_page_size_change_resets_sort(mock_df):
 
     # to_pandas_batches called again (reset)
     assert mock_df.to_pandas_batches.call_count >= 2
+
+
+def test_cell_execution_count_propagation(mock_df):
+    """Test that the captured cell_execution_count is propagated to to_pandas_batches."""
+    from bigframes.display.anywidget import TableWidget
+
+    # Mock IPython to return a specific execution count
+    mock_ipy = mock.Mock()
+    mock_ipy.execution_count = 42
+
+    with mock.patch("IPython.get_ipython", return_value=mock_ipy):
+        with bigframes.option_context("display.render_mode", "anywidget"):
+            widget = TableWidget(mock_df)
+
+    # Verify captured execution count
+    assert widget._cell_execution_count == 42
+
+    # to_pandas_batches should be called with cell_execution_count=42
+    mock_df.to_pandas_batches.assert_called_with(
+        page_size=widget.page_size,
+        cell_execution_count=42,
+    )