From 680f952192a7c7a47cc9658e20ad5434f7a9d370 Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Wed, 24 Mar 2021 17:34:58 +0100 Subject: [PATCH] Slightly simplify bits of page streaming logic --- google/cloud/bigquery/_pandas_helpers.py | 9 +++------ google/cloud/bigquery/table.py | 13 ++++++------- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py index 9cefa0a6c..184632973 100644 --- 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 index 64c7ee8a4..fa5b11f37 100644 --- 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,