Merge branch 'main' into b323176126-pandas-gbq

Author: tswast

.kokoro/requirements.txt

@@ -112,34 +112,38 @@ colorlog==6.8.2 \
     # via
     #   gcp-docuploader
     #   nox
-cryptography==43.0.1 \
-    --hash=sha256:014f58110f53237ace6a408b5beb6c427b64e084eb451ef25a28308270086494 \
-    --hash=sha256:1bbcce1a551e262dfbafb6e6252f1ae36a248e615ca44ba302df077a846a8806 \
-    --hash=sha256:203e92a75716d8cfb491dc47c79e17d0d9207ccffcbcb35f598fbe463ae3444d \
-    --hash=sha256:27e613d7077ac613e399270253259d9d53872aaf657471473ebfc9a52935c062 \
-    --hash=sha256:2bd51274dcd59f09dd952afb696bf9c61a7a49dfc764c04dd33ef7a6b502a1e2 \
-    --hash=sha256:38926c50cff6f533f8a2dae3d7f19541432610d114a70808f0926d5aaa7121e4 \
-    --hash=sha256:511f4273808ab590912a93ddb4e3914dfd8a388fed883361b02dea3791f292e1 \
-    --hash=sha256:58d4e9129985185a06d849aa6df265bdd5a74ca6e1b736a77959b498e0505b85 \
-    --hash=sha256:5b43d1ea6b378b54a1dc99dd8a2b5be47658fe9a7ce0a58ff0b55f4b43ef2b84 \
-    --hash=sha256:61ec41068b7b74268fa86e3e9e12b9f0c21fcf65434571dbb13d954bceb08042 \
-    --hash=sha256:666ae11966643886c2987b3b721899d250855718d6d9ce41b521252a17985f4d \
-    --hash=sha256:68aaecc4178e90719e95298515979814bda0cbada1256a4485414860bd7ab962 \
-    --hash=sha256:7c05650fe8023c5ed0d46793d4b7d7e6cd9c04e68eabe5b0aeea836e37bdcec2 \
-    --hash=sha256:80eda8b3e173f0f247f711eef62be51b599b5d425c429b5d4ca6a05e9e856baa \
-    --hash=sha256:8385d98f6a3bf8bb2d65a73e17ed87a3ba84f6991c155691c51112075f9ffc5d \
-    --hash=sha256:88cce104c36870d70c49c7c8fd22885875d950d9ee6ab54df2745f83ba0dc365 \
-    --hash=sha256:9d3cdb25fa98afdd3d0892d132b8d7139e2c087da1712041f6b762e4f807cc96 \
-    --hash=sha256:a575913fb06e05e6b4b814d7f7468c2c660e8bb16d8d5a1faf9b33ccc569dd47 \
-    --hash=sha256:ac119bb76b9faa00f48128b7f5679e1d8d437365c5d26f1c2c3f0da4ce1b553d \
-    --hash=sha256:c1332724be35d23a854994ff0b66530119500b6053d0bd3363265f7e5e77288d \
-    --hash=sha256:d03a475165f3134f773d1388aeb19c2d25ba88b6a9733c5c590b9ff7bbfa2e0c \
-    --hash=sha256:d75601ad10b059ec832e78823b348bfa1a59f6b8d545db3a24fd44362a1564cb \
-    --hash=sha256:de41fd81a41e53267cb020bb3a7212861da53a7d39f863585d13ea11049cf277 \
-    --hash=sha256:e710bf40870f4db63c3d7d929aa9e09e4e7ee219e703f949ec4073b4294f6172 \
-    --hash=sha256:ea25acb556320250756e53f9e20a4177515f012c9eaea17eb7587a8c4d8ae034 \
-    --hash=sha256:f98bf604c82c416bc829e490c700ca1553eafdf2912a91e23a79d97d9801372a \
-    --hash=sha256:fba1007b3ef89946dbbb515aeeb41e30203b004f0b4b00e5e16078b518563289
+cryptography==44.0.1 \
+    --hash=sha256:00918d859aa4e57db8299607086f793fa7813ae2ff5a4637e318a25ef82730f7 \
+    --hash=sha256:1e8d181e90a777b63f3f0caa836844a1182f1f265687fac2115fcf245f5fbec3 \
+    --hash=sha256:1f9a92144fa0c877117e9748c74501bea842f93d21ee00b0cf922846d9d0b183 \
+    --hash=sha256:21377472ca4ada2906bc313168c9dc7b1d7ca417b63c1c3011d0c74b7de9ae69 \
+    --hash=sha256:24979e9f2040c953a94bf3c6782e67795a4c260734e5264dceea65c8f4bae64a \
+    --hash=sha256:2a46a89ad3e6176223b632056f321bc7de36b9f9b93b2cc1cccf935a3849dc62 \
+    --hash=sha256:322eb03ecc62784536bc173f1483e76747aafeb69c8728df48537eb431cd1911 \
+    --hash=sha256:436df4f203482f41aad60ed1813811ac4ab102765ecae7a2bbb1dbb66dcff5a7 \
+    --hash=sha256:4f422e8c6a28cf8b7f883eb790695d6d45b0c385a2583073f3cec434cc705e1a \
+    --hash=sha256:53f23339864b617a3dfc2b0ac8d5c432625c80014c25caac9082314e9de56f41 \
+    --hash=sha256:5fed5cd6102bb4eb843e3315d2bf25fede494509bddadb81e03a859c1bc17b83 \
+    --hash=sha256:610a83540765a8d8ce0f351ce42e26e53e1f774a6efb71eb1b41eb01d01c3d12 \
+    --hash=sha256:6c8acf6f3d1f47acb2248ec3ea261171a671f3d9428e34ad0357148d492c7864 \
+    --hash=sha256:6f76fdd6fd048576a04c5210d53aa04ca34d2ed63336d4abd306d0cbe298fddf \
+    --hash=sha256:72198e2b5925155497a5a3e8c216c7fb3e64c16ccee11f0e7da272fa93b35c4c \
+    --hash=sha256:887143b9ff6bad2b7570da75a7fe8bbf5f65276365ac259a5d2d5147a73775f2 \
+    --hash=sha256:888fcc3fce0c888785a4876ca55f9f43787f4c5c1cc1e2e0da71ad481ff82c5b \
+    --hash=sha256:8e6a85a93d0642bd774460a86513c5d9d80b5c002ca9693e63f6e540f1815ed0 \
+    --hash=sha256:94f99f2b943b354a5b6307d7e8d19f5c423a794462bde2bf310c770ba052b1c4 \
+    --hash=sha256:9b336599e2cb77b1008cb2ac264b290803ec5e8e89d618a5e978ff5eb6f715d9 \
+    --hash=sha256:a2d8a7045e1ab9b9f803f0d9531ead85f90c5f2859e653b61497228b18452008 \
+    --hash=sha256:b8272f257cf1cbd3f2e120f14c68bff2b6bdfcc157fafdee84a1b795efd72862 \
+    --hash=sha256:bf688f615c29bfe9dfc44312ca470989279f0e94bb9f631f85e3459af8efc009 \
+    --hash=sha256:d9c5b9f698a83c8bd71e0f4d3f9f839ef244798e5ffe96febfa9714717db7af7 \
+    --hash=sha256:dd7c7e2d71d908dc0f8d2027e1604102140d84b155e658c20e8ad1304317691f \
+    --hash=sha256:df978682c1504fc93b3209de21aeabf2375cb1571d4e61907b3e7a2540e83026 \
+    --hash=sha256:e403f7f766ded778ecdb790da786b418a9f2394f36e8cc8b796cc056ab05f44f \
+    --hash=sha256:eb3889330f2a4a148abead555399ec9a32b13b7c8ba969b72d8e500eb7ef84cd \
+    --hash=sha256:f4daefc971c2d1f82f03097dc6f216744a6cd2ac0f04c68fb935ea2ba2a0d420 \
+    --hash=sha256:f51f5705ab27898afda1aaa430f34ad90dc117421057782022edf0600bec5f14 \
+    --hash=sha256:fd0ee90072861e276b0ff08bd627abec29e32a53b2be44e41dbcdf87cbee2b00
     # via
     #   -r requirements.in
     #   gcp-releasetool

