diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py index 62072cf88804..71fc9ef945d4 100644 --- a/bigquery/google/cloud/bigquery/table.py +++ b/bigquery/google/cloud/bigquery/table.py @@ -1630,6 +1630,14 @@ 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: + warnings.warn( + "Cannot use bqstorage_client if max_results is set, " + "reverting to fetching data with the tabledata.list endpoint.", + stacklevel=2, + ) + bqstorage_client = None + progress_bar = self._get_progress_bar(progress_bar_type) frames = [] diff --git a/bigquery/tests/system.py b/bigquery/tests/system.py index 3593e1ecb609..b1e583bbfaa3 100644 --- a/bigquery/tests/system.py +++ b/bigquery/tests/system.py @@ -2315,6 +2315,27 @@ def test_list_rows_page_size(self): page = next(pages) self.assertEqual(page.num_items, num_last_page) + @unittest.skipIf(pandas is None, "Requires `pandas`") + @unittest.skipIf( + bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" + ) + def test_list_rows_max_results_w_bqstorage(self): + table_ref = DatasetReference("bigquery-public-data", "utility_us").table( + "country_code_iso" + ) + bqstorage_client = bigquery_storage_v1beta1.BigQueryStorageClient( + credentials=Config.CLIENT._credentials + ) + + row_iterator = Config.CLIENT.list_rows( + table_ref, + selected_fields=[bigquery.SchemaField("country_name", "STRING")], + max_results=100, + ) + dataframe = row_iterator.to_dataframe(bqstorage_client=bqstorage_client) + + self.assertEqual(len(dataframe.index), 100) + def temp_dataset(self, dataset_id, location=None): dataset = Dataset(Config.CLIENT.dataset(dataset_id)) if location: diff --git a/bigquery/tests/unit/test_table.py b/bigquery/tests/unit/test_table.py index 8ba7fee892e5..d36eb43e4feb 100644 --- a/bigquery/tests/unit/test_table.py +++ b/bigquery/tests/unit/test_table.py @@ -2208,6 +2208,42 @@ def test_to_dataframe_error_if_pandas_is_none(self): with self.assertRaises(ValueError): row_iterator.to_dataframe() + @unittest.skipIf(pandas is None, "Requires `pandas`") + def test_to_dataframe_max_results_w_bqstorage_warning(self): + from google.cloud.bigquery.table import SchemaField + + schema = [ + SchemaField("name", "STRING", mode="REQUIRED"), + SchemaField("age", "INTEGER", mode="REQUIRED"), + ] + rows = [ + {"f": [{"v": "Phred Phlyntstone"}, {"v": "32"}]}, + {"f": [{"v": "Bharney Rhubble"}, {"v": "33"}]}, + ] + path = "/foo" + api_request = mock.Mock(return_value={"rows": rows}) + bqstorage_client = mock.Mock() + + row_iterator = self._make_one( + client=_mock_client(), + api_request=api_request, + path=path, + schema=schema, + max_results=42, + ) + + with warnings.catch_warnings(record=True) as warned: + row_iterator.to_dataframe(bqstorage_client=bqstorage_client) + + matches = [ + warning + for warning in warned + if warning.category is UserWarning + and "cannot use bqstorage_client" in str(warning).lower() + and "tabledata.list" in str(warning) + ] + self.assertEqual(len(matches), 1, msg="User warning was not emitted.") + @unittest.skipIf(pandas is None, "Requires `pandas`") @unittest.skipIf( bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`"