ENH: show progress bar when downloading data (#292)

* add progress bar

* update changelog

* default to use the progress bar

Author: dakl

docs/source/changelog.rst

@@ -11,6 +11,9 @@ Changelog
   ``max_results`` to 0 to ignore query outputs, such as for DML or DDL
   queries. (:issue:`102`)
 
+- Add ``progress_bar_type`` argument to :func:`~pandas_gbq.read_gbq()`. Use this
+  argument to display a progress bar when downloading data. (:issue:`182`)
+
 Documentation
 ~~~~~~~~~~~~~
 

pandas_gbq/gbq.py

@@ -436,7 +436,9 @@ def process_http_error(ex):
 
         raise GenericGBQException("Reason: {0}".format(ex))
 
-    def run_query(self, query, max_results=None, **kwargs):
+    def run_query(
+        self, query, max_results=None, progress_bar_type=None, **kwargs
+    ):
         from concurrent.futures import TimeoutError
         from google.auth.exceptions import RefreshError
         from google.cloud import bigquery
@@ -526,9 +528,15 @@ def run_query(self, query, max_results=None, **kwargs):
                 )
             )
 
-        return self._download_results(query_reply, max_results=max_results)
+        return self._download_results(
+            query_reply,
+            max_results=max_results,
+            progress_bar_type=progress_bar_type,
+        )
 
-    def _download_results(self, query_job, max_results=None):
+    def _download_results(
+        self, query_job, max_results=None, progress_bar_type=None
+    ):
         # No results are desired, so don't bother downloading anything.
         if max_results == 0:
             return None
@@ -552,7 +560,9 @@ def _download_results(self, query_job, max_results=None):
         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, bqstorage_client=bqstorage_client
+            dtypes=nullsafe_dtypes,
+            bqstorage_client=bqstorage_client,
+            progress_bar_type=progress_bar_type,
         )
 
         if df.empty:
@@ -833,6 +843,7 @@ def read_gbq(
     max_results=None,
     verbose=None,
     private_key=None,
+    progress_bar_type="tqdm",
 ):
     r"""Load data from Google BigQuery using google-cloud-python
 
@@ -952,6 +963,23 @@ def read_gbq(
         or string contents. This is useful for remote server
         authentication (eg. Jupyter/IPython notebook on remote host).
 
+    progress_bar_type (Optional[str]):
+        If set, use the `tqdm <https://tqdm.github.io/>`_ library to
+        display a progress bar while the data downloads. Install the
+        ``tqdm`` package to use this feature.
+        Possible values of ``progress_bar_type`` include:
+        ``None``
+            No progress bar.
+        ``'tqdm'``
+            Use the :func:`tqdm.tqdm` function to print a progress bar
+            to :data:`sys.stderr`.
+        ``'tqdm_notebook'``
+            Use the :func:`tqdm.tqdm_notebook` function to display a
+            progress bar as a Jupyter notebook widget.
+        ``'tqdm_gui'``
+            Use the :func:`tqdm.tqdm_gui` function to display a
+            progress bar as a graphical dialog box.
+
     Returns
     -------
     df: DataFrame
@@ -996,7 +1024,10 @@ def read_gbq(
     )
 
     final_df = connector.run_query(
-        query, configuration=configuration, max_results=max_results
+        query,
+        configuration=configuration,
+        max_results=max_results,
+        progress_bar_type=progress_bar_type,
     )
 
     # Reindex the DataFrame on the provided column

tests/unit/test_gbq.py

@@ -552,3 +552,22 @@ def test_load_does_not_modify_schema_arg(mock_bigquery_client):
         if_exists="append",
     )
     assert original_schema == original_schema_cp
+
+
+def test_read_gbq_calls_tqdm(
+    mock_bigquery_client, mock_service_account_credentials
+):
+    mock_service_account_credentials.project_id = "service_account_project_id"
+    df = gbq.read_gbq(
+        "SELECT 1",
+        dialect="standard",
+        credentials=mock_service_account_credentials,
+        progress_bar_type="foobar",
+    )
+    assert df is not None
+
+    mock_list_rows = mock_bigquery_client.list_rows("dest", max_results=100)
+
+    mock_list_rows.to_dataframe.assert_called_once_with(
+        dtypes=mock.ANY, bqstorage_client=mock.ANY, progress_bar_type="foobar"
+    )