From 89fccc107709744e2f1d31bf71abf4784a431d75 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 29 Mar 2019 09:44:24 -0700 Subject: [PATCH 1/3] Actually use the progress_bar_type argument in QueryJob.to_dataframe() --- bigquery/google/cloud/bigquery/job.py | 4 +++- bigquery/google/cloud/bigquery/table.py | 8 ++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index bc87f109a484..94a2290cc29e 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -2857,7 +2857,9 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non ValueError: If the `pandas` library cannot be imported. """ return self.result().to_dataframe( - bqstorage_client=bqstorage_client, dtypes=dtypes + bqstorage_client=bqstorage_client, + dtypes=dtypes, + progress_bar_type=progress_bar_type, ) def __iter__(self): diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py index dcb25d8bb3c6..5b0ad45af4be 100644 --- a/bigquery/google/cloud/bigquery/table.py +++ b/bigquery/google/cloud/bigquery/table.py @@ -1422,9 +1422,7 @@ def _get_progress_bar(self, progress_bar_type): desc=description, total=self.total_rows, unit=unit ) elif progress_bar_type == "tqdm_gui": - return tqdm.tqdm_gui( - desc=description, total=self.total_rows, unit=unit - ) + return tqdm.tqdm_gui(desc=description, total=self.total_rows, unit=unit) except (KeyError, TypeError): # Protect ourselves from any tqdm errors. In case of # unexpected tqdm behavior, just fall back to showing @@ -1515,7 +1513,7 @@ class _EmptyRowIterator(object): pages = () total_rows = 0 - def to_dataframe(self, bqstorage_client=None, dtypes=None): + def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=None): """Create an empty dataframe. Args: @@ -1523,6 +1521,8 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None): Ignored. Added for compatibility with RowIterator. dtypes (Any): Ignored. Added for compatibility with RowIterator. + progress_bar_type (Any): + Ignored. Added for compatibility with RowIterator. Returns: pandas.DataFrame: From 0b73e32e1c156036cad8ad2ccebabec34682cc17 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 29 Mar 2019 15:53:16 -0700 Subject: [PATCH 2/3] Test that a progress bar is constructed from QueryJob.to_dataframe --- bigquery/tests/unit/test_job.py | 34 ++++++++++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py index 833081ce066d..862f404d08c9 100644 --- a/bigquery/tests/unit/test_job.py +++ b/bigquery/tests/unit/test_job.py @@ -27,7 +27,10 @@ from google.cloud import bigquery_storage_v1beta1 except (ImportError, AttributeError): # pragma: NO COVER bigquery_storage_v1beta1 = None - +try: + from tqdm import tqdm +except (ImportError, AttributeError): # pragma: NO COVER + tqdm = None def _make_credentials(): import google.auth.credentials @@ -4699,6 +4702,35 @@ def test_to_dataframe_column_dtypes(self): self.assertEqual(df.complete.dtype.name, "bool") self.assertEqual(df.date.dtype.name, "object") + @unittest.skipIf(pandas is None, "Requires `pandas`") + @unittest.skipIf(tqdm is None, "Requires `tqdm`") + @mock.patch("tqdm.tqdm") + def test_to_dataframe_with_progress_bar(self, tqdm_mock): + begun_resource = self._make_resource() + query_resource = { + "jobComplete": True, + "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID}, + "totalRows": "4", + "schema": { + "fields": [ + {"name": "name", "type": "STRING", "mode": "NULLABLE"}, + ] + }, + } + done_resource = copy.deepcopy(begun_resource) + done_resource["status"] = {"state": "DONE"} + connection = _make_connection( + begun_resource, query_resource, done_resource, query_resource, query_resource, + ) + client = _make_client(project=self.PROJECT, connection=connection) + job = self._make_one(self.JOB_ID, self.QUERY, client) + + job.to_dataframe(progress_bar_type=None) + tqdm_mock.assert_not_called() + + job.to_dataframe(progress_bar_type="tqdm") + tqdm_mock.assert_called() + def test_iter(self): import types From 3b1d0d49ee570e04a0138c763a6aedd609bf7761 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 29 Mar 2019 15:55:25 -0700 Subject: [PATCH 3/3] Blacken --- bigquery/tests/unit/test_job.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py index 862f404d08c9..a42d9ffc311c 100644 --- a/bigquery/tests/unit/test_job.py +++ b/bigquery/tests/unit/test_job.py @@ -32,6 +32,7 @@ except (ImportError, AttributeError): # pragma: NO COVER tqdm = None + def _make_credentials(): import google.auth.credentials @@ -4712,15 +4713,17 @@ def test_to_dataframe_with_progress_bar(self, tqdm_mock): "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID}, "totalRows": "4", "schema": { - "fields": [ - {"name": "name", "type": "STRING", "mode": "NULLABLE"}, - ] + "fields": [{"name": "name", "type": "STRING", "mode": "NULLABLE"}] }, } done_resource = copy.deepcopy(begun_resource) done_resource["status"] = {"state": "DONE"} connection = _make_connection( - begun_resource, query_resource, done_resource, query_resource, query_resource, + begun_resource, + query_resource, + done_resource, + query_resource, + query_resource, ) client = _make_client(project=self.PROJECT, connection=connection) job = self._make_one(self.JOB_ID, self.QUERY, client)