test(bigquery): run CI on the Arrow query results stack (do not merge)

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

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
 

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

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
+    )