refactor(bigframes): make executor internals async (#17093)

Author: TrevorBergeron

packages/bigframes/bigframes/core/events.py

@@ -14,6 +14,8 @@
 
 from __future__ import annotations
 
+import asyncio
+import concurrent.futures
 import dataclasses
 import datetime
 import threading
@@ -68,6 +70,9 @@ class Publisher:
     def __init__(self):
         self._subscribers_lock = threading.Lock()
         self._subscribers: Set[Subscriber] = set()
+        self._executor: concurrent.futures.Executor = (
+            concurrent.futures.ThreadPoolExecutor()
+        )
 
     def subscribe(self, callback: Callable[[Event], None]) -> Subscriber:
         # TODO(b/448176657): figure out how to handle subscribers/publishers in
@@ -86,6 +91,16 @@ def publish(self, event: Event):
             for subscriber in self._subscribers:
                 subscriber(event)
 
+    async def publish_async(self, event: Event):
+        with self._subscribers_lock:
+            subscribers_snapshot = list(self._subscribers)
+        loop = asyncio.get_running_loop()
+        tasks = [
+            loop.run_in_executor(self._executor, subscriber, event)
+            for subscriber in subscribers_snapshot
+        ]
+        return await asyncio.gather(*tasks, return_exceptions=True)
+
 
 class Event:
     pass

packages/bigframes/bigframes/core/indexes/base.py

@@ -293,7 +293,8 @@ def get_loc(self, key) -> typing.Union[int, slice, "bigframes.series.Series"]:
 
         count_scalar = (
             self._block.session._executor.execute(
-                count_result, ex_spec.ExecutionSpec(promise_under_10gb=True)
+                count_result,
+                ex_spec.ExecutionSpec(promise_under_10gb=True),
             )
             .batches()
             .to_py_scalar()
@@ -308,7 +309,8 @@ def get_loc(self, key) -> typing.Union[int, slice, "bigframes.series.Series"]:
             position_result = filtered_block._expr.aggregate([(min_agg, "position")])
             position_scalar = (
                 self._block.session._executor.execute(
-                    position_result, ex_spec.ExecutionSpec(promise_under_10gb=True)
+                    position_result,
+                    ex_spec.ExecutionSpec(promise_under_10gb=True),
                 )
                 .batches()
                 .to_py_scalar()

packages/bigframes/bigframes/core/local_data.py

@@ -248,10 +248,11 @@ def iter_array(
         elif dtype == bigframes.dtypes.TIMEDELTA_DTYPE:
             if duration_type == "int":
                 yield from map(
-                    lambda x: ((x.days * 3600 * 24) + x.seconds) * 1_000_000
-                    + x.microseconds
-                    if x is not None
-                    else x,
+                    lambda x: (
+                        ((x.days * 3600 * 24) + x.seconds) * 1_000_000 + x.microseconds
+                        if x is not None
+                        else x
+                    ),
                     values,
                 )
             else:

packages/bigframes/bigframes/formatting_helpers.py

@@ -26,6 +26,8 @@
 import google.cloud.bigquery as bigquery
 import humanize
 
+import bigframes._config
+
 if TYPE_CHECKING:
     import bigframes.core.events
 
@@ -137,78 +139,85 @@ def repr_query_job_html(query_job: Optional[bigquery.QueryJob]):
 current_display_id: Optional[str] = None
 
 
-def progress_callback(
-    event: bigframes.core.events.Event,
-):
-    """Displays a progress bar while the query is running"""
-    global current_display_id
+def create_progress_callback():
+    # bind potentially thread-local config to the callback so that it uses the user thread
+    # config even if callback is invoked from a worker thread.
+    display_opts = bigframes._config.options.display
 
-    try:
-        import bigframes._config
-        import bigframes.core.events
-    except ImportError:
-        # Since this gets called from __del__, skip if the import fails to avoid
-        # ImportError: sys.meta_path is None, Python is likely shutting down.
-        # This will allow cleanup to continue.
-        return
+    def progress_callback(
+        event: bigframes.core.events.Event,
+    ):
+        """Displays a progress bar while the query is running"""
+        global current_display_id
 
-    progress_bar = bigframes._config.options.display.progress_bar
+        try:
+            import bigframes._config
+            import bigframes.core.events
+        except ImportError:
+            # Since this gets called from __del__, skip if the import fails to avoid
+            # ImportError: sys.meta_path is None, Python is likely shutting down.
+            # This will allow cleanup to continue.
+            return
 
-    if progress_bar == "auto":
-        progress_bar = "notebook" if in_ipython() else "terminal"
+        progress_bar = display_opts.progress_bar
+
+        if progress_bar == "auto":
+            progress_bar = "notebook" if in_ipython() else "terminal"
 
-    if progress_bar == "notebook":
-        import IPython.display as display
+        if progress_bar == "notebook":
+            import IPython.display as display
 
-        display_html = None
+            display_html = None
 
-        if isinstance(event, bigframes.core.events.ExecutionStarted):
-            # Start a new context for progress output.
-            current_display_id = None
+            if isinstance(event, bigframes.core.events.ExecutionStarted):
+                # Start a new context for progress output.
+                current_display_id = None
 
-        elif isinstance(event, bigframes.core.events.BigQuerySentEvent):
-            display_html = render_bqquery_sent_event_html(event)
+            elif isinstance(event, bigframes.core.events.BigQuerySentEvent):
+                display_html = render_bqquery_sent_event_html(event)
 
-        elif isinstance(event, bigframes.core.events.BigQueryRetryEvent):
-            display_html = render_bqquery_retry_event_html(event)
+            elif isinstance(event, bigframes.core.events.BigQueryRetryEvent):
+                display_html = render_bqquery_retry_event_html(event)
 
-        elif isinstance(event, bigframes.core.events.BigQueryReceivedEvent):
-            display_html = render_bqquery_received_event_html(event)
+            elif isinstance(event, bigframes.core.events.BigQueryReceivedEvent):
+                display_html = render_bqquery_received_event_html(event)
 
-        elif isinstance(event, bigframes.core.events.BigQueryFinishedEvent):
-            display_html = render_bqquery_finished_event_html(event)
+            elif isinstance(event, bigframes.core.events.BigQueryFinishedEvent):
+                display_html = render_bqquery_finished_event_html(event)
 
-        elif isinstance(event, bigframes.core.events.SessionClosed):
-            display_html = f"Session {event.session_id} closed."
+            elif isinstance(event, bigframes.core.events.SessionClosed):
+                display_html = f"Session {event.session_id} closed."
 
-        if display_html:
-            if current_display_id:
-                display.update_display(
-                    display.HTML(display_html),
-                    display_id=current_display_id,
-                )
-            else:
-                current_display_id = str(random.random())
-                display.display(
-                    display.HTML(display_html),
-                    display_id=current_display_id,
-                )
+            if display_html:
+                if current_display_id:
+                    display.update_display(
+                        display.HTML(display_html),
+                        display_id=current_display_id,
+                    )
+                else:
+                    current_display_id = str(random.random())
+                    display.display(
+                        display.HTML(display_html),
+                        display_id=current_display_id,
+                    )
 
-    elif progress_bar == "terminal":
-        message = None
-
-        if isinstance(event, bigframes.core.events.BigQuerySentEvent):
-            message = render_bqquery_sent_event_plaintext(event)
-            print(message)
-        elif isinstance(event, bigframes.core.events.BigQueryRetryEvent):
-            message = render_bqquery_retry_event_plaintext(event)
-            print(message)
-        elif isinstance(event, bigframes.core.events.BigQueryReceivedEvent):
-            message = render_bqquery_received_event_plaintext(event)
-            print(message)
-        elif isinstance(event, bigframes.core.events.BigQueryFinishedEvent):
-            message = render_bqquery_finished_event_plaintext(event)
-            print(message)
+        elif progress_bar == "terminal":
+            message = None
+
+            if isinstance(event, bigframes.core.events.BigQuerySentEvent):
+                message = render_bqquery_sent_event_plaintext(event)
+                print(message)
+            elif isinstance(event, bigframes.core.events.BigQueryRetryEvent):
+                message = render_bqquery_retry_event_plaintext(event)
+                print(message)
+            elif isinstance(event, bigframes.core.events.BigQueryReceivedEvent):
+                message = render_bqquery_received_event_plaintext(event)
+                print(message)
+            elif isinstance(event, bigframes.core.events.BigQueryFinishedEvent):
+                message = render_bqquery_finished_event_plaintext(event)
+                print(message)
+
+    return progress_callback
 
 
 def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None):

packages/bigframes/bigframes/operations/python_op_maps.py

@@ -71,11 +71,11 @@
     abs: numeric_ops.abs_op,
     pow: numeric_ops.pow_op,
     ### builtins -- iterable
-    all: array_ops.ArrayReduceOp(aggregations.all_op),
-    any: array_ops.ArrayReduceOp(aggregations.any_op),
-    sum: array_ops.ArrayReduceOp(aggregations.sum_op),
-    min: array_ops.ArrayReduceOp(aggregations.min_op),
-    max: array_ops.ArrayReduceOp(aggregations.max_op),
+    all: array_ops.ArrayReduceOp(aggregations.all_op),  # type: ignore
+    any: array_ops.ArrayReduceOp(aggregations.any_op),  # type: ignore
+    sum: array_ops.ArrayReduceOp(aggregations.sum_op),  # type: ignore
+    min: array_ops.ArrayReduceOp(aggregations.min_op),  # type: ignore
+    max: array_ops.ArrayReduceOp(aggregations.max_op),  # type: ignore
 }
 
 

packages/bigframes/bigframes/session/__init__.py

@@ -183,7 +183,9 @@ def __init__(
         # Publisher needs to be created before the other objects, especially
         # the executors, because they access it.
         self._publisher = bigframes.core.events.Publisher()
-        self._publisher.subscribe(bigframes.formatting_helpers.progress_callback)
+        self._publisher.subscribe(
+            bigframes.formatting_helpers.create_progress_callback()
+        )
 
         if context is None:
             context = bigquery_options.BigQueryOptions()

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

@@ -55,11 +55,6 @@ def create_job_configs_labels(
     if job_configs_labels is None:
         job_configs_labels = {}
 
-    # If the user has labels they wish to set, make sure we set those first so
-    # they are preserved.
-    for key, value in bigframes.options.compute.extra_query_labels.items():
-        job_configs_labels[key] = value
-
     if api_methods and "bigframes-api" not in job_configs_labels:
         job_configs_labels["bigframes-api"] = api_methods[0]
         del api_methods[0]
@@ -230,7 +225,9 @@ def format_option(key: str, value: Union[bool, str]) -> str:
     return f"{key}={repr(value)}"
 
 
-def add_and_trim_labels(job_config, session=None):
+def add_and_trim_labels(
+    job_config, session=None, extra_query_labels: Optional[Mapping[str, str]] = None
+):
     """
     Add additional labels to the job configuration and trim the total number of labels
     to ensure they do not exceed MAX_LABELS_COUNT labels per job.

packages/bigframes/bigframes/session/bigquery_session.py

@@ -119,7 +119,11 @@ def close(self):
             bfbqio.start_query_job_optional(
                 self.bqclient,
                 f"CALL BQ.ABORT_SESSION('{self._session_id}')",
-                job_config=bigquery.QueryJobConfig(),
+                # Assume this is being called in the user thread, so we can access
+                # this thread-local config.
+                job_config=bigquery.QueryJobConfig(
+                    labels=bigframes.options.compute.extra_query_labels
+                ),
                 location=self.location,
                 project=None,
                 timeout=None,