Make `max_stream_count` configurable when using Bigquery Storage API

[kien-truong]

Currently, for API that can use BQ Storage Client to fetch data like to_dataframe_iterable or to_arrow_iterable, the client library always uses the maximum number of read streams recommended by BQ server.

python-bigquery/google/cloud/bigquery/_pandas_helpers.py

Line 840 in ef8e927

requested_streams = 1 if preserve_order else 0

python-bigquery/google/cloud/bigquery/_pandas_helpers.py

Lines 854 to 858 in ef8e927

	session = bqstorage_client.create_read_session(
	parent="projects/{}".format(project_id),
	read_session=requested_session,
	max_stream_count=requested_streams,
	)

This behavior has the advantage of maximizing throughput but can lead to out-of-memory issue when there are too many streams being opened and result are not read fast enough: we've encountered queries that open hundreds of streams and consuming GBs of memory.

BQ Storage Client API also suggests capping max_stream_count when resource is constrained

https://cloud.google.com/bigquery/docs/reference/storage/rpc/google.cloud.bigquery.storage.v1#createreadsessionrequest

Typically, clients should either leave this unset to let the system to determine an upper bound OR set this a size for the maximum "units of work" it can gracefully handle.

This problem has been encountered by others before and can be worked-around by monkey-patching the create_read_session on the BQ Client object: #1292

However, it should really be fixed by allowing the max_stream_count parameter to be set through public API.

[kien-truong]

Using the default setting, in the worst-case scenario, for n-download streams, we would have to store 2n result pages in memory:

• 1 result page inside each download thread, times n threads
• n result page in the transfer queue between download threads and main thread

[chalmerlowe]

I am putting together a PR for this.
It is not as simple as just adding a max_stream_count argument because historically the argument preserve_order has been used to make some choices regarding the number of streams to use.

We have some logic to allow those two to interact in a way that makes sense and hopefully doesn't break backward compatibility.

We are also issuing a docstring to explain the method as a whole (this feels necessary since we have a new argument that may interact with another argument).

More to come, soon.

[kien-truong]

Right, you can't use more than 1 stream if you need to preserve row order when using ORDER BY, the concurrency will screw up the output.

[chalmerlowe]

Alright. Got a draft PR. Still working on the unit tests.
But I need to focus on something else right now. Will come back to this with a clear head.

[chalmerlowe]

@kien-truong let me know if the updates meet the need.
The code is available in main and will included in the next release. (I don't have a date yet for the next release.)

[kien-truong]

Thanks @chalmerlowe, can you also add the max_stream_count argument to 2 methods to_dataframe_iterable and to_arrow_iterable?

python-bigquery/google/cloud/bigquery/table.py

Line 1811 in 7372ad6

def to_arrow_iterable(

python-bigquery/google/cloud/bigquery/table.py

Line 1976 in 7372ad6

def to_dataframe_iterable(

Those 2 methods are where this is most useful to support incremental fetching use cases.