perf: Make executor data uploads async internally (#2529)

Author: TrevorBergeron

bigframes/session/bq_caching_executor.py

@@ -14,6 +14,7 @@
 
 from __future__ importannotations
 
+importconcurrent.futures
 importmath
 importthreading
 fromtypingimportLiteral, Mapping, Optional, Sequence, Tuple
@@ -514,13 +515,29 @@ def _substitute_large_local_sources(self, original_root: nodes.BigFrameNode):
  Replace large local sources with the uploaded version of those datasources.
  """
 # Step 1: Upload all previously un-uploaded data
+needs_upload= []
 forleafinoriginal_root.unique_nodes():
 ifisinstance(leaf, nodes.ReadLocalNode):
 if (
 leaf.local_data_source.metadata.total_bytes
 >bigframes.constants.MAX_INLINE_BYTES
  ):
-self._upload_local_data(leaf.local_data_source)
+needs_upload.append(leaf.local_data_source)
+
+futures: dict[concurrent.futures.Future, local_data.ManagedArrowTable] =dict()
+forlocal_sourceinneeds_upload:
+future=self.loader.read_data_async(
+local_source, bigframes.core.guid.generate_guid()
+ )
+futures[future] =local_source
+try:
+forfutureinconcurrent.futures.as_completed(futures.keys()):
+self.cache.cache_remote_replacement(futures[future], future.result())
+exceptExceptionase:
+# cancel all futures
+forfutureinfutures:
+future.cancel()
+raisee
 
 # Step 2: Replace local scans with remote scans
 defmap_local_scans(node: nodes.BigFrameNode):
@@ -550,18 +567,6 @@ def map_local_scans(node: nodes.BigFrameNode):
 
 returnoriginal_root.bottom_up(map_local_scans)
 
-def_upload_local_data(self, local_table: local_data.ManagedArrowTable):
-ifself.cache.get_uploaded_local_data(local_table) isnotNone:
-return
-# Lock prevents concurrent repeated work, but slows things down.
-# Might be better as a queue and a worker thread
-withself._upload_lock:
-ifself.cache.get_uploaded_local_data(local_table) isNone:
-uploaded=self.loader.load_data_or_write_data(
-local_table, bigframes.core.guid.generate_guid()
- )
-self.cache.cache_remote_replacement(local_table, uploaded)
-
 def_execute_plan_gbq(
 self,
 plan: nodes.BigFrameNode,

bigframes/session/loader.py

@@ -300,6 +300,17 @@ def __init__(
 self._session=session
 self._clock=session_time.BigQuerySyncedClock(bqclient)
 self._clock.sync()
+self._threadpool=concurrent.futures.ThreadPoolExecutor(
+max_workers=1, thread_name_prefix="bigframes-loader"
+ )
+
+defread_data_async(
+self, local_data: local_data.ManagedArrowTable, offsets_col: str
+ ) ->concurrent.futures.Future[bq_data.BigqueryDataSource]:
+future=self._threadpool.submit(
+self._load_data_or_write_data, local_data, offsets_col
+ )
+returnfuture
 
 defread_pandas(
 self,
@@ -350,7 +361,7 @@ def read_managed_data(
 session=self._session,
  )
 
-defload_data_or_write_data(
+def_load_data_or_write_data(
 self,
 data: local_data.ManagedArrowTable,
 offsets_col: str,