feat: universal per-cell execution history filtering

Author: shuoweil

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
@@ -129,6 +129,7 @@ class Event:
 class EventEnvelope:
     event: Event
     progress_bar: ProgressBarType = _DEFAULT
+    cell_execution_count: Optional[int] = None
 
 
 @dataclasses.dataclass(frozen=True)

packages/bigframes/bigframes/session/_io/bigquery/__init__.py

@@ -263,7 +263,7 @@ def add_and_trim_labels(
     )
 
 
-def create_bq_event_callback(publisher):
+def create_bq_event_callback(publisher, cell_execution_count=None):
     event_map = {
         google.cloud.bigquery._job_helpers.QueryFinishedEvent: (
             bigframes.core.events.BigQueryFinishedEvent
@@ -286,7 +286,9 @@ def publish_bq_event(event):
                 bf_event = bf_type.from_bqclient(event)  # type: ignore
                 break
         envelope = bigframes.core.events.EventEnvelope(
-            event=bf_event, progress_bar=bigframes.core.events._DEFAULT
+            event=bf_event,
+            progress_bar=bigframes.core.events._DEFAULT,
+            cell_execution_count=cell_execution_count,
         )
         publisher.publish(envelope)
 
@@ -309,10 +311,21 @@ def start_query_with_job(
     job_retry: google.api_core.retry.Retry = (third_party_gcb_retry.DEFAULT_JOB_RETRY),  # noqa: E501
     publisher: bigframes.core.events.Publisher,
     session=None,
+    cell_execution_count: Optional[int] = None,
 ) -> Tuple[google.cloud.bigquery.table.RowIterator, bigquery.QueryJob]:
     """
     Starts query job and waits for results.
     """
+    if cell_execution_count is None:
+        try:
+            import IPython
+
+            ipy = IPython.get_ipython()
+            if ipy is not None and hasattr(ipy, "execution_count"):
+                cell_execution_count = ipy.execution_count
+        except (ImportError, NameError):
+            pass
+
     # Note: Ensure no additional labels are added to job_config after this
     # point, as `add_and_trim_labels` ensures the label count does not
     # exceed MAX_LABELS_COUNT.
@@ -339,6 +352,7 @@ def start_query_with_job(
         sql=sql,
         publisher=publisher,
         metrics=metrics,
+        cell_execution_count=cell_execution_count,
     )
     return results_iterator, query_job
 
@@ -359,13 +373,24 @@ def start_query_job_optional(
     job_retry: google.api_core.retry.Retry = (third_party_gcb_retry.DEFAULT_JOB_RETRY),  # noqa: E501
     publisher: bigframes.core.events.Publisher,
     session=None,
+    cell_execution_count: Optional[int] = None,
 ) -> google.cloud.bigquery.table.RowIterator:
     """
     Run a bigquery query, with job optional.
 
     See:
     https://docs.cloud.google.com/bigquery/docs/running-queries#optional-job-creation
     """
+    if cell_execution_count is None:
+        try:
+            import IPython
+
+            ipy = IPython.get_ipython()
+            if ipy is not None and hasattr(ipy, "execution_count"):
+                cell_execution_count = ipy.execution_count
+        except (ImportError, NameError):
+            pass
+
     add_and_trim_labels(job_config, session=session)
     try:
         results_iterator = bq_client._query_and_wait_bigframes(
@@ -375,10 +400,14 @@ def start_query_job_optional(
             project=project,
             api_timeout=timeout,
             job_retry=job_retry,
-            callback=create_bq_event_callback(publisher),
+            callback=create_bq_event_callback(
+                publisher, cell_execution_count=cell_execution_count
+            ),
         )
         if metrics is not None:
-            metrics.count_job_stats(row_iterator=results_iterator)
+            metrics.count_job_stats(
+                row_iterator=results_iterator, cell_execution_count=cell_execution_count
+            )
         return results_iterator
     except google.api_core.exceptions.Forbidden as ex:
         if "Drive credentials" in ex.message:
@@ -392,35 +421,45 @@ def _publish_events(
     total_rows: Optional[int],
     publisher: bigframes.core.events.Publisher,
     metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
+    cell_execution_count: Optional[int] = None,
 ):
     if not query_job.configuration.dry_run:
         publisher.publish(
-            bigframes.core.events.BigQuerySentEvent(
-                sql,
-                billing_project=query_job.project,
-                location=query_job.location,
-                job_id=query_job.job_id,
-                request_id=None,
+            bigframes.core.events.EventEnvelope(
+                event=bigframes.core.events.BigQuerySentEvent(
+                    sql,
+                    billing_project=query_job.project,
+                    location=query_job.location,
+                    job_id=query_job.job_id,
+                    request_id=None,
+                ),
+                cell_execution_count=cell_execution_count,
             )
         )
     if not query_job.configuration.dry_run:
         publisher.publish(
-            bigframes.core.events.BigQueryFinishedEvent(
-                billing_project=query_job.project,
-                location=query_job.location,
-                job_id=query_job.job_id,
-                destination=query_job.destination,
-                total_rows=total_rows,
-                total_bytes_processed=query_job.total_bytes_processed,
-                slot_millis=query_job.slot_millis,
-                created=query_job.created,
-                started=query_job.started,
-                ended=query_job.ended,
+            bigframes.core.events.EventEnvelope(
+                event=bigframes.core.events.BigQueryFinishedEvent(
+                    billing_project=query_job.project,
+                    location=query_job.location,
+                    query_id=query_job.query_id,
+                    job_id=query_job.job_id,
+                    destination=query_job.destination,
+                    total_rows=total_rows,
+                    total_bytes_processed=query_job.total_bytes_processed,
+                    slot_millis=query_job.slot_millis,
+                    created=query_job.created,
+                    started=query_job.started,
+                    ended=query_job.ended,
+                ),
+                cell_execution_count=cell_execution_count,
             )
         )
 
     if metrics is not None:
-        metrics.count_job_stats(query_job=query_job)
+        metrics.count_job_stats(
+            query_job=query_job, cell_execution_count=cell_execution_count
+        )
 
 
 def delete_tables_matching_session_id(

packages/bigframes/bigframes/session/bq_caching_executor.py

@@ -208,8 +208,9 @@ async def _execute_async(
             execution_spec,
         )
         await self._publisher.publish_async(
-            bigframes.core.events.ExecutionFinished(
-                result=result,
+            bigframes.core.events.EventEnvelope(
+                event=bigframes.core.events.ExecutionFinished(result=result),
+                cell_execution_count=execution_spec.cell_execution_count,
             )
         )
         return result
@@ -224,8 +225,11 @@ async def _try_execute_semi_executors(
             maybe_result = await exec.execute(plan, execution_spec)
             if maybe_result:
                 await self._publisher.publish_async(
-                    bigframes.core.events.ExecutionFinished(
-                        result=maybe_result,
+                    bigframes.core.events.EventEnvelope(
+                        event=bigframes.core.events.ExecutionFinished(
+                            result=maybe_result,
+                        ),
+                        cell_execution_count=execution_spec.cell_execution_count,
                     )
                 )
                 return maybe_result

packages/bigframes/bigframes/session/direct_gbq_execution.py

@@ -109,6 +109,7 @@ async def execute(
             job_config=job_config,
             query_with_job=(not can_skip_job),
             session=plan.session,
+            cell_execution_count=spec.cell_execution_count,
         )
         result_bq_data = None
         if query_job and query_job.destination:
@@ -158,6 +159,7 @@ def _run_execute_query(
         job_config: bq_job.QueryJobConfig,
         query_with_job: bool,
         session,
+        cell_execution_count: Optional[int] = None,
     ) -> Tuple[bq_table.RowIterator, Optional[bigquery.QueryJob]]:
         """
         Starts BigQuery query job and waits for results.
@@ -171,6 +173,7 @@ def _run_execute_query(
                     metrics=self._metrics,
                     publisher=self._publisher,
                     session=session,
+                    cell_execution_count=cell_execution_count,
                 )
             else:
                 return (
@@ -181,6 +184,7 @@ def _run_execute_query(
                         metrics=self._metrics,
                         publisher=self._publisher,
                         session=session,
+                        cell_execution_count=cell_execution_count,
                     ),
                     None,
                 )

packages/bigframes/bigframes/session/execution_spec.py

@@ -60,6 +60,7 @@ class ExecutionSpec:
 
     # BigQuery specific options
     bigquery_config: Optional[BqComputeOptions] = None
+    cell_execution_count: Optional[int] = None
 
     def with_bq_labels(self, labels: Mapping[str, str]) -> ExecutionSpec:
         bq_config = self.bigquery_config or BqComputeOptions()
@@ -77,7 +78,23 @@ def with_compute_options(self, compute_options: ComputeOptions) -> ExecutionSpec
             new_bq_config = new_bq_config.push_labels(
                 dict(self.bigquery_config.extra_query_labels)
             )
-        return dataclasses.replace(self, bigquery_config=new_bq_config)
+
+        cell_execution_count = self.cell_execution_count
+        if cell_execution_count is None:
+            try:
+                import IPython
+
+                ipy = IPython.get_ipython()
+                if ipy is not None and hasattr(ipy, "execution_count"):
+                    cell_execution_count = ipy.execution_count
+            except (ImportError, NameError):
+                pass
+
+        return dataclasses.replace(
+            self,
+            bigquery_config=new_bq_config,
+            cell_execution_count=cell_execution_count,
+        )
 
 
 # Used internally by execution