feat(bigquery-storage): add Arrow options and cancel to read_rows

packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py

@@ -57,6 +57,8 @@ 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 +126,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:
@@ -137,13 +147,19 @@ 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,
             name,
             offset,
             {"retry": retry, "timeout": timeout, "metadata": metadata},
             retry_delay_callback=retry_delay_callback,
+            request_template=request_template,
         )
         stream._reconnect()
         return stream

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
@@ -86,7 +88,14 @@ 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,
+        *,
+        request_template=None,
     ):
         """Construct a ReadRowsStream.
 
@@ -114,6 +123,13 @@ def __init__(
                 ReadRowsStream will call retry_delay_callback with the delay
                 duration (in seconds) before it starts sleeping until the next
                 attempt.
+            request_template (Optional[Union[dict, \
+                ~google.cloud.bigquery_storage_v1.types.ReadRowsRequest \
+            ]]):
+                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[ \
@@ -129,6 +145,8 @@ def __init__(
         self._offset = offset
         self._read_rows_kwargs = read_rows_kwargs
         self._retry_delay_callback = retry_delay_callback
+        self._request_template = request_template
+        self._cancelled = False
         self._wrapped = None
 
     def __iter__(self):
@@ -140,19 +158,20 @@ 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
                     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(
@@ -165,22 +184,60 @@ 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 __del__(self):
+        if not getattr(self, "_cancelled", True):
+            self.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._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 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=request,
+                        **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
 

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) == []

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"],