ENH: Add use_bqstorage_api option to read_gbq

The BigQuery Storage API provides a way to read query results quickly
(and using multiple threads). It only works with large query results
(~125 MB), but as of 1.11.1, the google-cloud-bigquery library can
fallback to the BigQuery API to download results when a request to the
BigQuery Storage API fails.

As this API can increase costs (and may not be enabled on the user's
project), this option is disabled by default.

Author: tswast

pandas_gbq/gbq.py

@@ -5,6 +5,13 @@
 
 import numpy as np
 
+try:
+    # The BigQuery Storage API client is an optional dependency. It is only
+    # required when use_bqstorage_api=True.
+    from google.cloud import bigquery_storage_v1beta1
+except ImportError:  # pragma: NO COVER
+    bigquery_storage_v1beta1 = None
+
 from pandas_gbq.exceptions import AccessDenied
 
 logger = logging.getLogger(__name__)
@@ -302,6 +309,7 @@ def __init__(
         dialect="standard",
         location=None,
         credentials=None,
+        use_bqstorage_api=False,
     ):
         global context
         from google.api_core.exceptions import GoogleAPIError
@@ -352,6 +360,9 @@ def __init__(
             context.project = self.project_id
 
         self.client = self.get_client()
+        self.bqstorage_client = _make_bqstorage_client(
+            use_bqstorage_api, self.credentials
+        )
 
         # BQ Queries costs $5 per TB. First 1 TB per month is free
         # see here for more: https://cloud.google.com/bigquery/pricing
@@ -489,7 +500,9 @@ def run_query(self, query, **kwargs):
 
         schema_fields = [field.to_api_repr() for field in rows_iter.schema]
         nullsafe_dtypes = _bqschema_to_nullsafe_dtypes(schema_fields)
-        df = rows_iter.to_dataframe(dtypes=nullsafe_dtypes)
+        df = rows_iter.to_dataframe(
+            dtypes=nullsafe_dtypes, bqstorage_client=self.bqstorage_client
+        )
 
         if df.empty:
             df = _cast_empty_df_dtypes(schema_fields, df)
@@ -727,6 +740,21 @@ def _localize_df(schema_fields, df):
     return df
 
 
+def _make_bqstorage_client(use_bqstorage_api, credentials):
+    if not use_bqstorage_api:
+        return None
+
+    if bigquery_storage_v1beta1 is None:
+        raise ImportError(
+            "Install the google-cloud-bigquery-storage and fastavro packages "
+            "to use the BigQuery Storage API."
+        )
+
+    return bigquery_storage_v1beta1.BigQueryStorageClient(
+        credentials=credentials
+    )
+
+
 def read_gbq(
     query,
     project_id=None,
@@ -738,6 +766,7 @@ def read_gbq(
     location=None,
     configuration=None,
     credentials=None,
+    use_bqstorage_api=False,
     verbose=None,
     private_key=None,
 ):
@@ -815,6 +844,27 @@ def read_gbq(
         :class:`google.oauth2.service_account.Credentials` directly.
 
         .. versionadded:: 0.8.0
+    use_bqstorage_api : bool, default False
+        Use the `BigQuery Storage API
+        <https://cloud.google.com/bigquery/docs/reference/storage/>`__ to
+        download query results quickly, but at an increased cost. To use this
+        API, first `enable it in the Cloud Console
+        <https://console.cloud.google.com/apis/library/bigquerystorage.googleapis.com>`__.
+        You must also have the `bigquery.readsessions.create
+        <https://cloud.google.com/bigquery/docs/access-control#roles>`__
+        permission on the project you are billing queries to.
+
+        **Note:** Due to a `known issue in the ``google-cloud-bigquery``
+        package
+        <https://github.com/googleapis/google-cloud-python/pull/7633>`__
+        (fixed in version 1.11.0), you must write your query results to a
+        destination table. To do this with ``read_gbq``, supply a
+        ``configuration`` dictionary.
+
+        This feature requires the ``google-cloud-bigquery-storage`` and
+        ``fastavro`` packages.
+
+        .. versionadded:: 0.10.0
     verbose : None, deprecated
         Deprecated in Pandas-GBQ 0.4.0. Use the `logging module
         to adjust verbosity instead
@@ -835,6 +885,27 @@ def read_gbq(
     -------
     df: DataFrame
         DataFrame representing results of query.
+
+    Examples
+    --------
+
+    Use the BigQuery Storage API to fetch results quickly, but at an addition
+    cost. Due to a known issue in the BigQuery Storage API, you must write
+    your query results to a destination table.
+
+    >>> pandas_gbq.read_gbq(
+    ...     query_string,
+    ...     configuration={
+    ...         'query': {
+    ...             'destinationTable': {
+    ...                 'projectId': 'your-project',
+    ...                 'datasetId': 'destination_dataset',
+    ...                 'tableId': 'new_table_name',
+    ...             }
+    ...         }
+    ...     },
+    ...     use_bqstorage_api=True,
+    ... )
     """
     global context
 
@@ -871,6 +942,7 @@ def read_gbq(
         location=location,
         credentials=credentials,
         private_key=private_key,
+        use_bqstorage_api=use_bqstorage_api,
     )
 
     final_df = connector.run_query(query, configuration=configuration)

tests/system/test_gbq.py

@@ -895,10 +895,36 @@ def test_tokyo(self, tokyo_dataset, tokyo_table, private_key_path):
             location="asia-northeast1",
             private_key=private_key_path,
         )
-        print(df)
         assert df["max_year"][0] >= 2000
 
 
+@pytest.mark.skip(reason="large query for BQ Storage API tests")
+def test_read_gbq_w_bqstorage_api(credentials):
+    df = gbq.read_gbq(
+        """
+        SELECT
+            dependency_name,
+            dependency_platform,
+            project_name,
+            project_id,
+            version_number,
+            version_id,
+            dependency_kind,
+            optional_dependency,
+            dependency_requirements,
+            dependency_project_id
+        FROM
+            `bigquery-public-data.libraries_io.dependencies`
+        WHERE
+            LOWER(dependency_platform) = 'npm'
+        LIMIT 2500000
+        """,
+        use_bqstorage_api=True,
+        credentials=credentials,
+    )
+    assert len(df) == 2500000
+
+
 class TestToGBQIntegration(object):
     @pytest.fixture(autouse=True, scope="function")
     def setup(self, project, credentials, random_dataset_id):