Repository navigation
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 support for arrow_serialization_options in the BigQuery Storage read_rows client and stream reader, allowing users to specify options like buffer compression. It also implements a cancel() method on ReadRowsStream to stop reading before the end of a stream without reconnecting, along with corresponding unit and system tests. The review feedback highlights a potential race condition in _reconnect() where a newly established stream could be leaked if cancel() is called concurrently, and suggests a code change to safely assign and check the cancellation state.
| if self._cancelled: | ||
| cancel = getattr(wrapped, "cancel", None) | ||
| if cancel is not None: | ||
| cancel() | ||
| return | ||
| self._wrapped = wrapped | ||
| break |
There was a problem hiding this comment.
There is a potential race condition when cancel() is called concurrently from another thread (which is the standard way to cancel a blocking stream iterator) while _reconnect() is establishing a new stream.
If cancel() is called after the self._cancelled check on line 222 but before self._wrapped = wrapped on line 227, the cancel() call will execute on the old stream (or do nothing if it was None). Then, the new stream wrapped is assigned to self._wrapped and the loop breaks. Since self._cancelled is now True, the __iter__ loop will terminate, but the newly created wrapped stream is never cancelled or closed, leading to a resource leak.
To prevent this, assign self._wrapped = wrapped first, and then check self._cancelled. If it is cancelled, set self._wrapped = None (to preserve the expected state) and cancel the stream before returning.
| if self._cancelled: | |
| cancel = getattr(wrapped, "cancel", None) | |
| if cancel is not None: | |
| cancel() | |
| return | |
| self._wrapped = wrapped | |
| break | |
| self._wrapped = wrapped | |
| if self._cancelled: | |
| self._wrapped = None | |
| cancel = getattr(wrapped, "cancel", None) | |
| if cancel is not None: | |
| cancel() | |
| return | |
| break |
There was a problem hiding this comment.
Good catch, the window between the check and the assignment was real. Fixed in f759307: _reconnect() now stores the new call first and then checks _cancelled, so a concurrent cancel() either cancels it directly or gets picked up by the check.
Default job streams (projects/.../jobs/.../streams/_default) accept ArrowSerializationOptions on each ReadRowsRequest, but BigQueryReadClient.read_rows() had no way to send them, and ReadRowsStream dropped them when it reconnected after a transient error. There was also no way to stop a stream early without closing the whole client. read_rows() now takes arrow_serialization_options, and ReadRowsStream sends them on the first request and on every reconnect, so all batches in a stream use the same encoding. ReadRowsStream.cancel() cancels the current call and stops further reconnects, including a reconnect already in progress, without closing the client. google-cloud-bigquery needs both to read query results in Arrow format through the default job stream.
f759307 to
ea74a33
Compare
| offset, | ||
| read_rows_kwargs, | ||
| retry_delay_callback=None, | ||
| arrow_serialization_options=None, |
There was a problem hiding this comment.
Why aren't you using the existing read_rows_kwargs dictionary?
Edit: I see that client is the super() / generated client. Still, I'd prefer adding a more general "request_template" or something like that so that we can pass any new request parameters more easily in future.
There was a problem hiding this comment.
replaced arrow_serialization_options on ReadRowsStream.__init__ with *, request_template=None, and copied it via gapic_types.ReadRowsRequest(self._request_template) in _reconnect() before setting read_stream and offset so caller-supplied templates aren't mutated across reconnects.
|
|
||
| self._reconnect() | ||
|
|
||
| def cancel(self): |
There was a problem hiding this comment.
Thought: should we automatically call cancel in __del__? That might help in cases where folks are just reading one batch and returning.
There was a problem hiding this comment.
added __del__ to call self.cancel() when not getattr(self, "_cancelled", True), and cleared self._wrapped = None when iteration completes naturally so finished or already-cancelled streams aren't cancelled again on GC.
…cancel in __del__ Require keyword argument syntax for arrow_serialization_options on BigQueryReadClient.read_rows() and ReadRowsStream.__init__(). Cancel unfinished streams automatically in ReadRowsStream.__del__(), and clear _wrapped upon normal stream completion so completed or already-cancelled streams are not cancelled again during garbage collection.
Replace arrow_serialization_options on ReadRowsStream.__init__ with a keyword-only request_template parameter so future ReadRowsRequest options can be passed without changing the stream constructor. Copy the template via ReadRowsRequest(self._request_template) on reconnect to avoid mutating caller-supplied request objects.
This branch has not been deployed
2 participants
Adds
arrow_serialization_optionstoBigQueryReadClient.read_rows()and aReadRowsStream.cancel()method, for reading Arrow results from a query job's default stream.ReadRowscall, including reconnects.cancel()stops the stream, including during a reconnect, without closing the client.LZ4_FRAME/ZSTDandcancel()on a default job stream.google-cloud-bigquery will depend on the release that includes this.