Repository: googleapis/google-cloud-python

Issue:
BigQuery: `max_results` is ignored if `bqstorage_client` is used in `to_dataframe` or `to_arrow`
#### Steps to reproduce

1. Call `list_rows` with `max_results` set.
2. Call `to_dataframe` or `to_arrow`.
3. Observe that more rows were returned than were requested.

#### Code example

```python
from google.cloud import bigquery
from google.cloud import bigquery_storage
bqclient = bigquery.Client()
bqstorage_client = bigquery_storage.BigQueryStorageClient()

df_tabledata_list = bqclient.list_rows(
    "bigquery-public-data.utility_us.country_code_iso",
    selected_fields=[bigquery.SchemaField("country_name", "STRING")],
    max_results=100,
).to_dataframe()
print("tabledata.list: {} rows".format(len(df_tabledata_list.index)))

df_bqstorage = bqclient.list_rows(
    "bigquery-public-data.utility_us.country_code_iso",
    selected_fields=[bigquery.SchemaField("country_name", "STRING")],
    max_results=100,
).to_dataframe(bqstorage_client=bqstorage_client)
print("bqstorage: {} rows".format(len(df_bqstorage.index)))
```

#### Output

```
tabledata.list: 100 rows
bqstorage: 278 rows
```

#### Possible fixes

1. (Harder) Keep track of how many rows you've downloaded in a BQ Storage session so far. Once you've downloaded enough rows, close all streams (is this even possible?).
2. (Easier, but acceptable) If `max_results` is set, always download data with tabledata.list.

I think we should implement the fix with (2) because if `max_results` is set, it's unlikely that we are downloading all that many rows where using the BQ Storage API makes sense.



Discussion:
@tswast After implementing the second option, would it make sense to spend some time on researching the option 1? It might turn out that streams _can_ be closed at the right time reliably.
> After implementing the second option, would it make sense to spend some time on researching the option 1?

I worry that the additional logic of closing streams may be more complicated than it's worth, but it's worth investigating to see if it's possible.

Resolving patch:
diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py
--- 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 = []


Test patch:
diff --git a/bigquery/tests/system.py b/bigquery/tests/system.py
--- 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
--- 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`"