feat(bigquery): add create_bqstorage_client param to to_dataframe and to_arrow (#9573)

* feat(bigquery): add `create_bqstorage_client` param to `to_dataframe` and `to_arrow`

When the `create_bqstorage_client` parameter is set to `True`, the
BigQuery client constructs a BigQuery Storage API client for you. This
removes the need for boilerplate code to manually construct both clients
explitly with the same credentials.

Does this make the `bqstorage_client` parameter unnecessary? In most
cases, yes, but there are a few cases where we'll want to continue using
it.

* When partner tools use `to_dataframe`, they should continue to use
  `bqstorage_client` so that they can set the correct amended user-agent
  strings.
* When a developer needs to override the default API endpoint for the BQ
  Storage API, they'll need to manually supply a `bqstorage_client`.

* test for BQ Storage API usage in samples tests.

* fix: close bqstorage client if created by to_dataframe/to_arrow

* chore: blacken

* doc: update versionadded

* doc: update versionadded

Author: tswast

bigquery/google/cloud/bigquery/client.py

@@ -353,6 +353,19 @@ def dataset(self, dataset_id, project=None):
 
         return DatasetReference(project, dataset_id)
 
+    def _create_bqstorage_client(self):
+        """Create a BigQuery Storage API client using this client's credentials.
+
+        Returns:
+            google.cloud.bigquery_storage_v1beta1.BigQueryStorageClient:
+                A BigQuery Storage API client.
+        """
+        from google.cloud import bigquery_storage_v1beta1
+
+        return bigquery_storage_v1beta1.BigQueryStorageClient(
+            credentials=self._credentials
+        )
+
     def create_dataset(self, dataset, exists_ok=False, retry=DEFAULT_RETRY):
         """API call: create the dataset via a POST request.
 

bigquery/google/cloud/bigquery/job.py

@@ -3152,7 +3152,12 @@ def result(
 
     # If changing the signature of this method, make sure to apply the same
     # changes to table.RowIterator.to_arrow()
-    def to_arrow(self, progress_bar_type=None, bqstorage_client=None):
+    def to_arrow(
+        self,
+        progress_bar_type=None,
+        bqstorage_client=None,
+        create_bqstorage_client=False,
+    ):
         """[Beta] Create a class:`pyarrow.Table` by loading all pages of a
         table or query.
 
@@ -3185,6 +3190,16 @@ def to_arrow(self, progress_bar_type=None, bqstorage_client=None):
 
                 Reading from a specific partition or snapshot is not
                 currently supported by this method.
+            create_bqstorage_client (bool):
+                **Beta Feature** Optional. If ``True``, create a BigQuery
+                Storage API client using the default API settings. The
+                BigQuery Storage API is a faster way to fetch rows from
+                BigQuery. See the ``bqstorage_client`` parameter for more
+                information.
+
+                This argument does nothing if ``bqstorage_client`` is supplied.
+
+                ..versionadded:: 1.24.0
 
         Returns:
             pyarrow.Table
@@ -3199,12 +3214,20 @@ def to_arrow(self, progress_bar_type=None, bqstorage_client=None):
         ..versionadded:: 1.17.0
         """
         return self.result().to_arrow(
-            progress_bar_type=progress_bar_type, bqstorage_client=bqstorage_client
+            progress_bar_type=progress_bar_type,
+            bqstorage_client=bqstorage_client,
+            create_bqstorage_client=create_bqstorage_client,
         )
 
     # If changing the signature of this method, make sure to apply the same
     # changes to table.RowIterator.to_dataframe()
-    def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=None):
+    def to_dataframe(
+        self,
+        bqstorage_client=None,
+        dtypes=None,
+        progress_bar_type=None,
+        create_bqstorage_client=False,
+    ):
         """Return a pandas DataFrame from a QueryJob
 
         Args:
@@ -3237,6 +3260,16 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non
                 for details.
 
                 ..versionadded:: 1.11.0
+            create_bqstorage_client (bool):
+                **Beta Feature** Optional. If ``True``, create a BigQuery
+                Storage API client using the default API settings. The
+                BigQuery Storage API is a faster way to fetch rows from
+                BigQuery. See the ``bqstorage_client`` parameter for more
+                information.
+
+                This argument does nothing if ``bqstorage_client`` is supplied.
+
+                ..versionadded:: 1.24.0
 
         Returns:
             A :class:`~pandas.DataFrame` populated with row data and column
@@ -3250,6 +3283,7 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non
             bqstorage_client=bqstorage_client,
             dtypes=dtypes,
             progress_bar_type=progress_bar_type,
+            create_bqstorage_client=create_bqstorage_client,
         )
 
     def __iter__(self):

bigquery/google/cloud/bigquery/table.py

@@ -1456,7 +1456,12 @@ def _to_arrow_iterable(self, bqstorage_client=None):
 
     # If changing the signature of this method, make sure to apply the same
     # changes to job.QueryJob.to_arrow()
