WIP: Add progress bar when BQ Storage API is used.

tswast · tswast · commit 0d96dbd0a382 · 2019-04-11T14:30:37.000-07:00
diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py
@@ -21,6 +21,7 @@
 import datetime
 import json
 import operator
+import queue
 import warnings
 
 import six
@@ -56,6 +57,7 @@
 )
 _TABLE_HAS_NO_SCHEMA = 'Table has no schema:  call "client.get_table()"'
 _MARKER = object()
+_PROGRESS_INTERVAL = 1.0  # Time between tqdm updates, in seconds.
 
 
 def _reference_getter(table):
@@ -1254,6 +1256,16 @@ def __repr__(self):
         return "Row({}, {})".format(self._xxx_values, f2i)
 
 
+class _FakeQueue(object):
+    """A fake Queue class that does nothing.
+
+    This is used when there is no progress bar to send updates to.
+    """
+
+    def put(self, item):
+        """Don't actually do anything with the item."""
+
+
 class RowIterator(HTTPIterator):
     """A class for iterating through HTTP/JSON API row list responses.
 
@@ -1367,7 +1379,7 @@ def _to_dataframe_tabledata_list(self, dtypes, progress_bar=None):
 
         return pandas.concat(frames)
 
-    def _to_dataframe_bqstorage(self, bqstorage_client, dtypes):
+    def _to_dataframe_bqstorage(self, bqstorage_client, dtypes, progress_bar=None):
         """Use (faster, but billable) BQ Storage API to construct DataFrame."""
         import concurrent.futures
         from google.cloud import bigquery_storage_v1beta1
@@ -1402,16 +1414,57 @@ def _to_dataframe_bqstorage(self, bqstorage_client, dtypes):
         if not session.streams:
             return pandas.DataFrame(columns=columns)
 
+        exit_early = False
+        progress_queue = _FakeQueue()
+        if progress_bar is not None:
+            progress_queue = queue.Queue()
+
         def get_dataframe(stream):
             position = bigquery_storage_v1beta1.types.StreamPosition(stream=stream)
-            rowstream = bqstorage_client.read_rows(position)
-            return rowstream.to_dataframe(session, dtypes=dtypes)
+            rowstream = bqstorage_client.read_rows(position).rows(session)
+
+            frames = []
+            for page in rowstream.pages:
+                if exit_early:
+                    return pandas.DataFrame(columns=columns)
+                frames.append(page.to_dataframe(dtypes=dtypes))
+                progress_queue.put(page.num_items)
+
+            # page.to_dataframe() does not preserve column order. Rearrange at
+            # the end using manually-parsed schema.
+            return pandas.concat(frames)[columns]
+
+        def update_progress_bar():
+            try:
+                while True:
+                    next_update = progress_queue.get_nowait()
+                    progress_bar.update(next_update)
+            except queue.Empty:
+                pass
+
+        def frames_with_progress_bar(pool):
+            # Write progress to main thread.
+            frames = []
+            not_done  = [pool.submit(get_dataframe, stream) for stream in session.streams]
+
+            while not_done:
+                done, not_done = concurrent.futures.wait(not_done, timeout=_PROGRESS_INTERVAL)
+                frames.extend([future.result() for future in done])
+                update_progress_bar()
+
+            return frames
 
         with concurrent.futures.ThreadPoolExecutor() as pool:
-            frames = pool.map(get_dataframe, session.streams)
+            try:
+                if progress_bar is None:
+                    frames = pool.map(get_dataframe, session.streams)
+                else:
+                    frames = frames_with_progress_bar(pool)
+            except:
+                exit_early = True
+                raise
 
-        # rowstream.to_dataframe() does not preserve column order. Rearrange at
-        # the end using manually-parsed schema.
+        # Use [columns] to ensure column order matches manually-parsed schema.
         return pandas.concat(frames)[columns]
 
     def _get_progress_bar(self, progress_bar_type):
@@ -1508,7 +1561,7 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non
 
         if bqstorage_client is not None:
             try:
-                return self._to_dataframe_bqstorage(bqstorage_client, dtypes)
+                return self._to_dataframe_bqstorage(bqstorage_client, dtypes, progress_bar=progress_bar)
             except google.api_core.exceptions.Forbidden:
                 # Don't hide errors such as insufficient permissions to create
                 # a read session, or the API is not enabled. Both of those are
