Repository navigation
feat(bigquery): share one BigQuery Storage client per Client - #18652
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a shared, thread-safe BigQuery Storage client for ARROW query results, allowing multiple queries to reuse a single client instance. It also configures the shared client's gRPC transport with keepalive pings to prevent idle connections from being dropped, and ensures the shared client is closed when the main BigQuery client is closed. The feedback suggests defensively handling the options parameter in create_channel to prevent a potential TypeError if None is passed.
| def create_channel(*args, options=(), **channel_kwargs): | ||
| options = list(options) + [ | ||
| ("grpc.keepalive_time_ms", _BQSTORAGE_KEEPALIVE_TIME_MS) | ||
| ] | ||
| return transport_class.create_channel(*args, options=options, **channel_kwargs) |
There was a problem hiding this comment.
If options is explicitly passed as None to create_channel, calling list(options) will raise a TypeError: 'NoneType' object is not iterable. Using options or [] ensures defensive handling and prevents potential runtime crashes.
| def create_channel(*args, options=(), **channel_kwargs): | |
| options = list(options) + [ | |
| ("grpc.keepalive_time_ms", _BQSTORAGE_KEEPALIVE_TIME_MS) | |
| ] | |
| return transport_class.create_channel(*args, options=options, **channel_kwargs) | |
| def create_channel(*args, options=None, **channel_kwargs): | |
| options = list(options or []) + [ | |
| ("grpc.keepalive_time_ms", _BQSTORAGE_KEEPALIVE_TIME_MS) | |
| ] | |
| return transport_class.create_channel(*args, options=options, **channel_kwargs) |
References
- When processing optional parameters with a helper function, prefer placing the null or None check within the function body rather than at the call site to simplify the calling code and improve encapsulation.
47c18a9 to
18c7c06
Compare
3cba84a to
54bdf24
Compare
| options = list(options) + [ | ||
| ("grpc.keepalive_time_ms", _BQSTORAGE_KEEPALIVE_TIME_MS), | ||
| ("grpc.keepalive_timeout_ms", _BQSTORAGE_KEEPALIVE_TIMEOUT_MS), | ||
| ] | ||
| return transport_class.create_channel(*args, options=options, **channel_kwargs) |
There was a problem hiding this comment.
Would it make more sense to upstream this to the google-cloud-bigquery-storage client library?
There was a problem hiding this comment.
Yeah, I think it belongs in bigquery-storage long term. Any long ReadRows read could use it, not just ours.
There was a problem hiding this comment.
I'd rather do it as a follow-up though. We'd only get it here once our bqstorage minimum includes it, so it'd mean another release and another floor bump. Filed #18784 for it and left a TODO on _create_bqstorage_transport. There's already a ReadRowsStream change waiting on the next bqstorage release (#18781), so it could go out with that.
|
|
||
| return bqstorage_client | ||
|
|
||
| def _get_shared_bqstorage_client( |
There was a problem hiding this comment.
Why introduce a new method instead of doing this in_ensure_bqstorage_client or renaming _ensure_bqstorage_client to _get_shared_bqstorage_client?
There was a problem hiding this comment.
Every caller of _ensure_bqstorage_client() gets its own client and closes the transport when it's done (to_dataframe(), DB-API, magics). If it handed back the shared one, the first to_dataframe() would close the channel the Arrow queries are using.
There was a problem hiding this comment.
So I kept them separate: _ensure_bqstorage_client() gives you a client you own, and _get_shared_bqstorage_client() gives you one only Client.close() closes. Added a line to the docstring saying that.
There was a problem hiding this comment.
Gotcha. We should reconsider automatically closing in those methods at some point, but out of scope for these changes.
54bdf24 to
7e419d3
Compare
61eec40 to
eb2e0fc
Compare
30a1c50 to
928b0de
Compare
tswast
left a comment
tswast
left a comment
There was a problem hiding this comment.
Thanks! I suspect this will have a big impact on performance for long-lived client objects.
ARROW query results created a new BigQueryReadClient for every read, even ones that never opened a stream, and to_arrow_iterable() never closed it, so each read left a gRPC channel open until garbage collection. The new client also ignored the BigQuery client's universe domain. Client now creates one BigQueryReadClient under a lock the first time an ARROW result needs a stream, with its credentials and universe domain, and reuses it. close() closes it, and a later query creates a new one. Results that fit in the first page never create it. A bqstorage_client passed to to_arrow_iterable() is used instead and never closed. Other BigQuery Storage reads keep creating their own clients. The shared client's gRPC channel sets keepalive (60s ping time, 60s ping timeout) through the generated transport's channel callable. gRPC stops pinging after two pings while ReadRows sends nothing, so longer pauses rely on ReadRowsStream resuming at the next unread row.
bigquery-storage 2.43.0's ReadRowsStream.__del__ also calls cancel() on streams that finished normally. The call sends no RPC, but the system test counted it, so it failed with the released library. Count only cancels of streams that are still open, until the bqstorage extra requires a release with the fix (#18785).
Note that callers must not close the shared BigQuery Storage client, unlike the clients from _ensure_bqstorage_client(). Add a TODO to move keepalive into google-cloud-bigquery-storage (#18784).
The no-schema tests from main mock _ensure_bqstorage_client(), but download_arrow_from_job now gets its client from _get_shared_bqstorage_client(). Mock that instead, like the other Arrow tests.
Client.__init__ already sets _universe_domain, and _ensure_bqstorage_client() adds the universe domain when no client_options are passed. Drop the duplicate assignment and stop building ClientOptions in _get_shared_bqstorage_client().
2330555 to
68d6bfb
Compare
tswast
added
cla: yes
yoshi-kokoro
removed
the
kokoro:force-run
_get_shared_bqstorage_client() no longer passes client_options to _ensure_bqstorage_client(), so the system test's expected call failed. Match the unit tests and expect only the transport argument.
ohmayr
added
the
kokoro:force-run
yoshi-kokoro
removed
the
kokoro:force-run
Makes `BigQueryReadClient` send gRPC keepalive pings by default. `ReadRows` streams can sit idle for a long time, for example while BigQuery finishes a query job, and without pings a NAT or load balancer can drop the idle connection without telling either side. google-cloud-bigquery works around this with its own transport for Arrow query results (#18652). This moves the setting here, so every long read gets it. - Uses a gRPC channel with `grpc.keepalive_time_ms` and `grpc.keepalive_timeout_ms` set to 60s when no `transport` is passed (or `transport="grpc"`). - Leaves a caller's transport instance or callable unchanged. - Adds unit tests for the default channel options and for caller-provided transports. This changes the default channel for every bigquery-storage user, but normal table reads shouldn't notice: pings only go out while a call is open and the connection has been quiet for 60s (`keepalive_permit_without_calls` isn't set). A read on a dropped connection now fails with `UNAVAILABLE` after about 2 minutes instead of hanging, and `ReadRowsStream` already reconnects on that. google-cloud-bigtable and google-cloud-firestore already send keepalive pings by default, every 30s. Fixes #18784
3 participants
Shares one
BigQueryReadClientperClientfor Arrow results, instead of creating a new one for every read and never closing it.Client.close().bqstorage_clientpassed by the caller instead, without closing it.grpc.keepalive_time_ms: 60000,grpc.keepalive_timeout_ms: 60000).