-    def to_arrow(self, progress_bar_type=None, bqstorage_client=None):
+    def to_arrow(
+        self,
+        progress_bar_type=None,
+        bqstorage_client=None,
+        create_bqstorage_client=False,
+    ):
         """[Beta] Create a class:`pyarrow.Table` by loading all pages of a
         table or query.
 
@@ -1489,6 +1494,16 @@ def to_arrow(self, progress_bar_type=None, bqstorage_client=None):
 
                 Reading from a specific partition or snapshot is not
                 currently supported by this method.
+            create_bqstorage_client (bool):
+                **Beta Feature** Optional. If ``True``, create a BigQuery
+                Storage API client using the default API settings. The
+                BigQuery Storage API is a faster way to fetch rows from
+                BigQuery. See the ``bqstorage_client`` parameter for more
+                information.
+
+                This argument does nothing if ``bqstorage_client`` is supplied.
+
+                ..versionadded:: 1.24.0
 
         Returns:
             pyarrow.Table
@@ -1504,22 +1519,33 @@ def to_arrow(self, progress_bar_type=None, bqstorage_client=None):
         if pyarrow is None:
             raise ValueError(_NO_PYARROW_ERROR)
 
-        progress_bar = self._get_progress_bar(progress_bar_type)
+        owns_bqstorage_client = False
+        if not bqstorage_client and create_bqstorage_client:
+            owns_bqstorage_client = True
+            bqstorage_client = self.client._create_bqstorage_client()
 
-        record_batches = []
-        for record_batch in self._to_arrow_iterable(bqstorage_client=bqstorage_client):
-            record_batches.append(record_batch)
+        try:
+            progress_bar = self._get_progress_bar(progress_bar_type)
 
-            if progress_bar is not None:
-                # In some cases, the number of total rows is not populated
-                # until the first page of rows is fetched. Update the
-                # progress bar's total to keep an accurate count.
-                progress_bar.total = progress_bar.total or self.total_rows
-                progress_bar.update(record_batch.num_rows)
+            record_batches = []
+            for record_batch in self._to_arrow_iterable(
+                bqstorage_client=bqstorage_client
+            ):
+                record_batches.append(record_batch)
 
-        if progress_bar is not None:
-            # Indicate that the download has finished.
-            progress_bar.close()
+                if progress_bar is not None:
+                    # In some cases, the number of total rows is not populated
+                    # until the first page of rows is fetched. Update the
+                    # progress bar's total to keep an accurate count.
+                    progress_bar.total = progress_bar.total or self.total_rows
+                    progress_bar.update(record_batch.num_rows)
+
+            if progress_bar is not None:
+                # Indicate that the download has finished.
+                progress_bar.close()
+        finally:
+            if owns_bqstorage_client:
+                bqstorage_client.transport.channel.close()
 
         if record_batches:
             return pyarrow.Table.from_batches(record_batches)
@@ -1558,14 +1584,20 @@ def _to_dataframe_iterable(self, bqstorage_client=None, dtypes=None):
 
     # If changing the signature of this method, make sure to apply the same
     # changes to job.QueryJob.to_dataframe()
-    def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=None):
+    def to_dataframe(
+        self,
+        bqstorage_client=None,
+        dtypes=None,
+        progress_bar_type=None,
+        create_bqstorage_client=False,
+    ):
         """Create a pandas DataFrame by loading all pages of a query.
 
         Args:
             bqstorage_client (google.cloud.bigquery_storage_v1beta1.BigQueryStorageClient):
                 **Beta Feature** Optional. A BigQuery Storage API client. If
                 supplied, use the faster BigQuery Storage API to fetch rows
-                from BigQuery. This API is a billable API.
+                from BigQuery.
 
                 This method requires the ``pyarrow`` and
                 ``google-cloud-bigquery-storage`` libraries.
@@ -1602,6 +1634,16 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non
                   progress bar as a graphical dialog box.
 
                 ..versionadded:: 1.11.0
+            create_bqstorage_client (bool):
+                **Beta Feature** Optional. If ``True``, create a BigQuery
+                Storage API client using the default API settings. The
+                BigQuery Storage API is a faster way to fetch rows from
+                BigQuery. See the ``bqstorage_client`` parameter for more
+                information.
+
+                This argument does nothing if ``bqstorage_client`` is supplied.
+
+                ..versionadded:: 1.24.0
 
         Returns:
             pandas.DataFrame:
@@ -1621,32 +1663,44 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non
         if dtypes is None:
             dtypes = {}
 
-        if bqstorage_client and self.max_results is not None:
+        if (
+            bqstorage_client or create_bqstorage_client
+        ) and self.max_results is not None:
             warnings.warn(
                 "Cannot use bqstorage_client if max_results is set, "
                 "reverting to fetching data with the tabledata.list endpoint.",
                 stacklevel=2,
             )
+            create_bqstorage_client = False
             bqstorage_client = None
 
