From ea74a338c0ab47b931d3865c1540451de10d95a0 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 1/3] feat(bigquery-storage): add Arrow options and cancel to read_rows 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. --- .../cloud/bigquery_storage_v1/client.py | 10 +++ .../cloud/bigquery_storage_v1/reader.py | 68 ++++++++++++--- .../tests/system/reader/test_reader.py | 64 ++++++++++++++ .../tests/unit/test_read_client_v1.py | 52 ++++++++++++ .../tests/unit/test_reader_v1.py | 83 +++++++++++++++++++ 5 files changed, 267 insertions(+), 10 deletions(-) diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py index 0decb4f413ba..556318ce440e 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py @@ -57,6 +57,7 @@ def read_rows( timeout=google.api_core.gapic_v1.method.DEFAULT, metadata=(), retry_delay_callback=None, + arrow_serialization_options=None, ): """ Reads rows from the table in the format prescribed by the read @@ -124,6 +125,14 @@ def read_rows( BigQueryReadClient will call retry_delay_callback with the delay duration (in seconds) before it starts sleeping until the next attempt. + arrow_serialization_options (Optional[Union[dict, \ + ~google.cloud.bigquery_storage_v1.types.ArrowSerializationOptions \ + ]]): + Arrow serialization options, such as buffer compression, to + send on every ``ReadRows`` request, including after a + reconnect. Only job default streams + (``projects/{project}/locations/{location}/jobs/{job}/streams/_default``) + accept them. Returns: ~google.cloud.bigquery_storage_v1.reader.ReadRowsStream: @@ -144,6 +153,7 @@ def read_rows( offset, {"retry": retry, "timeout": timeout, "metadata": metadata}, retry_delay_callback=retry_delay_callback, + arrow_serialization_options=arrow_serialization_options, ) stream._reconnect() return stream diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py index 27ab7f55de48..5f49187211e9 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py @@ -86,7 +86,13 @@ class ReadRowsStream(object): """ def __init__( - self, client, name, offset, read_rows_kwargs, retry_delay_callback=None + self, + client, + name, + offset, + read_rows_kwargs, + retry_delay_callback=None, + arrow_serialization_options=None, ): """Construct a ReadRowsStream. @@ -114,6 +120,11 @@ def __init__( ReadRowsStream will call retry_delay_callback with the delay duration (in seconds) before it starts sleeping until the next attempt. + arrow_serialization_options (Optional[Union[dict, \ + ~google.cloud.bigquery_storage_v1.types.ArrowSerializationOptions \ + ]]): + Arrow serialization options to send on every ``ReadRows`` + request, including after a reconnect. Returns: Iterable[ \ @@ -129,6 +140,8 @@ def __init__( self._offset = offset self._read_rows_kwargs = read_rows_kwargs self._retry_delay_callback = retry_delay_callback + self._arrow_serialization_options = arrow_serialization_options + self._cancelled = False self._wrapped = None def __iter__(self): @@ -140,13 +153,13 @@ def __iter__(self): ]: A sequence of row messages. """ - # Infinite loop to reconnect on reconnectable errors while processing - # the row stream. + # Loop to reconnect on reconnectable errors while processing the row + # stream, until the stream ends or is cancelled. if self._wrapped is None: self._reconnect() - while True: + while not self._cancelled: try: for message in self._wrapped: rowcount = message.row_count @@ -165,22 +178,57 @@ def __iter__(self): # Transient error, so reconnect to the stream. pass except Exception as exc: + # cancel() makes the call fail, which isn't an error here. + if self._cancelled: + return if not self._resource_exhausted_exception_is_retryable(exc): raise self._reconnect() + def cancel(self): + """Cancel the ``ReadRows`` call. + + Use this to stop reading before the end of the stream. After + ``cancel()``, iterating the stream yields no more messages, and the + stream doesn't reconnect. The client stays open for other calls. + """ + self._cancelled = True + cancel = getattr(self._wrapped, "cancel", None) + if cancel is not None: + cancel() + def _reconnect(self): """Reconnect to the ReadRows stream using the most recent offset.""" - while True: + while not self._cancelled: try: - self._wrapped = self._client.read_rows( - read_stream=self._name, - offset=self._offset, - **self._read_rows_kwargs, - ) + if self._arrow_serialization_options is None: + wrapped = self._client.read_rows( + read_stream=self._name, + offset=self._offset, + **self._read_rows_kwargs, + ) + else: + # The flattened read_rows() arguments don't include + # arrow_serialization_options, so send a full request. + wrapped = self._client.read_rows( + request={ + "read_stream": self._name, + "offset": self._offset, + "arrow_serialization_options": self._arrow_serialization_options, + }, + **self._read_rows_kwargs, + ) + # Store the new call before checking for cancel(), so a + # concurrent cancel() either cancels it or is seen here. + self._wrapped = wrapped + if self._cancelled: + self.cancel() + return break except Exception as exc: + if self._cancelled: + return if not self._resource_exhausted_exception_is_retryable(exc): raise diff --git a/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py b/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py index e9269bda9d96..a9279c89f05c 100644 --- a/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py +++ b/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py @@ -536,3 +536,67 @@ def test_read_rows_to_dataframe_with_wide_table(client_and_types, project_id): some_rows = next(pages_iter) assert all(len(row["tract_geom"].as_py()) > 0 for row in some_rows) + + +def _arrow_table(responses): + import pyarrow + + schema, batches = None, [] + for response in responses: + if schema is None and response.arrow_schema.serialized_schema: + schema = pyarrow.ipc.read_schema( + pyarrow.py_buffer(response.arrow_schema.serialized_schema) + ) + if response.arrow_record_batch.serialized_record_batch: + batches.append( + pyarrow.ipc.read_record_batch( + pyarrow.py_buffer( + response.arrow_record_batch.serialized_record_batch + ), + schema, + ) + ) + return pyarrow.Table.from_batches(batches, schema=schema) + + +@pytest.fixture(scope="module") +def job_default_stream(bq_client): + job = bq_client.query( + "SELECT x, REPEAT('A', 1000) AS repeated_str " + "FROM UNNEST(GENERATE_ARRAY(1, 20000)) AS x ORDER BY x" + ) + job.result() + return ( + f"projects/{job.project}/locations/{job.location}" + f"/jobs/{job.job_id}/streams/_default" + ) + + +@pytest.mark.parametrize("codec", ["LZ4_FRAME", "ZSTD"]) +def test_read_rows_job_stream_with_arrow_serialization_options( + client_v1, job_default_stream, codec +): + plain = list(client_v1.read_rows(job_default_stream)) + compressed = list( + client_v1.read_rows( + job_default_stream, + arrow_serialization_options={"buffer_compression": codec}, + ) + ) + + def size(responses): + return sum(len(r.arrow_record_batch.serialized_record_batch) for r in responses) + + assert size(compressed) < size(plain) + assert _arrow_table(compressed).equals(_arrow_table(plain)) + + +def test_read_rows_job_stream_cancel(client_v1, job_default_stream): + stream = client_v1.read_rows(job_default_stream) + messages = iter(stream) + first = next(messages) + + stream.cancel() + + assert first.row_count > 0 + assert list(messages) == [] diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py b/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py index abb88ba36edc..bca9647bdffd 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py @@ -153,6 +153,58 @@ def test_read_rows(mock_transport, client_under_test): ) +def test_read_rows_with_arrow_serialization_options(mock_transport, client_under_test): + rpc_callable = mock.Mock() + mock_transport._wrapped_methods[mock_transport.read_rows] = rpc_callable + options = {"buffer_compression": "LZ4_FRAME"} + + client_under_test.read_rows( + "teststream", offset=10, arrow_serialization_options=options + ) + + rpc_callable.assert_called_once_with( + types.ReadRowsRequest( + read_stream="teststream", offset=10, arrow_serialization_options=options + ), + metadata=mock.ANY, + retry=mock.ANY, + timeout=mock.ANY, + ) + + +def test_read_rows_reconnect_resends_arrow_serialization_options( + mock_transport, client_under_test +): + def first_stream(): + yield types.ReadRowsResponse(row_count=5) + raise google.api_core.exceptions.ServiceUnavailable("transient disconnect") + + rpc_callable = mock.Mock( + side_effect=[first_stream(), iter([types.ReadRowsResponse(row_count=3)])] + ) + mock_transport._wrapped_methods[mock_transport.read_rows] = rpc_callable + options = {"buffer_compression": "ZSTD"} + + stream = client_under_test.read_rows( + "teststream", offset=10, arrow_serialization_options=options + ) + + assert [m.row_count for m in stream] == [5, 3] + assert rpc_callable.call_args_list == [ + mock.call( + types.ReadRowsRequest( + read_stream="teststream", + offset=offset, + arrow_serialization_options=options, + ), + metadata=mock.ANY, + retry=mock.ANY, + timeout=mock.ANY, + ) + for offset in (10, 15) + ] + + @pytest.mark.parametrize( "module_under_test", ["google.cloud.bigquery_storage_v1", "google.cloud.bigquery_storage_v1beta2"], diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py index d5a83b2a7cf7..6bc4115f9ce4 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py @@ -717,3 +717,86 @@ def test_to_dataframe_emits_pending_deprecation_warning(mut): assert actual_df == mock_df mock_parser.to_dataframe.assert_called_once_with(mock_message, dtypes=None) + + +def test_cancel_stops_iteration_without_reconnecting( + class_under_test, mock_gapic_client +): + first = types.ReadRowsResponse(row_count=1) + + def pages(): + yield first + raise google.api_core.exceptions.ServiceUnavailable("connection closed") + + wrapped = mock.MagicMock() + wrapped.__iter__.return_value = pages() + mock_gapic_client.read_rows.return_value = wrapped + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + messages = iter(reader) + assert next(messages) is first + reader.cancel() + + assert list(messages) == [] + wrapped.cancel.assert_called_once_with() + mock_gapic_client.read_rows.assert_called_once_with( + read_stream="teststream", offset=0 + ) + + +def test_cancel_ignores_the_cancelled_call_error(class_under_test, mock_gapic_client): + def pages(): + yield types.ReadRowsResponse(row_count=1) + raise google.api_core.exceptions.Cancelled("call cancelled") + + mock_gapic_client.read_rows.return_value = pages() + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + messages = iter(reader) + next(messages) + reader.cancel() + + assert list(messages) == [] + + +def test_cancel_before_iterating_yields_nothing(class_under_test, mock_gapic_client): + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + reader.cancel() + + assert list(reader) == [] + mock_gapic_client.read_rows.assert_not_called() + + +def test_cancel_during_reconnect_cancels_new_stream( + class_under_test, mock_gapic_client +): + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + wrapped = mock.MagicMock() + + def cancel_while_opening(**kwargs): + reader.cancel() + return wrapped + + mock_gapic_client.read_rows.side_effect = cancel_while_opening + + assert list(reader) == [] + wrapped.cancel.assert_called_once_with() + assert reader._wrapped is wrapped + + +def test_cancel_during_reconnect_retry_stops_without_retrying( + class_under_test, mock_gapic_client +): + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + def cancel_and_fail(**kwargs): + reader.cancel() + raise google.api_core.exceptions.Cancelled("call cancelled") + + mock_gapic_client.read_rows.side_effect = cancel_and_fail + + assert list(reader) == [] + mock_gapic_client.read_rows.assert_called_once_with( + read_stream="teststream", offset=0 + ) From d31e1dafa37ca33c0a49a4b9b2f7eaae287f7c84 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 20:55:19 +0000 Subject: [PATCH 2/3] feat(bigquery-storage): make arrow_serialization_options kw-only and 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. --- .../cloud/bigquery_storage_v1/client.py | 1 + .../cloud/bigquery_storage_v1/reader.py | 6 ++++ .../tests/unit/test_reader_v1.py | 32 +++++++++++++++++++ 3 files changed, 39 insertions(+) diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py index 556318ce440e..990c715d4a3e 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py @@ -57,6 +57,7 @@ def read_rows( timeout=google.api_core.gapic_v1.method.DEFAULT, metadata=(), retry_delay_callback=None, + *, arrow_serialization_options=None, ): """ diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py index 5f49187211e9..3c01b924910d 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py @@ -92,6 +92,7 @@ def __init__( offset, read_rows_kwargs, retry_delay_callback=None, + *, arrow_serialization_options=None, ): """Construct a ReadRowsStream. @@ -166,6 +167,7 @@ def __iter__(self): self._offset += rowcount yield message + self._wrapped = None return # Made it through the whole stream. except google.api_core.exceptions.InternalServerError as exc: resumable_error = any( @@ -198,6 +200,10 @@ def cancel(self): if cancel is not None: cancel() + def __del__(self): + if not getattr(self, "_cancelled", True): + self.cancel() + def _reconnect(self): """Reconnect to the ReadRows stream using the most recent offset.""" while not self._cancelled: diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py index 6bc4115f9ce4..c3f40d6d4c6d 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py @@ -800,3 +800,35 @@ def cancel_and_fail(**kwargs): mock_gapic_client.read_rows.assert_called_once_with( read_stream="teststream", offset=0 ) + + +def test_del_cancels_unfinished_stream_once(class_under_test, mock_gapic_client): + wrapped = mock.MagicMock() + mock_gapic_client.read_rows.return_value = wrapped + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + reader._reconnect() + + del reader + + wrapped.cancel.assert_called_once_with() + + +def test_del_does_not_cancel_twice_or_after_completion( + class_under_test, mock_gapic_client +): + cancelled = mock.MagicMock() + completed = mock.MagicMock() + completed.__iter__.return_value = iter([types.ReadRowsResponse(row_count=1)]) + mock_gapic_client.read_rows.side_effect = [cancelled, completed] + + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + reader._reconnect() + reader.cancel() + del reader + + finished_reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + assert len(list(finished_reader)) == 1 + del finished_reader + + cancelled.cancel.assert_called_once_with() + completed.cancel.assert_not_called() From 9d7cde017f7c81f139bd73684b4fe29e115fcfd3 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 21:22:28 +0000 Subject: [PATCH 3/3] refactor(bigquery-storage): use request_template in ReadRowsStream 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. --- .../cloud/bigquery_storage_v1/client.py | 7 +++- .../cloud/bigquery_storage_v1/reader.py | 31 +++++++++-------- .../tests/unit/test_reader_v1.py | 33 +++++++++++++++++++ 3 files changed, 56 insertions(+), 15 deletions(-) diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py index 990c715d4a3e..791be3f2a0dc 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py @@ -147,6 +147,11 @@ def read_rows( to a retryable error and retry attempts failed. ValueError: If the parameters are invalid. """ + request_template = None + if arrow_serialization_options is not None: + request_template = { + "arrow_serialization_options": arrow_serialization_options + } gapic_client = super(BigQueryReadClient, self) stream = reader.ReadRowsStream( gapic_client, @@ -154,7 +159,7 @@ def read_rows( offset, {"retry": retry, "timeout": timeout, "metadata": metadata}, retry_delay_callback=retry_delay_callback, - arrow_serialization_options=arrow_serialization_options, + request_template=request_template, ) stream._reconnect() return stream diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py index 3c01b924910d..ade01abb67b8 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py @@ -27,6 +27,8 @@ import google.api_core.exceptions import google.rpc.error_details_pb2 # type: ignore +from google.cloud.bigquery_storage_v1 import types as gapic_types + try: import pandas except ImportError: # pragma: NO COVER @@ -93,7 +95,7 @@ def __init__( read_rows_kwargs, retry_delay_callback=None, *, - arrow_serialization_options=None, + request_template=None, ): """Construct a ReadRowsStream. @@ -121,11 +123,13 @@ def __init__( ReadRowsStream will call retry_delay_callback with the delay duration (in seconds) before it starts sleeping until the next attempt. - arrow_serialization_options (Optional[Union[dict, \ - ~google.cloud.bigquery_storage_v1.types.ArrowSerializationOptions \ + request_template (Optional[Union[dict, \ + ~google.cloud.bigquery_storage_v1.types.ReadRowsRequest \ ]]): - Arrow serialization options to send on every ``ReadRows`` - request, including after a reconnect. + Request fields, such as ``arrow_serialization_options``, to + send on every ``ReadRows`` request, including after a + reconnect. ``read_stream`` and ``offset`` are set from + ``name`` and the current stream offset. Returns: Iterable[ \ @@ -141,7 +145,7 @@ def __init__( self._offset = offset self._read_rows_kwargs = read_rows_kwargs self._retry_delay_callback = retry_delay_callback - self._arrow_serialization_options = arrow_serialization_options + self._request_template = request_template self._cancelled = False self._wrapped = None @@ -208,21 +212,20 @@ def _reconnect(self): """Reconnect to the ReadRows stream using the most recent offset.""" while not self._cancelled: try: - if self._arrow_serialization_options is None: + if self._request_template is None: wrapped = self._client.read_rows( read_stream=self._name, offset=self._offset, **self._read_rows_kwargs, ) else: - # The flattened read_rows() arguments don't include - # arrow_serialization_options, so send a full request. + # The flattened read_rows() arguments only include + # read_stream and offset, so send a full request copy. + request = gapic_types.ReadRowsRequest(self._request_template) + request.read_stream = self._name + request.offset = self._offset wrapped = self._client.read_rows( - request={ - "read_stream": self._name, - "offset": self._offset, - "arrow_serialization_options": self._arrow_serialization_options, - }, + request=request, **self._read_rows_kwargs, ) # Store the new call before checking for cancel(), so a diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py index c3f40d6d4c6d..ed8220fe177f 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py @@ -832,3 +832,36 @@ def test_del_does_not_cancel_twice_or_after_completion( cancelled.cancel.assert_called_once_with() completed.cancel.assert_not_called() + + +def test_request_template_not_mutated_across_reconnects( + class_under_test, mock_gapic_client +): + template = types.ReadRowsRequest( + arrow_serialization_options={"buffer_compression": "LZ4_FRAME"} + ) + mock_gapic_client.read_rows.side_effect = [ + _pages_w_unavailable([types.ReadRowsResponse(row_count=4)]), + iter([types.ReadRowsResponse(row_count=2)]), + ] + reader = class_under_test( + mock_gapic_client, + "teststream", + 10, + {}, + request_template=template, + ) + + assert [m.row_count for m in reader] == [4, 2] + assert template.read_stream == "" + assert template.offset == 0 + assert mock_gapic_client.read_rows.call_args_list == [ + mock.call( + request=types.ReadRowsRequest( + read_stream="teststream", + offset=offset, + arrow_serialization_options={"buffer_compression": "LZ4_FRAME"}, + ) + ) + for offset in (10, 14) + ]