Slightly simplify bits of page streaming logic

googleapis · Mar 29, 2021 · 70553bc · 70553bc
1 parent 05cd336
commit 70553bc
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 13 deletions.
diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py
@@ -703,15 +703,12 @@ def _download_table_bqstorage(
                     continue
 
             # Return any remaining values after the workers finished.
-            while not worker_queue.empty():  # pragma: NO COVER
+            while True:  # pragma: NO COVER
                 try:
-                    # Include a timeout because even though the queue is
-                    # non-empty, it doesn't guarantee that a subsequent call to
-                    # get() will not block.
-                    frame = worker_queue.get(timeout=_PROGRESS_INTERVAL)
+                    frame = worker_queue.get_nowait()
                     yield frame
                 except queue.Empty:  # pragma: NO COVER
-                    continue
+                    break
         finally:
             # No need for a lock because reading/replacing a variable is
             # defined to be an atomic operation in the Python language

diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py
@@ -1490,13 +1490,12 @@ def _to_page_iterable(
         if not self._validate_bqstorage(bqstorage_client, False):
             bqstorage_client = None
 
-        if bqstorage_client is not None:
-            for item in bqstorage_download():
-                yield item
-            return
-
-        for item in tabledata_list_download():
-            yield item
+        result_pages = (
+            bqstorage_download()
+            if bqstorage_client is not None
+            else tabledata_list_download()
+        )
+        yield from result_pages
 
     def _to_arrow_iterable(
         self,