diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py index ab22407eff1a..ae090a4b0b83 100644 --- a/bigquery/google/cloud/bigquery/table.py +++ b/bigquery/google/cloud/bigquery/table.py @@ -35,6 +35,7 @@ except ImportError: # pragma: NO COVER tqdm = None +import google.api_core.exceptions from google.api_core.page_iterator import HTTPIterator import google.cloud._helpers @@ -1437,7 +1438,7 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non bqstorage_client ( \ google.cloud.bigquery_storage_v1beta1.BigQueryStorageClient \ ): - **Alpha Feature** Optional. A BigQuery Storage API client. If + **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. @@ -1448,8 +1449,9 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non currently supported by this method. **Caution**: There is a known issue reading small anonymous - query result tables with the BQ Storage API. Write your query - results to a destination table to work around this issue. + query result tables with the BQ Storage API. When a problem + is encountered reading a table, the tabledata.list method + from the BigQuery API is used, instead. dtypes ( \ Map[str, Union[str, pandas.Series.dtype]] \ ): @@ -1496,9 +1498,16 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non progress_bar = self._get_progress_bar(progress_bar_type) if bqstorage_client is not None: - return self._to_dataframe_bqstorage(bqstorage_client, dtypes) - else: - return self._to_dataframe_tabledata_list(dtypes, progress_bar=progress_bar) + try: + return self._to_dataframe_bqstorage(bqstorage_client, dtypes) + except google.api_core.exceptions.GoogleAPICallError: + # There is a known issue with reading from small anonymous + # query results tables, so some errors are expected. Rather + # than throw those errors, try reading the DataFrame again, but + # with the tabledata.list API. + pass + + return self._to_dataframe_tabledata_list(dtypes, progress_bar=progress_bar) class _EmptyRowIterator(object): diff --git a/bigquery/tests/system.py b/bigquery/tests/system.py index 04c4c3bf18fd..cceca192b8f7 100644 --- a/bigquery/tests/system.py +++ b/bigquery/tests/system.py @@ -1540,35 +1540,42 @@ def test_query_results_to_dataframe_w_bqstorage(self): bqstorage_client = bigquery_storage_v1beta1.BigQueryStorageClient( credentials=Config.CLIENT._credentials ) - df = ( - Config.CLIENT.query( - query, - # There is a known issue reading small anonymous query result - # tables with the BQ Storage API. Writing to a destination - # table works around this issue. - job_config=bigquery.QueryJobConfig( - destination=dest_ref, write_disposition="WRITE_TRUNCATE" - ), - ) - .result() - .to_dataframe(bqstorage_client) + + job_configs = ( + # There is a known issue reading small anonymous query result + # tables with the BQ Storage API. Writing to a destination + # table works around this issue. + bigquery.QueryJobConfig( + destination=dest_ref, write_disposition="WRITE_TRUNCATE" + ), + # Check that the client is able to work around the issue with + # reading small anonymous query result tables by falling back to + # the tabledata.list API. + None, ) - self.assertIsInstance(df, pandas.DataFrame) - self.assertEqual(len(df), 10) # verify the number of rows - column_names = ["id", "author", "time_ts", "dead"] - self.assertEqual(list(df), column_names) - exp_datatypes = { - "id": int, - "author": six.text_type, - "time_ts": pandas.Timestamp, - "dead": bool, - } - for index, row in df.iterrows(): - for col in column_names: - # all the schema fields are nullable, so None is acceptable - if not row[col] is None: - self.assertIsInstance(row[col], exp_datatypes[col]) + for job_config in job_configs: + df = ( + Config.CLIENT.query(query, job_config=job_config) + .result() + .to_dataframe(bqstorage_client) + ) + + self.assertIsInstance(df, pandas.DataFrame) + self.assertEqual(len(df), 10) # verify the number of rows + column_names = ["id", "author", "time_ts", "dead"] + self.assertEqual(list(df), column_names) + exp_datatypes = { + "id": int, + "author": six.text_type, + "time_ts": pandas.Timestamp, + "dead": bool, + } + for index, row in df.iterrows(): + for col in column_names: + # all the schema fields are nullable, so None is acceptable + if not row[col] is None: + self.assertIsInstance(row[col], exp_datatypes[col]) def test_insert_rows_nested_nested(self): # See #2951 diff --git a/bigquery/tests/unit/test_table.py b/bigquery/tests/unit/test_table.py index 4500856ec2a4..b0a1318c6f11 100644 --- a/bigquery/tests/unit/test_table.py +++ b/bigquery/tests/unit/test_table.py @@ -21,10 +21,13 @@ import pytest import six +import google.api_core.exceptions + try: from google.cloud import bigquery_storage_v1beta1 except ImportError: # pragma: NO COVER bigquery_storage_v1beta1 = None + try: import pandas except (ImportError, AttributeError): # pragma: NO COVER @@ -1748,6 +1751,48 @@ def test_to_dataframe_w_bqstorage_nonempty(self): self.assertEqual(list(got), column_names) self.assertEqual(len(got.index), 2) + @unittest.skipIf(pandas is None, "Requires `pandas`") + @unittest.skipIf( + bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" + ) + def test_to_dataframe_w_bqstorage_fallback_to_tabledata_list(self): + from google.cloud.bigquery import schema + from google.cloud.bigquery import table as mut + + bqstorage_client = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient + ) + bqstorage_client.create_read_session.side_effect = google.api_core.exceptions.InternalServerError( + "can't read with bqstorage_client" + ) + iterator_schema = [ + schema.SchemaField("name", "STRING", mode="REQUIRED"), + schema.SchemaField("age", "INTEGER", mode="REQUIRED"), + ] + rows = [ + {"f": [{"v": "Phred Phlyntstone"}, {"v": "32"}]}, + {"f": [{"v": "Bharney Rhubble"}, {"v": "33"}]}, + {"f": [{"v": "Wylma Phlyntstone"}, {"v": "29"}]}, + {"f": [{"v": "Bhettye Rhubble"}, {"v": "27"}]}, + ] + path = "/foo" + api_request = mock.Mock(return_value={"rows": rows}) + row_iterator = mut.RowIterator( + _mock_client(), + api_request, + path, + iterator_schema, + table=mut.Table("proj.dset.tbl"), + ) + + df = row_iterator.to_dataframe(bqstorage_client=bqstorage_client) + + self.assertIsInstance(df, pandas.DataFrame) + self.assertEqual(len(df), 4) # verify the number of rows + self.assertEqual(list(df), ["name", "age"]) # verify the column names + self.assertEqual(df.name.dtype.name, "object") + self.assertEqual(df.age.dtype.name, "int64") + @unittest.skipIf( bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" )