feat(bigquery): share one BigQuery Storage client per Client

packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py

@@ -160,23 +160,13 @@ def download_arrow_from_job(
     if pyarrow is None:
         raise ValueError(_NO_PYARROW_ERROR)
 
-    # Step 1: Ensure BigQuery Read API client is available upfront.
-    if bqstorage_client is None:
-        if row_iterator.client is None:
-            raise ValueError("RowIterator client is None.")
-        bqstorage_client = row_iterator.client._ensure_bqstorage_client()
-        if bqstorage_client is None:
-            raise ValueError(
-                "The google-cloud-bigquery-storage library is required to read Arrow results."
-            )
-
     offset = 0
     pa_schema = None
     total_rows = row_iterator.total_rows
     job_complete = False
     remaining = row_iterator.max_results
 
-    # Step 2: Process initial inline Arrow response from jobs.query if available.
+    # Step 1: Process initial inline Arrow response from jobs.query if available.
     if row_iterator._first_page_response:
         first_page = row_iterator._first_page_response
 
@@ -206,12 +196,22 @@ def download_arrow_from_job(
                 batch, remaining = _slice_batch(batch, remaining)
                 yield batch
 
-    # Step 3: Return early if all results were delivered in the first page
+    # Step 2: Return early if all results were delivered in the first page
     # response, or if the first page already reached max_results.
     if (job_complete and offset >= total_rows) or remaining == 0:
         return
 
-    # Step 4: Stream remaining Arrow record batches from the job default stream.
+    # Step 3: Stream remaining Arrow record batches from the job default stream,
+    # with the BigQuery client's shared BigQuery Storage client unless the
+    # caller passed one. Results that fit in the first page never create it.
+    if bqstorage_client is None:
+        if row_iterator.client is None:
+            raise ValueError("RowIterator client is None.")
+        bqstorage_client = row_iterator.client._get_shared_bqstorage_client()
+        if bqstorage_client is None:
+            raise ValueError(
+                "The google-cloud-bigquery-storage library is required to read Arrow results."
+            )
     project = row_iterator._project or (
         row_iterator.client.project if row_iterator.client else None
     )

packages/google-cloud-bigquery/google/cloud/bigquery/client.py

@@ -26,6 +26,7 @@
 import math
 import os
 import tempfile
+import threading
 import typing
 import uuid
 import warnings
@@ -166,6 +167,35 @@
 # https://github.com/googleapis/python-bigquery/issues/438
 _MIN_GET_QUERY_RESULTS_TIMEOUT = 120
 
+# While a stream is open and the connection is idle, the shared BigQuery
+# Storage client pings the server after a minute, and closes the connection
+# if a ping isn't answered within a minute. A dropped ReadRows stream is
+# resumed at the next unread row.
+_BQSTORAGE_KEEPALIVE_TIME_MS = 60 * 1000
+_BQSTORAGE_KEEPALIVE_TIMEOUT_MS = 60 * 1000
+
+
+def _create_bqstorage_transport(**kwargs):
+    """Create a BigQuery Storage gRPC transport whose channel sends keepalive
+    pings.
+
+    ``BigQueryReadClient`` calls this with the arguments it passes to its
+    default transport.
+    """
+    from google.cloud import bigquery_storage
+
+    transport_class = bigquery_storage.BigQueryReadClient.get_transport_class("grpc")
+
+    def create_channel(*args, options=(), **channel_kwargs):
+        options = list(options) + [
+            ("grpc.keepalive_time_ms", _BQSTORAGE_KEEPALIVE_TIME_MS),
+            ("grpc.keepalive_timeout_ms", _BQSTORAGE_KEEPALIVE_TIMEOUT_MS),
+        ]
+        return transport_class.create_channel(*args, options=options, **channel_kwargs)
+
+    return transport_class(channel=create_channel, **kwargs)
+
+
 _LOAD_TABLE_FROM_DATAFRAME_DEPRECATED = (
     "Loading DataFrames via google-cloud-bigquery is deprecated. "
     "For direct, optimized loading, please call 'pandas_gbq.to_gbq()' directly."
@@ -300,6 +330,14 @@ def __init__(
         self._default_load_job_config = copy.deepcopy(default_load_job_config)
         self.default_job_creation_mode = default_job_creation_mode
 
+        # ARROW query results share one BigQuery Storage client. It's created
+        # the first time a query needs a stream, and close() closes it.
+        self._universe_domain = _get_client_universe(client_options)
+        self._shared_bqstorage_client: Optional[
+            "google.cloud.bigquery_storage.BigQueryReadClient"
+        ] = None
+        self._shared_bqstorage_client_lock = threading.Lock()
+
         # Use property setter so validation can run.
         self.default_query_job_config = default_query_job_config
 
@@ -348,6 +386,9 @@ def default_load_job_config(self, value: LoadJobConfig):
     def close(self):
         """Close the underlying transport objects, releasing system resources.
 
+        This also closes the BigQuery Storage client that ``ARROW`` query
+        results share.
+
         .. note::
 
             The client instance can be used for making additional requests even
@@ -356,6 +397,11 @@ def close(self):
         """
         self._http._auth_request.session.close()
         self._http.close()
+        with self._shared_bqstorage_client_lock:
+            bqstorage_client = self._shared_bqstorage_client
+            self._shared_bqstorage_client = None
+        if bqstorage_client is not None:
+            bqstorage_client._transport.close()
 
     def get_service_account_email(
         self,
@@ -600,6 +646,7 @@ def _ensure_bqstorage_client(
         client_info: Optional[
             "google.api_core.gapic_v1.client_info.ClientInfo"
         ] = DEFAULT_BQSTORAGE_CLIENT_INFO,
+        transport: Optional[Callable[..., Any]] = None,
     ) -> Optional["google.cloud.bigquery_storage.BigQueryReadClient"]:
         """Create a BigQuery Storage API client using this client's credentials.
 
@@ -613,6 +660,9 @@ def _ensure_bqstorage_client(
             client_info:
                 The client info used with a new BigQuery Storage client
                 instance if one is created.
+            transport:
+                A callable that creates the transport of a new BigQuery Storage
+                client instance. If ``None``, the default transport is used.
 
         Returns:
             A BigQuery Storage API client.
@@ -680,10 +730,37 @@ def _ensure_bqstorage_client(
                 credentials=self._credentials,
                 client_options=client_options,
                 client_info=amended_client_info,
+                transport=transport,
             )
 
         return bqstorage_client
 
+    def _get_shared_bqstorage_client(
+        self,
+    ) -> Optional["google.cloud.bigquery_storage.BigQueryReadClient"]:
+        """Return the BigQuery Storage client that ``ARROW`` query results share.
+
+        The client is created on first use, with this client's credentials and
+        universe domain and a gRPC channel that sends keepalive pings.
+        :meth:`close` closes it, and the next call creates a new one.
+
+        Returns:
+            A BigQuery Storage API client, or ``None`` if
+            google-cloud-bigquery-storage is missing or outdated.
+        """
+        with self._shared_bqstorage_client_lock:
+            if self._shared_bqstorage_client is None:
+                client_options = None
+                if self._universe_domain != _DEFAULT_UNIVERSE:
+                    client_options = google.api_core.client_options.ClientOptions(
+                        universe_domain=self._universe_domain
+                    )
+                self._shared_bqstorage_client = self._ensure_bqstorage_client(
+                    client_options=client_options,
+                    transport=_create_bqstorage_transport,
+                )
+            return self._shared_bqstorage_client
+
     def _dataset_from_arg(self, dataset) -> Union[Dataset, DatasetReference]:
         if isinstance(dataset, str):
             dataset = DatasetReference.from_string(

packages/google-cloud-bigquery/google/cloud/bigquery/table.py

@@ -2317,7 +2317,8 @@ def _download_arrow_from_job_id(
 
         Args:
             bqstorage_client (Optional[bigquery_storage.BigQueryReadClient]):
-                Client for BigQuery Storage Read API. If None, one will be created.
+                Client for BigQuery Storage Read API. If None, the BigQuery
+                Storage client that the BigQuery client shares is used.
             timeout (Optional[float]):
                 Timeout in seconds for read operations.
 

packages/google-cloud-bigquery/tests/system/test_arrow.py

@@ -755,3 +755,69 @@ def test_query_and_wait_arrow_max_results(
     table = pyarrow.Table.from_batches(batches)
     assert table.column("x").to_pylist() == list(range(1, expected_rows + 1))
     assert read_rows.called == streamed
+
+
+def _query_arrow(client, row_count):
+    return client.query_and_wait(
+        f"SELECT x FROM UNNEST(GENERATE_ARRAY(1, {row_count})) AS x ORDER BY x",
+        job_config=bigquery.QueryJobConfig(use_query_cache=False),
+        page_size=100,
+        query_results_format=enums.QueryResultsFormat.ARROW,
+    )
+
+
+def test_query_and_wait_arrow_queries_share_one_bqstorage_client():
+    from google.cloud.bigquery import client as client_module
+
+    with bigquery.Client() as client:
+        # Wrap the client factory to count how many clients get created.
+        ensure = mock.Mock(wraps=client._ensure_bqstorage_client)
+        with mock.patch.object(client, "_ensure_bqstorage_client", ensure):
+            tables = [
+                pyarrow.Table.from_batches(
+                    list(_query_arrow(client, 1000).to_arrow_iterable())
+                ),
+                _query_arrow(client, 1000).to_arrow(),
+                pyarrow.Table.from_batches(
+                    list(_query_arrow(client, 1000).to_arrow_iterable())
+                ),
+            ]
+        shared = client._shared_bqstorage_client
+
+    for table in tables:
+        assert table.column("x").to_pylist() == list(range(1, 1001))
+    ensure.assert_called_once_with(
+        client_options=None, transport=client_module._create_bqstorage_transport
+    )
+    assert shared is not None
+    # Leaving the with block calls close(), which closes the shared channel.
+    assert client._shared_bqstorage_client is None
+    with pytest.raises(ValueError, match="(?i)closed"):
+        shared.read_rows("projects/p/locations/US/jobs/j/streams/_default")
+
+
+def test_query_and_wait_arrow_ending_a_read_early_cancels_only_its_stream():
+    from google.cloud.bigquery_storage_v1 import reader
+
+    # Wrap ReadRowsStream.cancel() to count calls without changing it.
+    cancel_stream = reader.ReadRowsStream.cancel
+    with bigquery.Client() as client, mock.patch.object(
+        reader.ReadRowsStream, "cancel", autospec=True, side_effect=cancel_stream
+    ) as cancel:
+        ended_early_rows = _query_arrow(client, 100000)
+        ended_early = ended_early_rows.to_arrow_iterable()
+        read_fully = _query_arrow(client, 100000).to_arrow_iterable()
+        # Two reads reach the stream, whether or not jobs.query returned a
+        # first page.
+        next(ended_early)
+        next(ended_early)
+        rows = [next(read_fully), next(read_fully)]
+
+        ended_early.close()
+        rows.extend(read_fully)
+
+    (stream,) = cancel.call_args.args
+    cancel.assert_called_once_with(stream)
+    assert stream._name.endswith(f"/jobs/{ended_early_rows.job_id}/streams/_default")
+    table = pyarrow.Table.from_batches(rows)
+    assert table.column("x").to_pylist() == list(range(1, 100001))