BigQuery: max_results is ignored if bqstorage_client is used in to_dataframe or to_arrow

[tswast]

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

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.

[plamut]

@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.

[tswast]

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.

[plamut]

There exists the FinalizeStream method, but is not a silver bullet:

• Cancelling a stream does not delete it, but just causes that no additional data is assigned to the stream. As per docs, the client must still continue reading the stream to process any data that has already been allocated to the stream.
• While we could read the additional data and discard it, users would then be charged for more rows / data that they requested (I presume?). Additionally, even a single response block might contain enough rows to exceed the max_results limit, again making the cost higher than expected.
• Ignoring any remaining data on the stream (by stopping reading from it) could leak resources on the backend, especially if done excessively by the client. We probably don't want to take this route anyway.

TL; DR - by trying to mimic the max_results in the client, it seems that we could get in more trouble than it's worth.