Adjust tests, add support for infinite queue size

tswast · Apr 14, 2021 · Mar 24, 2021 · Mar 24, 2021 · Mar 31, 2021 · Mar 31, 2021
commit 85929ed32599ceb586d189a84faf545affc9f63f
diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py
@@ -45,7 +45,7 @@
 
 _PROGRESS_INTERVAL = 0.2  # Maximum time between download status checks, in seconds.
 
-_DEFAULT_MAX_QUEUE_SIZE = 1  # max queue size for BQ Storage downloads
+_MAX_QUEUE_SIZE_DEFAULT = object()  # max queue size sentinel for BQ Storage downloads
 
 _PANDAS_DTYPE_TO_BQ = {
     "bool": "BOOLEAN",
@@ -610,7 +610,7 @@ def _download_table_bqstorage(
     preserve_order=False,
     selected_fields=None,
     page_to_item=None,
-    max_queue_size=_DEFAULT_MAX_QUEUE_SIZE,
+    max_queue_size=_MAX_QUEUE_SIZE_DEFAULT,
 ):
     """Use (faster, but billable) BQ Storage API to construct DataFrame."""
 
@@ -658,10 +658,15 @@ def _download_table_bqstorage(
 
     # Create a queue to collect frames as they are created in each thread.
     #
-    # The queue needs to be bounded, because if the user code processes the fetched
-    # result pages to slowly, but at the same time new pages are rapidly being fetched
-    # from the server, the queue can grow to the point where the process runs
+    # The queue needs to be bounded by default, because if the user code processes the
+    # fetched result pages too slowly, while at the same time new pages are rapidly being
+    # fetched from the server, the queue can grow to the point where the process runs
     # out of memory.
+    if max_queue_size is _MAX_QUEUE_SIZE_DEFAULT:
+        max_queue_size = total_streams
+    elif max_queue_size is None:
+        max_queue_size = 0  # unbounded
+
     worker_queue = queue.Queue(maxsize=max_queue_size)
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=total_streams) as pool:
@@ -741,7 +746,7 @@ def download_dataframe_bqstorage(
     dtypes,
     preserve_order=False,
     selected_fields=None,
-    max_queue_size=_DEFAULT_MAX_QUEUE_SIZE,
+    max_queue_size=_MAX_QUEUE_SIZE_DEFAULT,
 ):
     page_to_item = functools.partial(_bqstorage_page_to_dataframe, column_names, dtypes)
     return _download_table_bqstorage(

diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py
@@ -1625,7 +1625,7 @@ def to_dataframe_iterable(
         self,
         bqstorage_client=None,
         dtypes=None,
-        max_queue_size=_pandas_helpers._DEFAULT_MAX_QUEUE_SIZE,
+        max_queue_size=_pandas_helpers._MAX_QUEUE_SIZE_DEFAULT,
     ):
         """Create an iterable of pandas DataFrames, to process the table as a stream.
 
@@ -1651,6 +1651,10 @@ def to_dataframe_iterable(
                 streaming query results over the BigQuery Storage API. Ignored if
                 Storage API is not used.
 
+                By default, the max queue size is set to the number of BQ Storage streams
+                created by the server. If ``max_queue_size`` is :data:`None`, the queue
+                size is infinite.
+
                 ..versionadded:: 2.14.0
 
         Returns:

diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py
@@ -42,6 +42,11 @@
 from google.cloud.bigquery import schema
 from google.cloud.bigquery._pandas_helpers import _BIGNUMERIC_SUPPORT
 
+try:
+    from google.cloud import bigquery_storage
+except ImportError:  # pragma: NO COVER
+    bigquery_storage = None
+
 
 skip_if_no_bignumeric = pytest.mark.skipif(
     not _BIGNUMERIC_SUPPORT, reason="BIGNUMERIC support requires pyarrow>=3.0.0",
@@ -1266,43 +1271,64 @@ def test_dataframe_to_parquet_dict_sequence_schema(module_under_test):
     assert schema_arg == expected_schema_arg
 
 
-def test__download_table_bqstorage_bounded_queue_size(module_under_test):
+@pytest.mark.parametrize(
+    "stream_count,maxsize_kwarg,expected_call_count,expected_maxsize",
+    [
+        (3, {"max_queue_size": 2}, 3, 2),  # custom queue size
+        (4, {}, 4, 4),  # default queue size
+        (7, {"max_queue_size": None}, 7, 0),  # infinite queue size
+    ],
+)
+@pytest.mark.skipif(
+    bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`"
+)
+def test__download_table_bqstorage(
+    module_under_test,
+    stream_count,
+    maxsize_kwarg,
+    expected_call_count,
+    expected_maxsize,
+):
     from google.cloud.bigquery import dataset
     from google.cloud.bigquery import table
 
-    table_ref = table.TableReference(
-        dataset.DatasetReference("project-x", "dataset-y"), "table-z",
-    )
+    queue_used = None  # A reference to the queue used by code under test.
 
-    bigquery_storage = pytest.importorskip(
-        "google.cloud.bigquery_storage",
-        reason="Requires BigQuery Storage dependency."
-    )
     bqstorage_client = mock.create_autospec(
         bigquery_storage.BigQueryReadClient, instance=True
     )
-    fake_session = mock.Mock(streams=["stream/s1", "stream/s2", "stream/s3"])
+    fake_session = mock.Mock(streams=["stream/s{i}" for i in range(stream_count)])
     bqstorage_client.create_read_session.return_value = fake_session
 
+    table_ref = table.TableReference(
+        dataset.DatasetReference("project-x", "dataset-y"), "table-z",
+    )
+
     def fake_download_stream(
         download_state, bqstorage_client, session, stream, worker_queue, page_to_item
     ):
-        worker_queue.put_nowait("result")
+        nonlocal queue_used
+        queue_used = worker_queue
+        try:
+            worker_queue.put_nowait("result_page")
+        except queue.Full:  # pragma: NO COVER
+            pass
 
     download_stream = mock.Mock(side_effect=fake_download_stream)
 
-    with pytest.raises(queue.Full), mock.patch.object(
+    with mock.patch.object(
         module_under_test, "_download_table_bqstorage_stream", new=download_stream
     ):
         result_gen = module_under_test._download_table_bqstorage(
-            "some-project", table_ref, bqstorage_client, max_queue_size=2,
+            "some-project", table_ref, bqstorage_client, **maxsize_kwarg
         )
         list(result_gen)
 
     # Timing-safe, as the method under test should block until the pool shutdown is
     # complete, at which point all download stream workers have already been submitted
     # to the thread pool.
-    assert download_stream.call_count == 3
+    assert download_stream.call_count == stream_count  # once for each stream
+    assert queue_used.maxsize == expected_maxsize
 
 
 @pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")