docs/magics.rst

@@ -6,7 +6,7 @@ in a Jupyter notebook cell.
 
 .. code::
 
-    %load_ext google.cloud.bigquery
+    %load_ext bigquery_magics
 
 This makes the ``%%bigquery`` magic available.
 
@@ -27,8 +27,9 @@ Running a parameterized query:
    :start-after: [START bigquery_jupyter_query_params_scalars]
    :end-before: [END bigquery_jupyter_query_params_scalars]
 
-API Reference
--------------
+BigQuery Magics Reference
+-------------------------
 
-.. automodule:: google.cloud.bigquery.magics.magics
-    :members:
+-  `BigQuery Magics Documentation`_
+
+.. _BigQuery Magics Documentation: https://googleapis.dev/python/bigquery-magics/latest

google/cloud/bigquery/_pandas_helpers.py

@@ -25,6 +25,7 @@
 from itertools import islice
 import logging
 import queue
+import threading
 import warnings
 from typing import Any, Union, Optional, Callable, Generator, List
 
@@ -134,6 +135,21 @@ def __init__(self):
         # be an atomic operation in the Python language definition (enforced by
         # the global interpreter lock).
         self.done = False
+        # To assist with testing and understanding the behavior of the
+        # download, use this object as shared state to track how many worker
+        # threads have started and have gracefully shutdown.
+        self._started_workers_lock = threading.Lock()
+        self.started_workers = 0
+        self._finished_workers_lock = threading.Lock()
+        self.finished_workers = 0
+
+    def start(self):
+        with self._started_workers_lock:
+            self.started_workers += 1
+
+    def finish(self):
+        with self._finished_workers_lock:
+            self.finished_workers += 1
 
 
 BQ_FIELD_TYPE_TO_ARROW_FIELD_METADATA = {
@@ -819,25 +835,35 @@ def _bqstorage_page_to_dataframe(column_names, dtypes, page):
 def _download_table_bqstorage_stream(
     download_state, bqstorage_client, session, stream, worker_queue, page_to_item
 ):
-    reader = bqstorage_client.read_rows(stream.name)
+    download_state.start()
+    try:
+        reader = bqstorage_client.read_rows(stream.name)
 
-    # Avoid deprecation warnings for passing in unnecessary read session.
-    # https://github.com/googleapis/python-bigquery-storage/issues/229
-    if _versions_helpers.BQ_STORAGE_VERSIONS.is_read_session_optional:
-        rowstream = reader.rows()
-    else:
-        rowstream = reader.rows(session)
-
-    for page in rowstream.pages:
-        item = page_to_item(page)
-        while True:
-            if download_state.done:
-                return
-            try:
-                worker_queue.put(item, timeout=_PROGRESS_INTERVAL)
-                break
-            except queue.Full:  # pragma: NO COVER
-                continue
+        # Avoid deprecation warnings for passing in unnecessary read session.
+        # https://github.com/googleapis/python-bigquery-storage/issues/229
+        if _versions_helpers.BQ_STORAGE_VERSIONS.is_read_session_optional:
+            rowstream = reader.rows()
+        else:
+            rowstream = reader.rows(session)
+
+        for page in rowstream.pages:
+            item = page_to_item(page)
+
+            # Make sure we set a timeout on put() so that we give the worker
+            # thread opportunities to shutdown gracefully, for example if the
+            # parent thread shuts down or the parent generator object which
+            # collects rows from all workers goes out of scope. See:
+            # https://github.com/googleapis/python-bigquery/issues/2032
+            while True:
+                if download_state.done:
+                    return
+                try:
+                    worker_queue.put(item, timeout=_PROGRESS_INTERVAL)
+                    break
+                except queue.Full:
+                    continue
+    finally:
+        download_state.finish()
 
 
 def _nowait(futures):
@@ -863,6 +889,7 @@ def _download_table_bqstorage(
     page_to_item: Optional[Callable] = None,
     max_queue_size: Any = _MAX_QUEUE_SIZE_DEFAULT,
     max_stream_count: Optional[int] = None,
+    download_state: Optional[_DownloadState] = None,
 ) -> Generator[Any, None, None]:
     """Downloads a BigQuery table using the BigQuery Storage API.
 
@@ -890,6 +917,9 @@ def _download_table_bqstorage(
             is True, the requested streams are limited to 1 regardless of the
             `max_stream_count` value. If 0 or None, then the number of
             requested streams will be unbounded. Defaults to None.
+        download_state (Optional[_DownloadState]):
+            A threadsafe state object which can be used to observe the
+            behavior of the worker threads created by this method.
 
     Yields:
         pandas.DataFrame: Pandas DataFrames, one for each chunk of data
@@ -948,7 +978,8 @@ def _download_table_bqstorage(
 
     # Use _DownloadState to notify worker threads when to quit.
     # See: https://stackoverflow.com/a/29237343/101923
-    download_state = _DownloadState()
+    if download_state is None:
+        download_state = _DownloadState()
 
     # Create a queue to collect frames as they are created in each thread.
     #

samples/tests/test_download_public_data.py

@@ -12,29 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
-
 import pytest
 
 from .. import download_public_data
 
 pytest.importorskip("google.cloud.bigquery_storage_v1")
 
 
-def test_download_public_data(
-    caplog: pytest.LogCaptureFixture, capsys: pytest.CaptureFixture[str]
-) -> None:
-    # Enable debug-level logging to verify the BigQuery Storage API is used.
-    caplog.set_level(logging.DEBUG)
-
+def test_download_public_data(capsys: pytest.CaptureFixture[str]) -> None:
     download_public_data.download_public_data()
     out, _ = capsys.readouterr()
     assert "year" in out
     assert "gender" in out
     assert "name" in out
-
-    assert any(
-        "Started reading table 'bigquery-public-data.usa_names.usa_1910_current' with BQ Storage API session"
-        in message
-        for message in caplog.messages
-    )

samples/tests/test_download_public_data_sandbox.py

@@ -12,29 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
-
 import pytest
 
 from .. import download_public_data_sandbox
 
 pytest.importorskip("google.cloud.bigquery_storage_v1")
 
 
-def test_download_public_data_sandbox(
-    caplog: pytest.LogCaptureFixture, capsys: pytest.CaptureFixture[str]
-) -> None:
-    # Enable debug-level logging to verify the BigQuery Storage API is used.
-    caplog.set_level(logging.DEBUG)
-
+def test_download_public_data_sandbox(capsys: pytest.CaptureFixture[str]) -> None:
     download_public_data_sandbox.download_public_data_sandbox()
-    out, err = capsys.readouterr()
+    out, _ = capsys.readouterr()
     assert "year" in out
     assert "gender" in out
     assert "name" in out
-
-    assert any(
-        # An anonymous table is used because this sample reads from query results.
-        ("Started reading table" in message and "BQ Storage API session" in message)
-        for message in caplog.messages
-    )

tests/unit/test__pandas_helpers.py

@@ -16,6 +16,7 @@
 import datetime
 import decimal
 import functools
+import gc
 import operator
 import queue
 from typing import Union
@@ -1889,6 +1890,98 @@ def fake_download_stream(
     assert queue_used.maxsize == expected_maxsize
 
 
+@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
+def test__download_table_bqstorage_shuts_down_workers(
+    monkeypatch,
+    module_under_test,
+):
+    """Regression test for https://github.com/googleapis/python-bigquery/issues/2032
+
+    Make sure that when the top-level iterator goes out of scope (is deleted),
+    the child threads are also stopped.
+    """
+    from google.cloud.bigquery import dataset
+    from google.cloud.bigquery import table
+    import google.cloud.bigquery_storage_v1.reader
+    import google.cloud.bigquery_storage_v1.types
+
+    monkeypatch.setattr(
+        _versions_helpers.BQ_STORAGE_VERSIONS, "_installed_version", None
+    )
+    monkeypatch.setattr(bigquery_storage, "__version__", "2.5.0")
+
+    # Create a fake stream with a decent number of rows.
+    arrow_schema = pyarrow.schema(
+        [
+            ("int_col", pyarrow.int64()),
+            ("str_col", pyarrow.string()),
+        ]
+    )
+    arrow_rows = pyarrow.record_batch(
+        [
+            pyarrow.array([0, 1, 2], type=pyarrow.int64()),
+            pyarrow.array(["a", "b", "c"], type=pyarrow.string()),
+        ],
+        schema=arrow_schema,
+    )
+    session = google.cloud.bigquery_storage_v1.types.ReadSession()
+    session.data_format = "ARROW"
+    session.arrow_schema = {"serialized_schema": arrow_schema.serialize().to_pybytes()}
+    session.streams = [
+        google.cloud.bigquery_storage_v1.types.ReadStream(name=name)
+        for name in ("stream/s0", "stream/s1", "stream/s2")
+    ]
+    bqstorage_client = mock.create_autospec(
+        bigquery_storage.BigQueryReadClient, instance=True
+    )
+    reader = mock.create_autospec(
+        google.cloud.bigquery_storage_v1.reader.ReadRowsStream, instance=True
+    )
+    reader.__iter__.return_value = [
+        google.cloud.bigquery_storage_v1.types.ReadRowsResponse(
+            arrow_schema={"serialized_schema": arrow_schema.serialize().to_pybytes()},
+            arrow_record_batch={
+                "serialized_record_batch": arrow_rows.serialize().to_pybytes()
+            },
+        )
+        for _ in range(100)
+    ]
+    reader.rows.return_value = google.cloud.bigquery_storage_v1.reader.ReadRowsIterable(
+        reader, read_session=session
+    )
+    bqstorage_client.read_rows.return_value = reader
+    bqstorage_client.create_read_session.return_value = session
+    table_ref = table.TableReference(
+        dataset.DatasetReference("project-x", "dataset-y"),
+        "table-z",
+    )
+    download_state = module_under_test._DownloadState()
+    assert download_state.started_workers == 0
+    assert download_state.finished_workers == 0
+
+    result_gen = module_under_test._download_table_bqstorage(
+        "some-project",
+        table_ref,
+        bqstorage_client,
+        max_queue_size=1,
+        page_to_item=module_under_test._bqstorage_page_to_arrow,
+        download_state=download_state,
+    )
+
+    result_gen_iter = iter(result_gen)
+    next(result_gen_iter)
+    assert download_state.started_workers == 3
+    assert download_state.finished_workers == 0
+
+    # Stop iteration early and simulate the variables going out of scope
+    # to be doubly sure that the worker threads are supposed to be cleaned up.
+    del result_gen, result_gen_iter
+    gc.collect()
+
+    assert download_state.started_workers == 3
+    assert download_state.finished_workers == 3
+
+
 @pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
 def test_download_arrow_row_iterator_unknown_field_type(module_under_test):
     fake_page = api_core.page_iterator.Page(