-        progress_bar = self._get_progress_bar(progress_bar_type)
+        owns_bqstorage_client = False
+        if not bqstorage_client and create_bqstorage_client:
+            owns_bqstorage_client = True
+            bqstorage_client = self.client._create_bqstorage_client()
 
-        frames = []
-        for frame in self._to_dataframe_iterable(
-            bqstorage_client=bqstorage_client, dtypes=dtypes
-        ):
-            frames.append(frame)
+        try:
+            progress_bar = self._get_progress_bar(progress_bar_type)
 
-            if progress_bar is not None:
-                # In some cases, the number of total rows is not populated
-                # until the first page of rows is fetched. Update the
-                # progress bar's total to keep an accurate count.
-                progress_bar.total = progress_bar.total or self.total_rows
-                progress_bar.update(len(frame))
+            frames = []
+            for frame in self._to_dataframe_iterable(
+                bqstorage_client=bqstorage_client, dtypes=dtypes
+            ):
+                frames.append(frame)
+
+                if progress_bar is not None:
+                    # In some cases, the number of total rows is not populated
+                    # until the first page of rows is fetched. Update the
+                    # progress bar's total to keep an accurate count.
+                    progress_bar.total = progress_bar.total or self.total_rows
+                    progress_bar.update(len(frame))
 
-        if progress_bar is not None:
-            # Indicate that the download has finished.
-            progress_bar.close()
+            if progress_bar is not None:
+                # Indicate that the download has finished.
+                progress_bar.close()
+        finally:
+            if owns_bqstorage_client:
+                bqstorage_client.transport.channel.close()
 
         # Avoid concatting an empty list.
         if not frames:
@@ -1667,11 +1721,18 @@ class _EmptyRowIterator(object):
     pages = ()
     total_rows = 0
 
-    def to_arrow(self, progress_bar_type=None):
+    def to_arrow(
+        self,
+        progress_bar_type=None,
+        bqstorage_client=None,
+        create_bqstorage_client=False,
+    ):
         """[Beta] Create an empty class:`pyarrow.Table`.
 
         Args:
             progress_bar_type (Optional[str]): Ignored. Added for compatibility with RowIterator.
+            bqstorage_client (Any): Ignored. Added for compatibility with RowIterator.
+            create_bqstorage_client (bool): Ignored. Added for compatibility with RowIterator.
 
         Returns:
             pyarrow.Table: An empty :class:`pyarrow.Table`.
@@ -1680,13 +1741,20 @@ def to_arrow(self, progress_bar_type=None):
             raise ValueError(_NO_PYARROW_ERROR)
         return pyarrow.Table.from_arrays(())
 
-    def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=None):
+    def to_dataframe(
+        self,
+        bqstorage_client=None,
+        dtypes=None,
+        progress_bar_type=None,
+        create_bqstorage_client=False,
+    ):
         """Create an empty dataframe.
 
         Args:
             bqstorage_client (Any): Ignored. Added for compatibility with RowIterator.
             dtypes (Any): Ignored. Added for compatibility with RowIterator.
             progress_bar_type (Any): Ignored. Added for compatibility with RowIterator.
+            create_bqstorage_client (bool): Ignored. Added for compatibility with RowIterator.
 
         Returns:
             pandas.DataFrame: An empty :class:`~pandas.DataFrame`.

bigquery/samples/download_public_data.py

@@ -0,0 +1,33 @@
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def download_public_data(client):
+
+    # [START bigquery_pandas_public_data]
+    # TODO(developer): Import the client library.
+    # from google.cloud import bigquery
+
+    # TODO(developer): Construct a BigQuery client object.
+    # client = bigquery.Client()
+
+    # TODO(developer): Set table_id to the fully-qualified table ID in standard
+    # SQL format, including the project ID and dataset ID.
+    table_id = "bigquery-public-data.usa_names.usa_1910_current"
+
+    # Use the BigQuery Storage API to speed-up downloads of large tables.
+    dataframe = client.list_rows(table_id).to_dataframe(create_bqstorage_client=True)
+
+    print(dataframe.info())
+    # [END bigquery_pandas_public_data]

bigquery/samples/download_public_data_sandbox.py

@@ -0,0 +1,34 @@
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def download_public_data_sandbox(client):
+
+    # [START bigquery_pandas_public_data_sandbox]
+    # TODO(developer): Import the client library.
+    # from google.cloud import bigquery
+
+    # TODO(developer): Construct a BigQuery client object.
+    # client = bigquery.Client()
+
+    # `SELECT *` is an anti-pattern in BigQuery because it is cheaper and
+    # faster to use the BigQuery Storage API directly, but BigQuery Sandbox
+    # users can only use the BigQuery Storage API to download query results.
+    query_string = "SELECT * FROM `bigquery-public-data.usa_names.usa_1910_current`"
+
+    # Use the BigQuery Storage API to speed-up downloads of large tables.
+    dataframe = client.query(query_string).to_dataframe(create_bqstorage_client=True)
+
+    print(dataframe.info())
+    # [END bigquery_pandas_public_data_sandbox]