Fix list_rows() max results with BQ storage client (#9178)

If max results is set, the streaming API should not be used, and a
user warning is issued.

Author: plamut

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 = []

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:

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`"