From ea74a338c0ab47b931d3865c1540451de10d95a0 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 01/10] 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 1851fa95e254ff4fafee304e741ded7aded55aab Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 02/10] refactor(bigquery): move Arrow query result reading to _arrow_query RowIterator._download_arrow_from_job_id() holds the whole ARROW read path inline in table.py, and upcoming Arrow work (serialization options, row limits, explicit destinations, empty results) would grow it further. _pyarrow_helpers isn't an option, because it's deprecated in favor of pandas-gbq. Move the method body unchanged into a new private module, _arrow_query.download_arrow_from_job(), with self renamed to row_iterator. The RowIterator method keeps its signature and docstring and delegates to it. Tests that mocked table.pyarrow for this path now mock _arrow_query.pyarrow. No behavior change. --- .../google/cloud/bigquery/_arrow_query.py | 125 ++++++++++++++++++ .../google/cloud/bigquery/table.py | 83 +----------- .../unit/test_query_results_format_arrow.py | 14 +- 3 files changed, 136 insertions(+), 86 deletions(-) create mode 100644 packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py new file mode 100644 index 000000000000..25738a2e818d --- /dev/null +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -0,0 +1,125 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Read query results returned in Apache Arrow format.""" + +from __future__ import annotations + +import base64 +from typing import Iterator, Optional, TYPE_CHECKING + +try: + import pyarrow # type: ignore +except ImportError: + pyarrow = None # type: ignore[assignment] + +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud import bigquery_storage + from google.cloud.bigquery.table import RowIterator + +_NO_PYARROW_ERROR = ( + "The pyarrow library is not installed, please install " + "pyarrow to use the to_arrow() function." +) + + +def download_arrow_from_job( + row_iterator: "RowIterator", + bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None, + timeout: Optional[float] = None, +) -> Iterator["pyarrow.RecordBatch"]: + """Yield Arrow record batches for ``row_iterator``'s ARROW query results. + + See :meth:`google.cloud.bigquery.table.RowIterator._download_arrow_from_job_id`. + """ + if pyarrow is None: + raise ValueError(_NO_PYARROW_ERROR) + + # Step 1: Ensure BigQuery Read API client is available upfront. + if bqstorage_client is None: + if row_iterator.client is None: + raise ValueError("RowIterator client is None.") + bqstorage_client = row_iterator.client._ensure_bqstorage_client() + if bqstorage_client is None: + raise ValueError( + "The google-cloud-bigquery-storage library is required to read Arrow results." + ) + + offset = 0 + pa_schema = None + total_rows = row_iterator.total_rows + job_complete = False + + # Step 2: Process initial inline Arrow response from jobs.query if available. + if row_iterator._first_page_response: + first_page = row_iterator._first_page_response + row_iterator._first_page_response = None + + job_complete = bool(first_page.get("jobComplete", False)) + if job_complete: + total_rows = int(first_page.get("totalRows", 0)) + + arrow_schema_json = first_page.get("arrowSchema") + if isinstance(arrow_schema_json, dict): + schema_bytes = arrow_schema_json.get("serializedSchema") + if schema_bytes: + if isinstance(schema_bytes, str): + schema_bytes = base64.b64decode(schema_bytes) + pa_schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes)) + + arrow_batch_json = first_page.get("arrowRecordBatch") + if isinstance(arrow_batch_json, dict) and pa_schema is not None: + batch_bytes = arrow_batch_json.get("serializedRecordBatch") + if batch_bytes: + if isinstance(batch_bytes, str): + batch_bytes = base64.b64decode(batch_bytes) + batch = pyarrow.ipc.read_record_batch( + pyarrow.py_buffer(batch_bytes), + pa_schema, + ) + offset += batch.num_rows + yield batch + + # Step 3: Return early if all results were delivered in the first page response. + if job_complete and offset >= total_rows: + return + + # Step 4: Stream remaining Arrow record batches from the job default stream. + project = row_iterator._project or ( + row_iterator.client.project if row_iterator.client else None + ) + location = row_iterator._location or ( + row_iterator.client.location if row_iterator.client else None + ) + stream_name = f"projects/{project}/locations/{location}/jobs/{row_iterator._job_id}/streams/_default" + reader = bqstorage_client.read_rows(stream_name, offset=offset, timeout=timeout) + for response in reader: + if ( + response.arrow_schema + and response.arrow_schema.serialized_schema + and pa_schema is None + ): + pa_schema = pyarrow.ipc.read_schema( + pyarrow.py_buffer(response.arrow_schema.serialized_schema) + ) + if ( + response.arrow_record_batch + and response.arrow_record_batch.serialized_record_batch + and pa_schema is not None + ): + batch = pyarrow.ipc.read_record_batch( + pyarrow.py_buffer(response.arrow_record_batch.serialized_record_batch), + pa_schema, + ) + yield batch diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index d1d31b14ffd8..14f996b894f2 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -16,7 +16,6 @@ from __future__ import absolute_import -import base64 import copy import datetime import functools @@ -59,6 +58,7 @@ import google.cloud._helpers # type: ignore from google.api_core.page_iterator import HTTPIterator from google.cloud.bigquery import ( + _arrow_query, _helpers, _pandas_helpers, _string_references, @@ -2325,84 +2325,9 @@ def _download_arrow_from_job_id( Yields: pyarrow.RecordBatch: Record batches generated from the query result. """ - if pyarrow is None: - raise ValueError(_NO_PYARROW_ERROR) - - # Step 1: Ensure BigQuery Read API client is available upfront. - if bqstorage_client is None: - if self.client is None: - raise ValueError("RowIterator client is None.") - bqstorage_client = self.client._ensure_bqstorage_client() - if bqstorage_client is None: - raise ValueError( - "The google-cloud-bigquery-storage library is required to read Arrow results." - ) - - offset = 0 - pa_schema = None - total_rows = self.total_rows - job_complete = False - - # Step 2: Process initial inline Arrow response from jobs.query if available. - if self._first_page_response: - first_page = self._first_page_response - self._first_page_response = None - - job_complete = bool(first_page.get("jobComplete", False)) - if job_complete: - total_rows = int(first_page.get("totalRows", 0)) - - arrow_schema_json = first_page.get("arrowSchema") - if isinstance(arrow_schema_json, dict): - schema_bytes = arrow_schema_json.get("serializedSchema") - if schema_bytes: - if isinstance(schema_bytes, str): - schema_bytes = base64.b64decode(schema_bytes) - pa_schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes)) - - arrow_batch_json = first_page.get("arrowRecordBatch") - if isinstance(arrow_batch_json, dict) and pa_schema is not None: - batch_bytes = arrow_batch_json.get("serializedRecordBatch") - if batch_bytes: - if isinstance(batch_bytes, str): - batch_bytes = base64.b64decode(batch_bytes) - batch = pyarrow.ipc.read_record_batch( - pyarrow.py_buffer(batch_bytes), - pa_schema, - ) - offset += batch.num_rows - yield batch - - # Step 3: Return early if all results were delivered in the first page response. - if job_complete and offset >= total_rows: - return - - # Step 4: Stream remaining Arrow record batches from the job default stream. - project = self._project or (self.client.project if self.client else None) - location = self._location or (self.client.location if self.client else None) - stream_name = f"projects/{project}/locations/{location}/jobs/{self._job_id}/streams/_default" - reader = bqstorage_client.read_rows(stream_name, offset=offset, timeout=timeout) - for response in reader: - if ( - response.arrow_schema - and response.arrow_schema.serialized_schema - and pa_schema is None - ): - pa_schema = pyarrow.ipc.read_schema( - pyarrow.py_buffer(response.arrow_schema.serialized_schema) - ) - if ( - response.arrow_record_batch - and response.arrow_record_batch.serialized_record_batch - and pa_schema is not None - ): - batch = pyarrow.ipc.read_record_batch( - pyarrow.py_buffer( - response.arrow_record_batch.serialized_record_batch - ), - pa_schema, - ) - yield batch + return _arrow_query.download_arrow_from_job( + self, bqstorage_client=bqstorage_client, timeout=timeout + ) # If changing the signature of this method, make sure to apply the same # changes to job.QueryJob.to_arrow() diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index 3bc19b4d3ccd..1e22381419c7 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -193,7 +193,7 @@ def test_download_arrow_from_job_id_constructs_stream_and_reads(self): mock_response.arrow_record_batch = None mock_bqstorage.read_rows.return_value = [mock_response] - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, []) @@ -247,7 +247,7 @@ def test_download_arrow_from_job_id_with_first_page_response(self): mock_bqstorage.read_rows.return_value = [mock_response] - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema" mock_pyarrow.ipc.read_record_batch.side_effect = [ @@ -306,7 +306,7 @@ def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): mock_bqstorage.read_rows.return_value = [mock_response] - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema" mock_pyarrow.ipc.read_record_batch.return_value = mock_stream_batch @@ -339,7 +339,7 @@ def test_download_arrow_from_job_id_missing_storage_client_raises_value_error(se query_results_format="ARROW", ) - with mock.patch("google.cloud.bigquery.table.pyarrow"): + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow"): with self.assertRaises(ValueError) as ctx: list(iterator._download_arrow_from_job_id()) self.assertIn( @@ -374,7 +374,7 @@ def test_download_arrow_from_job_id_with_schema_and_batch(self): mock_bqstorage.read_rows.return_value = [mock_response] - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x mock_pyarrow.ipc.read_schema.return_value = "fake_schema" mock_pyarrow.ipc.read_record_batch.return_value = "fake_batch" @@ -422,7 +422,7 @@ def test_download_arrow_from_job_id_avoids_read_rows_when_all_rows_present(self) mock_first_batch = mock.MagicMock() mock_first_batch.num_rows = 10 - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema" mock_pyarrow.ipc.read_record_batch.return_value = mock_first_batch @@ -464,7 +464,7 @@ def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self): mock_first_batch.num_rows = 10 mock_bqstorage.read_rows.return_value = [] - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema" mock_pyarrow.ipc.read_record_batch.return_value = mock_first_batch From 2f149614891db71b67b4ead02440e7cb8325f24f Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 03/10] feat(bigquery): validate Arrow queries before sending them query_and_wait() with query_results_format=ARROW only found a missing pyarrow or google-cloud-bigquery-storage when results were read, after the query had already run. A query could work while its results fit in the jobs.query response and then fail once they needed a stream. It also accepted WRITE_APPEND, whose results can't be read on their own, and compression_codec without ARROW, where it does nothing. Check all three before the request is built. compression_codec without ARROW and an ARROW query with WRITE_APPEND (including through default_query_job_config) raise ValueError. An ARROW query checks pyarrow and google-cloud-bigquery-storage with the existing version helpers, and also checks that BigQueryReadClient.read_rows() accepts arrow_serialization_options. That check looks for ReadRowsStream.cancel(), which shipped in the same change, rather than a version number, so it's correct for the unreleased in-repo build that the system tests install until the next storage release. --- .../google/cloud/bigquery/_arrow_query.py | 67 ++++++++ .../google/cloud/bigquery/_job_helpers.py | 11 ++ .../cloud/bigquery/_versions_helpers.py | 15 ++ .../google/cloud/bigquery/client.py | 7 + packages/google-cloud-bigquery/noxfile.py | 10 ++ .../tests/system/test_arrow.py | 27 +++ .../tests/unit/test__versions_helpers.py | 25 +++ .../unit/test_query_results_format_arrow.py | 157 ++++++++++++++++++ 8 files changed, 319 insertions(+) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index 25738a2e818d..618ff031e9b6 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -19,6 +19,10 @@ import base64 from typing import Iterator, Optional, TYPE_CHECKING +from google.cloud.bigquery import _versions_helpers +from google.cloud.bigquery import enums +from google.cloud.bigquery import exceptions + try: import pyarrow # type: ignore except ImportError: @@ -26,6 +30,7 @@ if TYPE_CHECKING: # pragma: NO COVER from google.cloud import bigquery_storage + from google.cloud.bigquery import job from google.cloud.bigquery.table import RowIterator _NO_PYARROW_ERROR = ( @@ -34,6 +39,68 @@ ) +def validate_query( + query_results_format: Optional[str], + job_config: Optional["job.QueryJobConfig"], + *, + compression_codec: Optional[str] = None, +) -> None: + """Raise before sending a query whose ``ARROW`` results can't be read. + + Args: + query_results_format: The ``query_results_format`` passed to + ``query_and_wait()``. + job_config: The job config, already merged with + ``Client.default_query_job_config``. + compression_codec: The ``compression_codec`` passed to + ``query_and_wait()``. + + Raises: + ValueError: + If ``compression_codec`` is set without ``ARROW``, or if an + ``ARROW`` query uses ``WRITE_APPEND``. + google.cloud.bigquery.exceptions.BigQueryError: + If pyarrow or google-cloud-bigquery-storage is missing or too old + to read ``ARROW`` results. + """ + if query_results_format != enums.QueryResultsFormat.ARROW: + if compression_codec is not None: + raise ValueError("compression_codec requires query_results_format=ARROW.") + return + + # Results that don't fit in the jobs.query response are read from the + # destination table, which with WRITE_APPEND also holds older rows. + if ( + job_config is not None + and job_config.write_disposition == enums.WriteDisposition.WRITE_APPEND + ): + raise ValueError( + "query_results_format=ARROW doesn't support WRITE_APPEND, because " + "the results would also include the rows already in the " + "destination table." + ) + + # Check before the query runs, so a query doesn't succeed while its + # results fit in the first page and then fail once they need a stream. + _check_dependencies() + + +def _check_dependencies() -> None: + """Raise if pyarrow or google-cloud-bigquery-storage can't read ``ARROW`` + results.""" + _versions_helpers.PYARROW_VERSIONS.try_import(raise_if_error=True) + _versions_helpers.BQ_STORAGE_VERSIONS.try_import(raise_if_error=True) + if not _versions_helpers.BQ_STORAGE_VERSIONS.supports_arrow_serialization_options: + raise exceptions.LegacyBigQueryStorageError( + "query_results_format=ARROW needs a google-cloud-bigquery-storage " + "release whose BigQueryReadClient.read_rows() accepts " + "arrow_serialization_options, please upgrade it " + "(`pip install --upgrade google-cloud-bigquery-storage`) " + "(version found: " + f"{_versions_helpers.BQ_STORAGE_VERSIONS.installed_version})." + ) + + def download_arrow_from_job( row_iterator: "RowIterator", bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None, diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py index a7655f4bddf7..a952eff60f8c 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -49,6 +49,7 @@ import google.api_core.exceptions as core_exceptions from google.api_core import retry as retries +from google.cloud.bigquery import _arrow_query from google.cloud.bigquery import enums from google.cloud.bigquery import job import google.cloud.bigquery.job.query @@ -501,7 +502,17 @@ def query_and_wait( If ``job_config`` is not an instance of :class:`~google.cloud.bigquery.job.QueryJobConfig` class. + ValueError: + If ``compression_codec`` is set without ``ARROW``, or if an + ``ARROW`` query uses ``WRITE_APPEND``. + google.cloud.bigquery.exceptions.BigQueryError: + If pyarrow or google-cloud-bigquery-storage is missing or too old + to read ``ARROW`` results. The query isn't sent. """ + _arrow_query.validate_query( + query_results_format, job_config, compression_codec=compression_codec + ) + request_body = _to_query_request( query=query, job_config=job_config, location=location, timeout=api_timeout ) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py index d9afc9781532..1dc1848957fb 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py @@ -128,6 +128,21 @@ def is_read_session_optional(self) -> bool: """ return self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION + @property + def supports_arrow_serialization_options(self) -> bool: + """True if ``BigQueryReadClient.read_rows()`` accepts + ``arrow_serialization_options``, which ``ARROW`` query results need. + + That option was added in the same release as + ``ReadRowsStream.cancel()``, which ``ARROW`` query results also use, so + this checks for ``cancel()`` instead of comparing version numbers. That + also gives the right answer for a build of + google-cloud-bigquery-storage that hasn't been released yet. + """ + from google.cloud.bigquery_storage_v1 import reader + + return hasattr(reader.ReadRowsStream, "cancel") + def try_import(self, raise_if_error: bool = False) -> Any: """Tries to import the bigquery_storage module, and returns results accordingly. It also verifies the module version is recent enough. diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index e1cb5e1bdcdf..1322a7538a40 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -3746,6 +3746,13 @@ def query_and_wait( If ``job_config`` is not an instance of :class:`~google.cloud.bigquery.job.QueryJobConfig` class. + ValueError: + If ``compression_codec`` is set without ``ARROW``, or if an + ``ARROW`` query uses ``WRITE_APPEND``, including through + ``default_query_job_config``. + google.cloud.bigquery.exceptions.BigQueryError: + If pyarrow or google-cloud-bigquery-storage is missing or too + old to read ``ARROW`` results. The query isn't sent. """ return self._query_and_wait_bigframes( query, diff --git a/packages/google-cloud-bigquery/noxfile.py b/packages/google-cloud-bigquery/noxfile.py index af201dfb8a05..2cd266371c7c 100644 --- a/packages/google-cloud-bigquery/noxfile.py +++ b/packages/google-cloud-bigquery/noxfile.py @@ -288,6 +288,16 @@ def system(session): extras = "[all]" session.install("-e", f".{extras}", "-c", constraints_path) + # TODO: Remove once google-cloud-bigquery requires a + # google-cloud-bigquery-storage release whose read_rows() accepts + # arrow_serialization_options. Until then, query_results_format=ARROW + # needs the in-repo version. + session.install( + "-e", + str(CURRENT_DIRECTORY.parent / "google-cloud-bigquery-storage"), + "--no-deps", + ) + # Test with some broken "extras" in case the user didn't install the extra # directly. For example, pandas-gbq is recommended for pandas features, but # we want to test that we fallback to the previous behavior. For context, diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index 9db96503b00b..f9d4ade03f18 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -274,3 +274,30 @@ def test_query_and_wait_arrow_format_to_dataframe_iterable(bigquery_client): concat_df = pandas.concat(dfs, ignore_index=True) assert concat_df.columns.tolist() == ["num", "text"] assert concat_df.to_dict(orient="records") == [{"num": 200, "text": "python"}] + + +def test_query_and_wait_arrow_rejects_write_append(bigquery_client, dataset_id): + table_id = f"{bigquery_client.project}.{dataset_id}.arrow_write_append" + bigquery_client.query_and_wait( + f"CREATE OR REPLACE TABLE `{table_id}` AS SELECT 1 AS x" + ) + append = bigquery.QueryJobConfig( + destination=table_id, + write_disposition=enums.WriteDisposition.WRITE_APPEND, + ) + + with pytest.raises(ValueError, match="WRITE_APPEND"): + bigquery_client.query_and_wait( + "SELECT 2 AS x", + job_config=append, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + with bigquery.Client(default_query_job_config=append) as client: + with pytest.raises(ValueError, match="WRITE_APPEND"): + client.query_and_wait( + "SELECT 2 AS x", query_results_format=enums.QueryResultsFormat.ARROW + ) + + # Neither query ran, so the table still holds only its original row. + rows = bigquery_client.query_and_wait(f"SELECT x FROM `{table_id}`") + assert [row.x for row in rows] == [1] diff --git a/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py b/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py index aec553bb7b19..2dc8ff8e0a93 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py @@ -200,6 +200,31 @@ def test_bqstorage_is_read_session_optional_false(): assert not bqstorage_versions.is_read_session_optional +@pytest.mark.skipif( + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" +) +def test_bqstorage_supports_arrow_serialization_options_true(): + from google.cloud.bigquery_storage_v1 import reader + + bqstorage_versions = _versions_helpers.BQStorageVersions() + with mock.patch.object(reader.ReadRowsStream, "cancel", create=True): + assert bqstorage_versions.supports_arrow_serialization_options + + +@pytest.mark.skipif( + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" +) +def test_bqstorage_supports_arrow_serialization_options_false(): + from google.cloud.bigquery_storage_v1 import reader + + class ReadRowsStream: + """Stands in for a ReadRowsStream from before cancel() was added.""" + + bqstorage_versions = _versions_helpers.BQStorageVersions() + with mock.patch.object(reader, "ReadRowsStream", ReadRowsStream): + assert not bqstorage_versions.supports_arrow_serialization_options + + @pytest.mark.skipif(pandas is None, reason="pandas is not installed") @pytest.mark.parametrize("version", ["1.1.5", "2.0.0", "2.1.0"]) def test_try_import_raises_no_error_w_recent_pandas(version): diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index 1e22381419c7..fda8477ad3c8 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -16,12 +16,30 @@ import unittest from unittest import mock +import packaging.version import pytest +from google.cloud.bigquery import _arrow_query from google.cloud.bigquery import _job_helpers +from google.cloud.bigquery import _versions_helpers +from google.cloud.bigquery import exceptions from google.cloud.bigquery.client import Client +from google.cloud.bigquery.enums import WriteDisposition +from google.cloud.bigquery.job import QueryJobConfig from google.cloud.bigquery.table import RowIterator, _EmptyRowIterator +from .helpers import make_client + +# The real check, saved before the autouse fixture below replaces it. +_check_dependencies = _arrow_query._check_dependencies + + +@pytest.fixture(autouse=True) +def check_dependencies(): + """Skip the ARROW dependency check, which depends on installed packages.""" + with mock.patch.object(_arrow_query, "_check_dependencies") as check: + yield check + class TestQueryResultsFormatOption1(unittest.TestCase): def test_supported_by_jobs_query_includes_query_results_format(self): @@ -569,5 +587,144 @@ def test_row_iterator_to_dataframe_when_format_is_arrow(self): self.assertEqual(df, "full_df") +def _query_and_wait(client, **kwargs): + return _job_helpers.query_and_wait( + client, + "SELECT 1", + project="p", + location="us", + retry=None, + job_retry=None, + **{"job_config": None, **kwargs}, + ) + + +def test_query_and_wait_compression_codec_requires_arrow(check_dependencies): + client = mock.MagicMock(spec=Client) + + with pytest.raises(ValueError, match="compression_codec requires"): + _query_and_wait(client, compression_codec="LZ4_FRAME") + + client._call_api.assert_not_called() + check_dependencies.assert_not_called() + + +def test_query_and_wait_without_arrow_skips_dependency_check(check_dependencies): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "rows": [], + "schema": {"fields": []}, + } + + _query_and_wait(client) + + check_dependencies.assert_not_called() + + +def test_query_and_wait_arrow_rejects_write_append(): + client = mock.MagicMock(spec=Client) + job_config = QueryJobConfig(write_disposition=WriteDisposition.WRITE_APPEND) + + with pytest.raises(ValueError, match="WRITE_APPEND"): + _query_and_wait(client, job_config=job_config, query_results_format="ARROW") + + client._call_api.assert_not_called() + + +def test_query_and_wait_arrow_rejects_default_write_append(): + client = make_client( + default_query_job_config=QueryJobConfig( + destination="p.d.t", write_disposition=WriteDisposition.WRITE_APPEND + ) + ) + client._call_api = mock.Mock() + + with pytest.raises(ValueError, match="WRITE_APPEND"): + client.query_and_wait("SELECT 1", query_results_format="ARROW") + + client._call_api.assert_not_called() + + +@pytest.mark.parametrize( + "write_disposition", + [None, WriteDisposition.WRITE_TRUNCATE, WriteDisposition.WRITE_EMPTY], +) +def test_validate_query_allows_other_write_dispositions( + check_dependencies, write_disposition +): + job_config = QueryJobConfig(write_disposition=write_disposition) + + _arrow_query.validate_query("ARROW", job_config, compression_codec="ZSTD") + + check_dependencies.assert_called_once_with() + + +def test_query_and_wait_arrow_checks_dependencies_before_sending(check_dependencies): + check_dependencies.side_effect = exceptions.BigQueryStorageNotFoundError("missing") + client = mock.MagicMock(spec=Client) + + with pytest.raises(exceptions.BigQueryStorageNotFoundError): + _query_and_wait(client, query_results_format="ARROW") + + client._call_api.assert_not_called() + + +def test_check_dependencies_requires_pyarrow(): + error = exceptions.LegacyPyarrowError("pyarrow is not installed") + with mock.patch.object( + _versions_helpers.PYARROW_VERSIONS, "try_import", side_effect=error + ) as try_import: + with pytest.raises(exceptions.LegacyPyarrowError): + _check_dependencies() + + try_import.assert_called_once_with(raise_if_error=True) + + +def test_check_dependencies_requires_bigquery_storage(): + error = exceptions.BigQueryStorageNotFoundError("not installed") + with ( + mock.patch.object(_versions_helpers.PYARROW_VERSIONS, "try_import"), + mock.patch.object( + _versions_helpers.BQ_STORAGE_VERSIONS, "try_import", side_effect=error + ) as try_import, + ): + with pytest.raises(exceptions.BigQueryStorageNotFoundError): + _check_dependencies() + + try_import.assert_called_once_with(raise_if_error=True) + + +@pytest.mark.parametrize("supported", [True, False]) +def test_check_dependencies_requires_read_rows_arrow_options(supported): + versions_type = type(_versions_helpers.BQ_STORAGE_VERSIONS) + with ( + mock.patch.object(_versions_helpers.PYARROW_VERSIONS, "try_import"), + mock.patch.object(_versions_helpers.BQ_STORAGE_VERSIONS, "try_import"), + mock.patch.object( + versions_type, + "supports_arrow_serialization_options", + new_callable=mock.PropertyMock, + return_value=supported, + ), + mock.patch.object( + versions_type, + "installed_version", + new_callable=mock.PropertyMock, + return_value=packaging.version.Version("2.42.0"), + ), + ): + if supported: + _check_dependencies() + else: + with pytest.raises( + exceptions.LegacyBigQueryStorageError, + match=r"pip install --upgrade google-cloud-bigquery-storage.*" + r"\(version found: 2\.42\.0\)", + ): + _check_dependencies() + + if __name__ == "__main__": unittest.main() From 6543176fb266de3c2ea11e032fbbf0f3cffc0a3d Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 04/10] fix(bigquery): send the same Arrow options on jobs.query and ReadRows query_and_wait() sent compression_codec inside formatOptions, which only has useInt64Timestamp and timestampOutputFormat, so jobs.query ignored it, and ReadRows never got it at all. The first page and the stream each used their own endpoint's defaults, so one result could mix encodings. Build one arrowSerializationOptions value for ARROW queries and send it at the top level of the jobs.query request, where it also joins the jobs.query allowlist. Options the caller didn't set are sent with their defaults (COMPRESSION_UNSPECIFIED, TIMESTAMP_PRECISION_MICROS), so the encoding never depends on server defaults. RowIterator keeps the value on both the jobs.query and the waited-job paths and sends it on read_rows(). jobs.insert never gets these fields, since they come from query_and_wait() arguments rather than QueryJobConfig. The mypy session now installs the in-repo google-cloud-bigquery-storage too, so the new read_rows() argument is type-checked. --- .../google/cloud/bigquery/_arrow_query.py | 35 ++- .../google/cloud/bigquery/_job_helpers.py | 17 +- .../google/cloud/bigquery/table.py | 5 + packages/google-cloud-bigquery/noxfile.py | 8 + .../tests/system/test_arrow.py | 36 +++ .../unit/test_query_results_format_arrow.py | 228 +++++++++++++++++- 6 files changed, 310 insertions(+), 19 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index 618ff031e9b6..1cd1089d7f8c 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -17,7 +17,7 @@ from __future__ import annotations import base64 -from typing import Iterator, Optional, TYPE_CHECKING +from typing import Dict, Iterator, Optional, TYPE_CHECKING from google.cloud.bigquery import _versions_helpers from google.cloud.bigquery import enums @@ -101,6 +101,30 @@ def _check_dependencies() -> None: ) +def serialization_options(compression_codec: Optional[str]) -> Dict[str, str]: + """Build the ``arrowSerializationOptions`` sent on ``jobs.query``. + + ``ReadRows`` gets the same options. Options the caller didn't set are sent + with their defaults, so every batch uses one encoding, whatever each + endpoint would default to. + """ + return { + "bufferCompression": getattr(compression_codec, "value", compression_codec) + or "COMPRESSION_UNSPECIFIED", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", + } + + +def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]: + """Convert ``jobs.query`` ``arrowSerializationOptions`` to the + ``ArrowSerializationOptions`` fields that ``ReadRows`` takes.""" + options = options or serialization_options(None) + return { + "buffer_compression": options["bufferCompression"], + "picos_timestamp_precision": options["picosTimestampPrecision"], + } + + def download_arrow_from_job( row_iterator: "RowIterator", bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None, @@ -170,7 +194,14 @@ def download_arrow_from_job( row_iterator.client.location if row_iterator.client else None ) stream_name = f"projects/{project}/locations/{location}/jobs/{row_iterator._job_id}/streams/_default" - reader = bqstorage_client.read_rows(stream_name, offset=offset, timeout=timeout) + reader = bqstorage_client.read_rows( + stream_name, + offset=offset, + timeout=timeout, + arrow_serialization_options=_read_rows_options( + row_iterator._arrow_serialization_options + ), + ) for response in reader: if ( response.arrow_schema diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py index a952eff60f8c..384df0cbe8ed 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -518,11 +518,12 @@ def query_and_wait( ) if query_results_format is not None: request_body["queryResultsFormat"] = query_results_format - if compression_codec is not None: - request_body.setdefault("formatOptions", {}) - request_body["formatOptions"]["arrowSerializationOptions"] = { - "bufferCompression": compression_codec - } + arrow_serialization_options = None + if query_results_format == enums.QueryResultsFormat.ARROW: + arrow_serialization_options = _arrow_query.serialization_options( + compression_codec + ) + request_body["arrowSerializationOptions"] = arrow_serialization_options # Some API parameters aren't supported by the jobs.query API. In these # cases, fallback to a jobs.insert call. @@ -547,6 +548,7 @@ def query_and_wait( page_size=page_size, max_results=max_results, query_results_format=query_results_format, + arrow_serialization_options=arrow_serialization_options, callback=callback, ) @@ -620,6 +622,7 @@ def do_query(): page_size=page_size, max_results=max_results, query_results_format=query_results_format, + arrow_serialization_options=arrow_serialization_options, callback=callback, ) @@ -660,6 +663,7 @@ def do_query(): started=query_results.started, ended=query_results.ended, query_results_format=query_results_format, + arrow_serialization_options=arrow_serialization_options, ) if job_retry is not None: @@ -701,6 +705,7 @@ def _supported_by_jobs_query(request_body: Dict[str, Any]) -> bool: "reservation", "maxSlots", "queryResultsFormat", + "arrowSerializationOptions", } unsupported_keys = request_keys - keys_allowlist @@ -716,6 +721,7 @@ def _wait_or_cancel( max_results: Optional[int], *, query_results_format: Optional[str] = None, + arrow_serialization_options: Optional[Dict[str, str]] = None, callback: Callable = lambda _: None, ) -> table.RowIterator: """Wait for a job to complete and return the results. @@ -761,6 +767,7 @@ def _wait_or_cancel( ) ) query_results._query_results_format = query_results_format + query_results._arrow_serialization_options = arrow_serialization_options return query_results except Exception: # Attempt to cancel the job since we can't return the results. diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index 14f996b894f2..56e7accd5485 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -1916,6 +1916,9 @@ class RowIterator(HTTPIterator): If representing query results, the start time of the associated query. ended (Optional[datetime.datetime]): If representing query results, the end time of the associated query. + arrow_serialization_options (Optional[Dict[str, str]]): + The ``arrowSerializationOptions`` sent on ``jobs.query`` for + ``ARROW`` results. ``ReadRows`` requests the same options. """ def __init__( @@ -1944,6 +1947,7 @@ def __init__( started: Optional[datetime.datetime] = None, ended: Optional[datetime.datetime] = None, query_results_format: Optional[str] = None, + arrow_serialization_options: Optional[Dict[str, str]] = None, ): super(RowIterator, self).__init__( client, @@ -1978,6 +1982,7 @@ def __init__( self._job_started = started self._job_ended = ended self._query_results_format = query_results_format + self._arrow_serialization_options = arrow_serialization_options @property @_disallow_when_arrow diff --git a/packages/google-cloud-bigquery/noxfile.py b/packages/google-cloud-bigquery/noxfile.py index 2cd266371c7c..5f62a16aeaae 100644 --- a/packages/google-cloud-bigquery/noxfile.py +++ b/packages/google-cloud-bigquery/noxfile.py @@ -222,6 +222,14 @@ def mypy(session): """Run type checks with mypy.""" session.install("-e", ".[all]") + # TODO: Remove once google-cloud-bigquery requires a + # google-cloud-bigquery-storage release whose read_rows() accepts + # arrow_serialization_options. Not editable, because mypy can't follow the + # import hook of an editable install. + session.install( + str(CURRENT_DIRECTORY.parent / "google-cloud-bigquery-storage"), + "--no-deps", + ) session.install(MYPY_VERSION) # Just install the dependencies' type info directly, since "mypy --install-types" diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index f9d4ade03f18..756f37cadf2f 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -301,3 +301,39 @@ def test_query_and_wait_arrow_rejects_write_append(bigquery_client, dataset_id): # Neither query ran, so the table still holds only its original row. rows = bigquery_client.query_and_wait(f"SELECT x FROM `{table_id}`") assert [row.x for row in rows] == [1] + + +@pytest.mark.parametrize( + "codec", + [ + enums.QueryResultsCompressionCodec.LZ4_FRAME, + enums.QueryResultsCompressionCodec.ZSTD, + ], +) +@pytest.mark.parametrize( + "page_size", [None, 100], ids=["first_page_only", "with_stream"] +) +def test_query_and_wait_arrow_compression_matches_uncompressed( + bigquery_client, codec, page_size +): + query = ( + "SELECT x, CAST(x AS STRING) AS s, x / 7 AS f, " + "TIMESTAMP_ADD(TIMESTAMP '2026-01-01', INTERVAL x SECOND) AS ts " + "FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=False) + + def read(**kwargs): + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + page_size=page_size, + query_results_format=enums.QueryResultsFormat.ARROW, + **kwargs, + ) + return pyarrow.Table.from_batches(list(rows.to_arrow_iterable())) + + compressed = read(compression_codec=codec) + + assert compressed.num_rows == 1000 + assert compressed.equals(read()) diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index fda8477ad3c8..8b58886bd9d1 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -24,12 +24,24 @@ from google.cloud.bigquery import _versions_helpers from google.cloud.bigquery import exceptions from google.cloud.bigquery.client import Client +from google.cloud.bigquery.enums import QueryResultsCompressionCodec from google.cloud.bigquery.enums import WriteDisposition from google.cloud.bigquery.job import QueryJobConfig from google.cloud.bigquery.table import RowIterator, _EmptyRowIterator from .helpers import make_client +# The Arrow serialization options that jobs.query and ReadRows get when the +# caller sets none. +_DEFAULT_QUERY_OPTIONS = { + "bufferCompression": "COMPRESSION_UNSPECIFIED", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", +} +_DEFAULT_READ_OPTIONS = { + "buffer_compression": "COMPRESSION_UNSPECIFIED", + "picos_timestamp_precision": "TIMESTAMP_PRECISION_MICROS", +} + # The real check, saved before the autouse fixture below replaces it. _check_dependencies = _arrow_query._check_dependencies @@ -93,12 +105,15 @@ def test_job_helpers_query_and_wait_sets_compression_codec(self): ) call_args = client._call_api.call_args - self.assertIn("formatOptions", call_args.kwargs["data"]) + self.assertNotIn( + "arrowSerializationOptions", call_args.kwargs["data"]["formatOptions"] + ) self.assertEqual( - call_args.kwargs["data"]["formatOptions"]["arrowSerializationOptions"][ - "bufferCompression" - ], - "LZ4_FRAME", + call_args.kwargs["data"]["arrowSerializationOptions"], + { + "bufferCompression": "LZ4_FRAME", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", + }, ) def test_job_helpers_query_and_wait_fallback_preserves_query_results_format(self): @@ -220,7 +235,10 @@ def test_download_arrow_from_job_id_constructs_stream_and_reads(self): "projects/test-proj/locations/US/jobs/test-job-456/streams/_default" ) mock_bqstorage.read_rows.assert_called_once_with( - expected_stream_name, offset=0, timeout=5.0 + expected_stream_name, + offset=0, + timeout=5.0, + arrow_serialization_options=_DEFAULT_READ_OPTIONS, ) def test_download_arrow_from_job_id_with_first_page_response(self): @@ -290,7 +308,10 @@ def test_download_arrow_from_job_id_with_first_page_response(self): "projects/test-proj/locations/US/jobs/test-job-sync/streams/_default" ) mock_bqstorage.read_rows.assert_called_once_with( - expected_stream_name, offset=10, timeout=5.0 + expected_stream_name, + offset=10, + timeout=5.0, + arrow_serialization_options=_DEFAULT_READ_OPTIONS, ) def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): @@ -339,7 +360,10 @@ def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): ) expected_stream_name = "projects/test-proj/locations/US/jobs/test-job-schema-only/streams/_default" mock_bqstorage.read_rows.assert_called_once_with( - expected_stream_name, offset=0, timeout=5.0 + expected_stream_name, + offset=0, + timeout=5.0, + arrow_serialization_options=_DEFAULT_READ_OPTIONS, ) def test_download_arrow_from_job_id_missing_storage_client_raises_value_error(self): @@ -492,7 +516,10 @@ def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self): mock_client._ensure_bqstorage_client.assert_called_once() expected_stream = "projects/test-proj/locations/US/jobs/test-job-incomplete/streams/_default" mock_bqstorage.read_rows.assert_called_once_with( - expected_stream, offset=10, timeout=5.0 + expected_stream, + offset=10, + timeout=5.0, + arrow_serialization_options=_DEFAULT_READ_OPTIONS, ) def test_enums_values(self): @@ -535,9 +562,7 @@ def test_job_helpers_query_and_wait_accepts_enums(self): call_args = client._call_api.call_args self.assertEqual(call_args.kwargs["data"]["queryResultsFormat"], "ARROW") self.assertEqual( - call_args.kwargs["data"]["formatOptions"]["arrowSerializationOptions"][ - "bufferCompression" - ], + call_args.kwargs["data"]["arrowSerializationOptions"]["bufferCompression"], "LZ4_FRAME", ) @@ -726,5 +751,184 @@ def test_check_dependencies_requires_read_rows_arrow_options(supported): _check_dependencies() +def test_query_and_wait_arrow_sends_default_serialization_options(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "schema": {"fields": []}, + } + + row_iterator = _query_and_wait(client, query_results_format="ARROW") + + data = client._call_api.call_args.kwargs["data"] + assert data["arrowSerializationOptions"] == _DEFAULT_QUERY_OPTIONS + assert "arrowSerializationOptions" not in data["formatOptions"] + assert row_iterator._arrow_serialization_options == _DEFAULT_QUERY_OPTIONS + + +def test_query_and_wait_without_arrow_sends_no_serialization_options(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "rows": [], + "schema": {"fields": []}, + } + + row_iterator = _query_and_wait(client) + + assert "arrowSerializationOptions" not in client._call_api.call_args.kwargs["data"] + assert row_iterator._arrow_serialization_options is None + + +def test_query_and_wait_arrow_waits_with_the_same_options(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": False, + } + + with ( + mock.patch.object(_job_helpers, "_to_query_job") as to_query_job, + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, + ): + _query_and_wait(client, query_results_format="ARROW", compression_codec="ZSTD") + + wait_or_cancel.assert_called_once_with( + to_query_job.return_value, + api_timeout=mock.ANY, + wait_timeout=mock.ANY, + retry=mock.ANY, + page_size=mock.ANY, + max_results=mock.ANY, + query_results_format="ARROW", + arrow_serialization_options={ + "bufferCompression": "ZSTD", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", + }, + callback=mock.ANY, + ) + + +def test_wait_or_cancel_sets_arrow_serialization_options(): + job = mock.MagicMock(dry_run=False) + job.result.return_value = RowIterator( + client=mock.MagicMock(), api_request=mock.MagicMock(), path=None, schema=() + ) + + rows = _job_helpers._wait_or_cancel( + job, + api_timeout=None, + wait_timeout=None, + retry=None, + page_size=None, + max_results=None, + query_results_format="ARROW", + arrow_serialization_options=_DEFAULT_QUERY_OPTIONS, + ) + + assert rows._query_results_format == "ARROW" + assert rows._arrow_serialization_options == _DEFAULT_QUERY_OPTIONS + + +def test_query_and_wait_arrow_jobs_insert_omits_arrow_fields(): + client = make_client() + client._call_api = mock.Mock( + return_value={ + "jobReference": {"projectId": "PROJECT", "jobId": "j", "location": "US"}, + "configuration": {"query": {"query": "SELECT 1"}}, + "status": {"state": "RUNNING"}, + } + ) + # jobs.query doesn't accept a destination table, so this uses jobs.insert. + job_config = QueryJobConfig(destination="PROJECT.dataset.table") + + with mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel: + client.query_and_wait( + "SELECT 1", job_config=job_config, query_results_format="ARROW" + ) + + (insert_call,) = client._call_api.call_args_list + assert insert_call.kwargs["path"] == "/projects/PROJECT/jobs" + query_config = insert_call.kwargs["data"]["configuration"]["query"] + assert "queryResultsFormat" not in query_config + assert "arrowSerializationOptions" not in query_config + wait_or_cancel.assert_called_once_with( + mock.ANY, + api_timeout=mock.ANY, + wait_timeout=mock.ANY, + retry=mock.ANY, + page_size=mock.ANY, + max_results=mock.ANY, + query_results_format="ARROW", + arrow_serialization_options=_DEFAULT_QUERY_OPTIONS, + callback=mock.ANY, + ) + + +def test_supported_by_jobs_query_includes_arrow_serialization_options(): + body = { + "query": "SELECT 1", + "queryResultsFormat": "ARROW", + "arrowSerializationOptions": _DEFAULT_QUERY_OPTIONS, + } + + assert _job_helpers._supported_by_jobs_query(body) + + +@pytest.mark.parametrize("codec", ["LZ4_FRAME", QueryResultsCompressionCodec.LZ4_FRAME]) +def test_serialization_options_accepts_str_and_enum(codec): + assert _arrow_query.serialization_options(codec) == { + "bufferCompression": "LZ4_FRAME", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", + } + + +def test_read_rows_options(): + assert _arrow_query._read_rows_options(None) == _DEFAULT_READ_OPTIONS + assert _arrow_query._read_rows_options( + { + "bufferCompression": "ZSTD", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_PICOS", + } + ) == { + "buffer_compression": "ZSTD", + "picos_timestamp_precision": "TIMESTAMP_PRECISION_PICOS", + } + + +def test_download_arrow_from_job_id_sends_query_options_on_read_rows(): + bqstorage_client = mock.MagicMock() + bqstorage_client.read_rows.return_value = [] + iterator = RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=(), + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + arrow_serialization_options={ + "bufferCompression": "LZ4_FRAME", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", + }, + ) + + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow"): + list(iterator._download_arrow_from_job_id(bqstorage_client=bqstorage_client)) + + bqstorage_client.read_rows.assert_called_once_with( + "projects/p/locations/US/jobs/j/streams/_default", + offset=0, + timeout=None, + arrow_serialization_options={ + "buffer_compression": "LZ4_FRAME", + "picos_timestamp_precision": "TIMESTAMP_PRECISION_MICROS", + }, + ) + + if __name__ == "__main__": unittest.main() From 660ff1dcaf6c32f1ede9f354e88ec44896ca31f8 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 05/10] feat(bigquery): add timestamp_precision for Arrow query results ARROW results always truncated picosecond TIMESTAMP columns to microseconds, because query_and_wait() had no way to ask for another precision. Client.query() and Client.list_rows() already take timestamp_precision for STRUCT_ENCODING results. Add the same timestamp_precision parameter, with the existing TimestampPrecision enum, to query_and_wait(). PICOSECOND maps to TIMESTAMP_PRECISION_PICOS, which returns picosecond columns at full precision as ISO 8601 strings, and anything else maps to TIMESTAMP_PRECISION_MICROS. Nanoseconds aren't offered, since the picosecond strings already keep full precision. The value goes into the arrowSerializationOptions sent on jobs.query and ReadRows, and setting it without ARROW raises ValueError, like compression_codec. It's labeled [Private Preview], as on Client.query(). --- .../google/cloud/bigquery/_arrow_query.py | 22 +++- .../google/cloud/bigquery/_job_helpers.py | 17 ++- .../google/cloud/bigquery/client.py | 14 ++- .../tests/system/test_arrow.py | 102 ++++++++++++++++++ .../unit/test_query_results_format_arrow.py | 48 +++++++++ 5 files changed, 193 insertions(+), 10 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index 1cd1089d7f8c..138f78078355 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -44,6 +44,7 @@ def validate_query( job_config: Optional["job.QueryJobConfig"], *, compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, ) -> None: """Raise before sending a query whose ``ARROW`` results can't be read. @@ -54,11 +55,13 @@ def validate_query( ``Client.default_query_job_config``. compression_codec: The ``compression_codec`` passed to ``query_and_wait()``. + timestamp_precision: The ``timestamp_precision`` passed to + ``query_and_wait()``. Raises: ValueError: - If ``compression_codec`` is set without ``ARROW``, or if an - ``ARROW`` query uses ``WRITE_APPEND``. + If ``compression_codec`` or ``timestamp_precision`` is set + without ``ARROW``, or if an ``ARROW`` query uses ``WRITE_APPEND``. google.cloud.bigquery.exceptions.BigQueryError: If pyarrow or google-cloud-bigquery-storage is missing or too old to read ``ARROW`` results. @@ -66,6 +69,8 @@ def validate_query( if query_results_format != enums.QueryResultsFormat.ARROW: if compression_codec is not None: raise ValueError("compression_codec requires query_results_format=ARROW.") + if timestamp_precision is not None: + raise ValueError("timestamp_precision requires query_results_format=ARROW.") return # Results that don't fit in the jobs.query response are read from the @@ -101,7 +106,10 @@ def _check_dependencies() -> None: ) -def serialization_options(compression_codec: Optional[str]) -> Dict[str, str]: +def serialization_options( + compression_codec: Optional[str], + timestamp_precision: Optional[enums.TimestampPrecision] = None, +) -> Dict[str, str]: """Build the ``arrowSerializationOptions`` sent on ``jobs.query``. ``ReadRows`` gets the same options. Options the caller didn't set are sent @@ -111,7 +119,13 @@ def serialization_options(compression_codec: Optional[str]) -> Dict[str, str]: return { "bufferCompression": getattr(compression_codec, "value", compression_codec) or "COMPRESSION_UNSPECIFIED", - "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", + # Picosecond values come back as ISO 8601 strings at full precision, + # so there's no need to offer nanoseconds. + "picosTimestampPrecision": ( + "TIMESTAMP_PRECISION_PICOS" + if timestamp_precision == enums.TimestampPrecision.PICOSECOND + else "TIMESTAMP_PRECISION_MICROS" + ), } diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py index 384df0cbe8ed..fd0339f33e1d 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -433,6 +433,7 @@ def query_and_wait( max_results: Optional[int] = None, query_results_format: Optional[str] = None, compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, callback: Callable = lambda _: None, ) -> table.RowIterator: """Run the query, wait for it to finish, and return the results. @@ -482,6 +483,11 @@ def query_and_wait( [Beta] The format for query results (e.g. "ARROW" or :class:`~google.cloud.bigquery.enums.QueryResultsFormat.ARROW`). compression_codec (Optional[Union[str, google.cloud.bigquery.enums.QueryResultsCompressionCodec]]): [Beta] Compression codec for Arrow serialization (e.g. "LZ4_FRAME" or :class:`~google.cloud.bigquery.enums.QueryResultsCompressionCodec.LZ4_FRAME`). + timestamp_precision (Optional[enums.TimestampPrecision]): + [Private Preview] Precision of picosecond ``TIMESTAMP`` columns in + ``ARROW`` results. With ``PICOSECOND``, they arrive at full + precision as ISO 8601 strings. Otherwise, they're truncated to + microseconds. Requires ``query_results_format=ARROW``. callback (Callable): A callback function used by bigframes to report query progress. @@ -503,14 +509,17 @@ def query_and_wait( :class:`~google.cloud.bigquery.job.QueryJobConfig` class. ValueError: - If ``compression_codec`` is set without ``ARROW``, or if an - ``ARROW`` query uses ``WRITE_APPEND``. + If ``compression_codec`` or ``timestamp_precision`` is set + without ``ARROW``, or if an ``ARROW`` query uses ``WRITE_APPEND``. google.cloud.bigquery.exceptions.BigQueryError: If pyarrow or google-cloud-bigquery-storage is missing or too old to read ``ARROW`` results. The query isn't sent. """ _arrow_query.validate_query( - query_results_format, job_config, compression_codec=compression_codec + query_results_format, + job_config, + compression_codec=compression_codec, + timestamp_precision=timestamp_precision, ) request_body = _to_query_request( @@ -521,7 +530,7 @@ def query_and_wait( arrow_serialization_options = None if query_results_format == enums.QueryResultsFormat.ARROW: arrow_serialization_options = _arrow_query.serialization_options( - compression_codec + compression_codec, timestamp_precision ) request_body["arrowSerializationOptions"] = arrow_serialization_options diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index 1322a7538a40..b8f0e9b99db2 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -3677,6 +3677,7 @@ def query_and_wait( max_results: Optional[int] = None, query_results_format: Optional[str] = None, compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, ) -> RowIterator: """Run the query, wait for it to finish, and return the results. @@ -3728,6 +3729,11 @@ def query_and_wait( [Beta] The format for query results (e.g. "ARROW" or :class:`~google.cloud.bigquery.enums.QueryResultsFormat.ARROW`). compression_codec (Optional[Union[str, google.cloud.bigquery.enums.QueryResultsCompressionCodec]]): [Beta] Compression codec for Arrow serialization (e.g. "LZ4_FRAME" or :class:`~google.cloud.bigquery.enums.QueryResultsCompressionCodec.LZ4_FRAME`). + timestamp_precision (Optional[enums.TimestampPrecision]): + [Private Preview] Precision of picosecond ``TIMESTAMP`` columns + in ``ARROW`` results. With ``PICOSECOND``, they arrive at full + precision as ISO 8601 strings. Otherwise, they're truncated to + microseconds. Requires ``query_results_format=ARROW``. Returns: google.cloud.bigquery.table.RowIterator: @@ -3747,8 +3753,9 @@ def query_and_wait( :class:`~google.cloud.bigquery.job.QueryJobConfig` class. ValueError: - If ``compression_codec`` is set without ``ARROW``, or if an - ``ARROW`` query uses ``WRITE_APPEND``, including through + If ``compression_codec`` or ``timestamp_precision`` is set + without ``ARROW``, or if an ``ARROW`` query uses + ``WRITE_APPEND``, including through ``default_query_job_config``. google.cloud.bigquery.exceptions.BigQueryError: If pyarrow or google-cloud-bigquery-storage is missing or too @@ -3767,6 +3774,7 @@ def query_and_wait( max_results=max_results, query_results_format=query_results_format, compression_codec=compression_codec, + timestamp_precision=timestamp_precision, ) def _query_and_wait_bigframes( @@ -3784,6 +3792,7 @@ def _query_and_wait_bigframes( max_results: Optional[int] = None, query_results_format: Optional[str] = None, compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, callback: Callable = lambda _: None, ) -> RowIterator: """See query_and_wait. @@ -3818,6 +3827,7 @@ def _query_and_wait_bigframes( max_results=max_results, query_results_format=query_results_format, compression_codec=compression_codec, + timestamp_precision=timestamp_precision, callback=callback, ) diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index 756f37cadf2f..48c223f23d31 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -14,11 +14,13 @@ """System tests for Arrow connector.""" +import datetime from typing import Optional import pyarrow import pytest +from google.api_core import exceptions from google.cloud import bigquery from google.cloud.bigquery import enums @@ -337,3 +339,103 @@ def read(**kwargs): assert compressed.num_rows == 1000 assert compressed.equals(read()) + + +_MICROSECOND_VALUES = [ + datetime.datetime(2025, 1, day, 0, 0, 0, 123456, tzinfo=datetime.timezone.utc) + for day in (1, 2, 3) +] +_TIMESTAMP_PRECISIONS = pytest.mark.parametrize( + "timestamp_precision", + [ + None, + enums.TimestampPrecision.MICROSECOND, + enums.TimestampPrecision.PICOSECOND, + ], + ids=["default", "microsecond", "picosecond"], +) +_PAGE_SIZES = pytest.mark.parametrize( + "page_size", [None, 1], ids=["first_page_only", "with_stream"] +) + + +def _read_timestamps(bigquery_client, query, timestamp_precision, page_size): + rows = bigquery_client.query_and_wait( + query, + job_config=bigquery.QueryJobConfig(use_query_cache=False), + page_size=page_size, + query_results_format=enums.QueryResultsFormat.ARROW, + timestamp_precision=timestamp_precision, + ) + table = pyarrow.Table.from_batches(list(rows.to_arrow_iterable())) + return table.column("ts").to_pylist() + + +@_PAGE_SIZES +@_TIMESTAMP_PRECISIONS +def test_query_and_wait_arrow_timestamp_precision_microsecond_column( + bigquery_client, timestamp_precision, page_size +): + values = _read_timestamps( + bigquery_client, + "SELECT TIMESTAMP_ADD(TIMESTAMP '2025-01-01 00:00:00.123456 UTC', " + "INTERVAL x DAY) AS ts FROM UNNEST([0, 1, 2]) AS x ORDER BY x", + timestamp_precision, + page_size, + ) + + assert values == _MICROSECOND_VALUES + + +@_PAGE_SIZES +@_TIMESTAMP_PRECISIONS +def test_query_and_wait_arrow_timestamp_precision_picosecond_column( + bigquery_client, timestamp_precision, page_size +): + try: + values = _read_timestamps( + bigquery_client, + "SELECT CAST(FORMAT('2025-01-0%dT00:00:00.123456789012Z', x) " + "AS TIMESTAMP(12)) AS ts FROM UNNEST([1, 2, 3]) AS x ORDER BY x", + timestamp_precision, + page_size, + ) + except exceptions.BadRequest as exc: + # TIMESTAMP(12) is enabled per project. The test project has it, but + # a developer's own project may not. + if "Timestamp precision type parameter is not supported" not in str(exc): + raise + pytest.xfail(f"TIMESTAMP(12) isn't enabled in this project: {exc.message}") + + if timestamp_precision == enums.TimestampPrecision.PICOSECOND: + assert [value[:10] for value in values] == [ + "2025-01-01", + "2025-01-02", + "2025-01-03", + ] + assert all(".123456789012" in value for value in values) + else: + assert values == _MICROSECOND_VALUES + + +@_PAGE_SIZES +@_TIMESTAMP_PRECISIONS +def test_query_and_wait_arrow_timestamp_precision_picosecond_table( + bigquery_client, scalars_table_pico, timestamp_precision, page_size +): + values = _read_timestamps( + bigquery_client, + f"SELECT pico_col AS ts FROM `{scalars_table_pico}` ORDER BY pico_col", + timestamp_precision, + page_size, + ) + + if timestamp_precision == enums.TimestampPrecision.PICOSECOND: + assert [value[:10] for value in values] == [ + "2025-01-01", + "2025-01-02", + "2025-01-03", + ] + assert all(".123456789012" in value for value in values) + else: + assert values == _MICROSECOND_VALUES diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index 8b58886bd9d1..8a51d552528f 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -25,6 +25,7 @@ from google.cloud.bigquery import exceptions from google.cloud.bigquery.client import Client from google.cloud.bigquery.enums import QueryResultsCompressionCodec +from google.cloud.bigquery.enums import TimestampPrecision from google.cloud.bigquery.enums import WriteDisposition from google.cloud.bigquery.job import QueryJobConfig from google.cloud.bigquery.table import RowIterator, _EmptyRowIterator @@ -930,5 +931,52 @@ def test_download_arrow_from_job_id_sends_query_options_on_read_rows(): ) +@pytest.mark.parametrize( + "timestamp_precision", + [TimestampPrecision.MICROSECOND, TimestampPrecision.PICOSECOND], +) +def test_query_and_wait_timestamp_precision_requires_arrow( + check_dependencies, timestamp_precision +): + client = mock.MagicMock(spec=Client) + + with pytest.raises(ValueError, match="timestamp_precision requires"): + _query_and_wait(client, timestamp_precision=timestamp_precision) + + client._call_api.assert_not_called() + check_dependencies.assert_not_called() + + +@pytest.mark.parametrize( + ("timestamp_precision", "expected"), + [ + (None, "TIMESTAMP_PRECISION_MICROS"), + (TimestampPrecision.MICROSECOND, "TIMESTAMP_PRECISION_MICROS"), + (TimestampPrecision.PICOSECOND, "TIMESTAMP_PRECISION_PICOS"), + ], +) +def test_client_query_and_wait_arrow_timestamp_precision(timestamp_precision, expected): + client = make_client() + client._call_api = mock.Mock( + return_value={ + "jobReference": {"projectId": "PROJECT", "jobId": "j", "location": "US"}, + "jobComplete": True, + "schema": {"fields": []}, + } + ) + + rows = client.query_and_wait( + "SELECT 1", + query_results_format="ARROW", + timestamp_precision=timestamp_precision, + ) + + data = client._call_api.call_args.kwargs["data"] + assert data["arrowSerializationOptions"]["picosTimestampPrecision"] == expected + assert rows._arrow_serialization_options["picosTimestampPrecision"] == expected + # formatOptions only shapes STRUCT_ENCODING rows, which ARROW doesn't return. + assert "timestampOutputFormat" not in data["formatOptions"] + + if __name__ == "__main__": unittest.main() From ab029f6b8e81a84ceb963eb78996c5d8cb59d45b Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:41:05 +0000 Subject: [PATCH 06/10] fix(bigquery): keep the Arrow first page and wait without rows When an ARROW query's results didn't fit in the jobs.query response, QueryJob.result() dropped that response, because it only kept first pages that had STRUCT_ENCODING rows. The stream then reread every row from offset 0. While waiting for the job, query_and_wait() also passed page_size to jobs.getQueryResults, which returned STRUCT_ENCODING rows that ARROW results never use. For ARROW queries, decide whether rows remain by comparing pageRowCount with totalRows instead of checking for a pageToken, wait with page_size=0 so jobs.getQueryResults is called with maxResults=0, and keep a jobs.query response that has arrowRecordBatch as the first page. RowIterator then returns the first page and streams the rest from the next row, also when it's read again. Queries that don't use ARROW are unchanged. --- .../google/cloud/bigquery/_arrow_query.py | 13 +- .../google/cloud/bigquery/_job_helpers.py | 10 +- .../google/cloud/bigquery/job/query.py | 8 +- .../tests/system/test_arrow.py | 314 ++++++++++++++++++ .../tests/system/test_query.py | 35 +- .../tests/unit/job/test_query.py | 36 ++ .../unit/test_query_results_format_arrow.py | 75 ++++- 7 files changed, 482 insertions(+), 9 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index 138f78078355..491eb270fc81 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -17,7 +17,7 @@ from __future__ import annotations import base64 -from typing import Dict, Iterator, Optional, TYPE_CHECKING +from typing import Any, Dict, Iterator, Optional, TYPE_CHECKING from google.cloud.bigquery import _versions_helpers from google.cloud.bigquery import enums @@ -129,6 +129,16 @@ def serialization_options( } +def has_more_rows(response: Dict[str, Any]) -> bool: + """Return True if rows remain after an ``ARROW`` ``jobs.query`` response. + + ``ARROW`` responses say how many rows the first page holds in + ``pageRowCount``, so compare that with ``totalRows`` instead of checking + for a ``pageToken``. + """ + return int(response.get("pageRowCount", 0)) < int(response.get("totalRows", 0)) + + def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]: """Convert ``jobs.query`` ``arrowSerializationOptions`` to the ``ArrowSerializationOptions`` fields that ``ReadRows`` takes.""" @@ -169,7 +179,6 @@ def download_arrow_from_job( # Step 2: Process initial inline Arrow response from jobs.query if available. if row_iterator._first_page_response: first_page = row_iterator._first_page_response - row_iterator._first_page_response = None job_complete = bool(first_page.get("jobComplete", False)) if job_complete: diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py index fd0339f33e1d..617efd35a911 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -534,6 +534,10 @@ def query_and_wait( ) request_body["arrowSerializationOptions"] = arrow_serialization_options + # jobs.getQueryResults only returns STRUCT_ENCODING rows, so while waiting + # for an ARROW query, request none of them. + wait_page_size = 0 if arrow_serialization_options is not None else page_size + # Some API parameters aren't supported by the jobs.query API. In these # cases, fallback to a jobs.insert call. if not _supported_by_jobs_query(request_body): @@ -554,7 +558,7 @@ def query_and_wait( api_timeout=api_timeout, wait_timeout=wait_timeout, retry=retry, - page_size=page_size, + page_size=wait_page_size, max_results=max_results, query_results_format=query_results_format, arrow_serialization_options=arrow_serialization_options, @@ -617,6 +621,8 @@ def do_query(): ) page_token = query_results.page_token more_pages = page_token is not None + if arrow_serialization_options is not None: + more_pages = _arrow_query.has_more_rows(response) if more_pages or not query_results.complete: # TODO(swast): Avoid a call to jobs.get in some cases (few @@ -628,7 +634,7 @@ def do_query(): api_timeout=api_timeout, wait_timeout=wait_timeout, retry=retry, - page_size=page_size, + page_size=wait_page_size, max_results=max_results, query_results_format=query_results_format, arrow_serialization_options=arrow_serialization_options, diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py index e85e43f6a164..0b7af6293851 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py @@ -1842,9 +1842,13 @@ def is_job_done(): # RowIterator response if there are any rows in it. This prevents us # from stopping the iteration early in the cases where we set # maxResults=0. In that case, we're missing rows and there's no next - # page token. + # page token. ARROW responses from jobs.query carry their rows in + # arrowRecordBatch instead of rows. first_page_response = self._query_results._properties - if "rows" not in first_page_response: + if ( + "rows" not in first_page_response + and "arrowRecordBatch" not in first_page_response + ): first_page_response = None rows = self._client._list_rows_from_query_results( diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index 48c223f23d31..e838fc8e5282 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -15,10 +15,15 @@ """System tests for Arrow connector.""" import datetime +import json +import math from typing import Optional +from unittest import mock +import uuid import pyarrow import pytest +from dateutil import relativedelta from google.api_core import exceptions from google.cloud import bigquery @@ -439,3 +444,312 @@ def test_query_and_wait_arrow_timestamp_precision_picosecond_table( assert all(".123456789012" in value for value in values) else: assert values == _MICROSECOND_VALUES + + +@pytest.mark.parametrize( + ("page_size", "priority", "expected_offset"), + [ + (100, None, 100), + # jobs.query doesn't accept a priority, so this case runs through + # jobs.insert, which returns no rows inline. + (None, enums.QueryPriority.INTERACTIVE, 0), + ], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_matches_struct_encoding( + bigquery_client, bqstorage_client, page_size, priority, expected_offset +): + query = ( + "SELECT x, CAST(x AS STRING) AS s, MOD(x, 2) = 0 AS even, " + "DATE_ADD(DATE '2026-01-01', INTERVAL x DAY) AS d " + "FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + # Wrap read_rows() to check where the stream starts. + read_rows = mock.Mock(wraps=bqstorage_client.read_rows) + + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + page_size=page_size, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + with mock.patch.object(bqstorage_client, "read_rows", read_rows): + batches = list(rows.to_arrow_iterable(bqstorage_client=bqstorage_client)) + + assert pyarrow.Table.from_batches(batches).to_pylist() == expected + read_rows.assert_called_once_with( + f"projects/{bigquery_client.project}/locations/{rows.location}" + f"/jobs/{rows.job_id}/streams/_default", + offset=expected_offset, + timeout=mock.ANY, + arrow_serialization_options=mock.ANY, + ) + + +def _interval_to_relativedelta(value): + return relativedelta.relativedelta( + months=value.months, days=value.days, microseconds=value.nanoseconds // 1000 + ) + + +def _arrow_rows(table): + """Return the rows of an Arrow table as STRUCT_ENCODING returns them. + + Arrow has JSON values as strings and INTERVAL values as MonthDayNano, + while STRUCT_ENCODING rows have parsed JSON and relativedelta values. + """ + converters = {} + for field in table.schema: + if field.metadata == {b"ARROW:extension:name": b"google:sqlType:json"}: + converters[field.name] = json.loads + elif pyarrow.types.is_interval(field.type): + converters[field.name] = _interval_to_relativedelta + rows = table.to_pylist() + for row in rows: + for name, convert in converters.items(): + if row[name] is not None: + row[name] = convert(row[name]) + return rows + + +@pytest.mark.parametrize( + ("page_size", "priority", "expected_offset"), + [ + (1, None, 1), + # jobs.query doesn't accept a priority, so this case runs through + # jobs.insert, which returns no rows inline. + (None, enums.QueryPriority.INTERACTIVE, 0), + ], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_locations( + bigquery_client, + bqstorage_client, + scalars_table_multi_location, + page_size, + priority, + expected_offset, +): + location, table_id = scalars_table_multi_location + query = f"SELECT * FROM `{table_id}` ORDER BY rowindex" + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + read_rows = mock.Mock(wraps=bqstorage_client.read_rows) + + # The query doesn't name a location, so only jobReference has it. + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + page_size=page_size, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + with mock.patch.object(bqstorage_client, "read_rows", read_rows): + batches = list(rows.to_arrow_iterable(bqstorage_client=bqstorage_client)) + + assert rows.location == location + assert _arrow_rows(pyarrow.Table.from_batches(batches)) == expected + read_rows.assert_called_once_with( + f"projects/{bigquery_client.project}/locations/{location}" + f"/jobs/{rows.job_id}/streams/_default", + offset=expected_offset, + timeout=mock.ANY, + arrow_serialization_options=mock.ANY, + ) + + +_DATA_TYPES_QUERY = """ +SELECT + i, + IF(i = 2, NULL, i = 1) AS bool_col, + IF(i = 2, NULL, i * 1000000000000) AS int64_col, + IF(i = 2, NULL, i / 3) AS float64_col, + IF(i = 2, NULL, NUMERIC '-123456789.123456789') AS numeric_col, + IF(i = 2, NULL, BIGNUMERIC '1.23456789012345678901234567890123456789') AS bignumeric_col, + IF(i = 2, NULL, FORMAT('row %d', i)) AS string_col, + IF(i = 2, NULL, CAST(FORMAT('row %d', i) AS BYTES)) AS bytes_col, + IF(i = 2, NULL, DATE_ADD(DATE '2026-01-01', INTERVAL i DAY)) AS date_col, + IF(i = 2, NULL, DATETIME '2026-01-01 12:34:56.789012') AS datetime_col, + IF(i = 2, NULL, TIME '12:34:56.789012') AS time_col, + IF(i = 2, NULL, TIMESTAMP '2026-01-01 12:34:56.789012 UTC') AS timestamp_col, + IF(i = 2, NULL, INTERVAL '1-2 3 4:5:6.789012' YEAR TO SECOND) AS interval_col, + IF(i = 2, NULL, ST_GEOGPOINT(-122.0838511, 37.3860517)) AS geography_col, + IF(i = 2, NULL, JSON '{"a": [1, 2.5, "x"], "b": null, "c": {"d": true}}') AS json_col, + IF(i = 2, NULL, RANGE(DATE '2026-01-01', DATE '2026-02-01')) AS range_date_col, + IF(i = 2, NULL, RANGE(DATETIME '2026-01-01 00:00:00', NULL)) AS range_datetime_col, + IF( + i = 2, + NULL, + RANGE(TIMESTAMP '2026-01-01 00:00:00 UTC', TIMESTAMP '2026-01-02 00:00:00 UTC') + ) AS range_timestamp_col, + IF(i = 2, [], [i, i + 1]) AS array_col, + IF(i = 2, NULL, STRUCT(i AS x, [STRUCT('a' AS s, [1.5, 2.5] AS f)] AS nested)) AS struct_col +FROM UNNEST([1, 2, 3]) AS i +ORDER BY i +""" + + +def test_query_and_wait_arrow_data_types(bigquery_client): + job_config = bigquery.QueryJobConfig(use_query_cache=False) + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait( + _DATA_TYPES_QUERY, job_config=job_config + ) + ] + + rows = bigquery_client.query_and_wait( + _DATA_TYPES_QUERY, + job_config=job_config, + # The first row arrives in the first page, the others in the stream. + page_size=1, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + batches = list(rows.to_arrow_iterable()) + + assert batches[0].num_rows == 1 + assert _arrow_rows(pyarrow.Table.from_batches(batches)) == expected + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["jobs_query", "jobs_insert"], +) +def test_query_and_wait_arrow_unnamed_struct_fields(bigquery_client, priority): + query = "SELECT (1, 'a') AS s, STRUCT(2 AS x, 'b') AS p" + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + # jobs.query doesn't accept a priority, so this runs through jobs.insert. + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + + try: + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + # jobs.query fails with a jobInternalError, which job_retry would + # retry until its deadline. + job_retry=None, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + except exceptions.BadRequest as exc: + if not any(error.get("reason") == "jobInternalError" for error in exc.errors): + raise + pytest.xfail(f"jobs.query can't send unnamed STRUCT fields: {exc.message}") + table = pyarrow.Table.from_batches(list(rows.to_arrow_iterable())) + + assert table.to_pylist() == expected + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["first_page_only", "no_rows_inline"], +) +def test_query_and_wait_arrow_extreme_values( + bigquery_client, scalars_extreme_table, priority +): + # BigQuery fails the extreme INTERVAL values with "Out of range conversion + # for microseconds value", since Arrow's interval type can't hold them. + query = ( + f"SELECT * EXCEPT (interval_col) FROM `{scalars_extreme_table}` " + "ORDER BY rowindex" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + actual = _arrow_rows(pyarrow.Table.from_batches(list(rows.to_arrow_iterable()))) + + # NaN never equals itself, so check it on its own. + assert math.isnan(actual[2].pop("float64_col")) + assert math.isnan(expected[2].pop("float64_col")) + assert actual == expected + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_cached_results(bigquery_client, priority): + # No other query uses this literal, so only the second run can hit the + # cache. + query = ( + f"SELECT x, '{uuid.uuid4().hex}' AS key " + "FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=True) + if priority is not None: + job_config.priority = priority + + def read(): + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + page_size=100, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + table = pyarrow.Table.from_batches(list(rows.to_arrow_iterable())) + return bigquery_client.get_job(rows.job_id, location=rows.location), table + + first_job, first = read() + cached_job, cached = read() + + assert not first_job.cache_hit + assert cached_job.cache_hit + assert cached.num_rows == 1000 + assert cached.equals(first) + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["jobs_query", "jobs_insert"], +) +def test_query_and_wait_arrow_failed_query(bigquery_client, priority): + # The last row fails, so the query fails while it runs. + query = ( + "SELECT IF(x < 1000, x, ERROR(FORMAT('failed at row %d', x))) AS x " + "FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + + with pytest.raises(exceptions.GoogleAPICallError) as struct_encoding_error: + bigquery_client.query_and_wait(query, job_config=job_config) + with pytest.raises(exceptions.GoogleAPICallError) as arrow_error: + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + list(rows.to_arrow_iterable()) + + assert type(arrow_error.value) is type(struct_encoding_error.value) + assert "failed at row 1000" in arrow_error.value.message diff --git a/packages/google-cloud-bigquery/tests/system/test_query.py b/packages/google-cloud-bigquery/tests/system/test_query.py index 69ff08b66073..43ce8d293403 100644 --- a/packages/google-cloud-bigquery/tests/system/test_query.py +++ b/packages/google-cloud-bigquery/tests/system/test_query.py @@ -14,6 +14,7 @@ import datetime import decimal +import json from typing import Tuple import pytest @@ -162,7 +163,7 @@ def test_query_statistics(bigquery_client, query_api_method): assert len(plan) > stages_with_inputs -@pytest.mark.parametrize( +_QUERY_PARAMETER_CASES = pytest.mark.parametrize( ("sql", "expected", "query_parameters"), ( ( @@ -438,6 +439,9 @@ def test_query_statistics(bigquery_client, query_api_method): ), ), ) + + +@_QUERY_PARAMETER_CASES def test_query_parameters( bigquery_client, query_api_method, sql, expected, query_parameters ): @@ -454,6 +458,35 @@ def test_query_parameters( assert rows[0][0] == expected +@_QUERY_PARAMETER_CASES +def test_query_parameters_arrow(bigquery_client, sql, expected, query_parameters): + job_config = bigquery.QueryJobConfig( + query_parameters=query_parameters, use_query_cache=False + ) + try: + rows = bigquery_client.query_and_wait( + sql, + job_config=job_config, + # jobs.query fails unnamed STRUCT fields with a jobInternalError, + # which job_retry would retry until its deadline. + job_retry=None, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + except exceptions.BadRequest as exc: + if not any(error.get("reason") == "jobInternalError" for error in exc.errors): + raise + pytest.xfail(f"jobs.query can't send unnamed STRUCT fields: {exc.message}") + (batch,) = rows.to_arrow_iterable() + assert batch.num_rows == 1 + assert batch.num_columns == 1 + (value,) = batch.column(0).to_pylist() + if batch.schema.field(0).metadata == { + b"ARROW:extension:name": b"google:sqlType:json" + }: + value = json.loads(value) + assert value == expected + + def test_dry_run( bigquery_client: bigquery.Client, query_api_method: str, diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_query.py b/packages/google-cloud-bigquery/tests/unit/job/test_query.py index d6727ae0355e..daa3d4bd4dc7 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_query.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_query.py @@ -1279,6 +1279,42 @@ def test_result_with_done_jobs_query_response_doesnt_call_get_query_results(self timeout=DEFAULT_GET_JOB_TIMEOUT, ) + def test_result_with_done_jobs_query_response_keeps_arrow_first_page(self): + """An ARROW response from jobs.query is kept as the first page, like + rows are, so the stream can continue after it. + """ + job_resource = self._make_resource(started=True, ended=True, location="EU") + conn = make_connection(job_resource) + client = _make_client(self.PROJECT, connection=conn) + query_resource_done = { + "jobComplete": True, + "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID}, + "arrowSchema": {"serializedSchema": "c2NoZW1h"}, + "arrowRecordBatch": {"serializedRecordBatch": "YmF0Y2g="}, + "pageRowCount": "1", + "totalRows": "2", + } + job = google.cloud.bigquery._job_helpers._to_query_job( + client, + "SELECT 1", + request_config=None, + query_response=query_resource_done, + ) + + result = job.result(page_size=0) + + self.assertEqual( + result._first_page_response["arrowRecordBatch"], + query_resource_done["arrowRecordBatch"], + ) + # Only jobs.get, no jobs.getQueryResults. + conn.api_request.assert_called_once_with( + method="GET", + path=f"/projects/{self.PROJECT}/jobs/{self.JOB_ID}", + query_params={"projection": "full"}, + timeout=DEFAULT_GET_JOB_TIMEOUT, + ) + def test_result_with_none_timeout(self): # Verifies that with an intentional None timeout, get job uses None # instead of the default timeout. diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index 8a51d552528f..426d55506da3 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -295,7 +295,8 @@ def test_download_arrow_from_job_id_with_first_page_response(self): batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_first_batch, mock_second_batch]) - self.assertIsNone(iterator._first_page_response) + # The first page is kept for the next read. + self.assertIsNotNone(iterator._first_page_response) mock_pyarrow.ipc.read_schema.assert_called_once_with(raw_schema_bytes) mock_pyarrow.ipc.read_record_batch.assert_has_calls( @@ -354,7 +355,8 @@ def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_stream_batch]) - self.assertIsNone(iterator._first_page_response) + # The first page is kept for the next read. + self.assertIsNotNone(iterator._first_page_response) mock_pyarrow.ipc.read_schema.assert_called_once_with(raw_schema_bytes) mock_pyarrow.ipc.read_record_batch.assert_called_once_with( b"stream_batch_bytes", "deserialized_schema" @@ -473,6 +475,8 @@ def test_download_arrow_from_job_id_avoids_read_rows_when_all_rows_present(self) batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_first_batch]) mock_client._ensure_bqstorage_client.assert_called_once() + batches_second = list(iterator._download_arrow_from_job_id(timeout=5.0)) + self.assertEqual(batches_second, [mock_first_batch]) def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self): mock_client = mock.MagicMock() @@ -866,6 +870,7 @@ def test_query_and_wait_arrow_jobs_insert_omits_arrow_fields(): arrow_serialization_options=_DEFAULT_QUERY_OPTIONS, callback=mock.ANY, ) + assert wait_or_cancel.call_args.kwargs["page_size"] == 0 def test_supported_by_jobs_query_includes_arrow_serialization_options(): @@ -978,5 +983,71 @@ def test_client_query_and_wait_arrow_timestamp_precision(timestamp_precision, ex assert "timestampOutputFormat" not in data["formatOptions"] +def _arrow_query_response(**fields): + return { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "arrowSchema": {"serializedSchema": "c2NoZW1h"}, + **fields, + } + + +@pytest.mark.parametrize( + "response", + [ + _arrow_query_response(pageRowCount="10", totalRows="25"), + _arrow_query_response(jobComplete=False), + ], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_waits_without_fetching_rows(response): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = response + + with ( + mock.patch.object(_job_helpers, "_to_query_job") as to_query_job, + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, + ): + rows = _query_and_wait(client, query_results_format="ARROW", page_size=10) + + assert rows is wait_or_cancel.return_value + assert wait_or_cancel.call_args.args[0] is to_query_job.return_value + assert wait_or_cancel.call_args.kwargs["page_size"] == 0 + # page_size still sets the size of the first page from jobs.query. + assert client._call_api.call_args.kwargs["data"]["maxResults"] == 10 + + +def test_query_and_wait_arrow_all_rows_inline_ignores_page_token(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = _arrow_query_response( + pageRowCount="25", totalRows="25", pageToken="next" + ) + + with mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel: + rows = _query_and_wait(client, query_results_format="ARROW") + + wait_or_cancel.assert_not_called() + assert rows._first_page_response["pageRowCount"] == "25" + + +def test_query_and_wait_without_arrow_waits_with_page_size(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "rows": [], + "schema": {"fields": []}, + "pageToken": "next", + } + + with ( + mock.patch.object(_job_helpers, "_to_query_job"), + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, + ): + _query_and_wait(client, page_size=10) + + assert wait_or_cancel.call_args.kwargs["page_size"] == 10 + + if __name__ == "__main__": unittest.main() From 06c54256e7136aef3fbea6c6ee9cd92d02e843d8 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 07/10] fix(bigquery): apply max_results to the whole Arrow result With ARROW, max_results only limited the first page from jobs.query. Every remaining row was then streamed, so a query with max_results=150 and a 100-row first page returned the whole result. to_arrow() and to_dataframe() also warned that max_results made them fall back to REST, which isn't true for ARROW, and to_arrow() then dropped the bqstorage_client the caller passed. RowIterator now counts rows across the first page and the stream, slices the batch that reaches max_results, and stops. If the first page already reaches the limit, it doesn't open a stream, and query_and_wait() doesn't wait for rows it won't read. A stream that ends early, at the limit, on an error, or because the caller closed the iterator, is cancelled with ReadRowsStream.cancel(), and errors are raised after the batches that arrived before them. For ARROW results, to_arrow() and to_dataframe() skip the warning and keep the caller's bqstorage_client. --- .../google/cloud/bigquery/_arrow_query.py | 84 ++++++--- .../google/cloud/bigquery/_job_helpers.py | 2 +- .../google/cloud/bigquery/table.py | 18 +- .../tests/system/test_arrow.py | 29 +++ .../unit/test_query_results_format_arrow.py | 172 ++++++++++++++++++ 5 files changed, 273 insertions(+), 32 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index 491eb270fc81..78ebf9b69953 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -17,7 +17,7 @@ from __future__ import annotations import base64 -from typing import Any, Dict, Iterator, Optional, TYPE_CHECKING +from typing import Any, Dict, Iterator, Optional, Tuple, TYPE_CHECKING from google.cloud.bigquery import _versions_helpers from google.cloud.bigquery import enums @@ -129,14 +129,29 @@ def serialization_options( } -def has_more_rows(response: Dict[str, Any]) -> bool: +def has_more_rows(response: Dict[str, Any], max_results: Optional[int]) -> bool: """Return True if rows remain after an ``ARROW`` ``jobs.query`` response. ``ARROW`` responses say how many rows the first page holds in - ``pageRowCount``, so compare that with ``totalRows`` instead of checking - for a ``pageToken``. + ``pageRowCount``, so compare that with ``totalRows``, or ``max_results`` if + it's smaller, instead of checking for a ``pageToken``. """ - return int(response.get("pageRowCount", 0)) < int(response.get("totalRows", 0)) + rows_wanted = int(response.get("totalRows", 0)) + if max_results is not None: + rows_wanted = min(rows_wanted, max_results) + return int(response.get("pageRowCount", 0)) < rows_wanted + + +def _slice_batch( + batch: "pyarrow.RecordBatch", remaining: Optional[int] +) -> Tuple["pyarrow.RecordBatch", Optional[int]]: + """Cut ``batch`` to at most ``remaining`` rows, and return it with the + number of rows still wanted after it (``None`` for no limit).""" + if remaining is None: + return batch, None + if batch.num_rows > remaining: + batch = batch.slice(0, remaining) + return batch, remaining - batch.num_rows def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]: @@ -175,6 +190,7 @@ def download_arrow_from_job( pa_schema = None total_rows = row_iterator.total_rows job_complete = False + remaining = row_iterator.max_results # Step 2: Process initial inline Arrow response from jobs.query if available. if row_iterator._first_page_response: @@ -203,10 +219,12 @@ def download_arrow_from_job( pa_schema, ) offset += batch.num_rows + batch, remaining = _slice_batch(batch, remaining) yield batch - # Step 3: Return early if all results were delivered in the first page response. - if job_complete and offset >= total_rows: + # Step 3: Return early if all results were delivered in the first page + # response, or if the first page already reached max_results. + if (job_complete and offset >= total_rows) or remaining == 0: return # Step 4: Stream remaining Arrow record batches from the job default stream. @@ -225,22 +243,36 @@ def download_arrow_from_job( row_iterator._arrow_serialization_options ), ) - for response in reader: - if ( - response.arrow_schema - and response.arrow_schema.serialized_schema - and pa_schema is None - ): - pa_schema = pyarrow.ipc.read_schema( - pyarrow.py_buffer(response.arrow_schema.serialized_schema) - ) - if ( - response.arrow_record_batch - and response.arrow_record_batch.serialized_record_batch - and pa_schema is not None - ): - batch = pyarrow.ipc.read_record_batch( - pyarrow.py_buffer(response.arrow_record_batch.serialized_record_batch), - pa_schema, - ) - yield batch + finished = False + try: + for response in reader: + if ( + response.arrow_schema + and response.arrow_schema.serialized_schema + and pa_schema is None + ): + pa_schema = pyarrow.ipc.read_schema( + pyarrow.py_buffer(response.arrow_schema.serialized_schema) + ) + if ( + response.arrow_record_batch + and response.arrow_record_batch.serialized_record_batch + and pa_schema is not None + ): + batch = pyarrow.ipc.read_record_batch( + pyarrow.py_buffer( + response.arrow_record_batch.serialized_record_batch + ), + pa_schema, + ) + batch, remaining = _slice_batch(batch, remaining) + yield batch + if remaining == 0: + break + else: + finished = True + finally: + # Stop the stream if reading ends before it does: at max_results, on + # an error, or when the caller closes the iterator. + if not finished: + reader.cancel() diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py index 617efd35a911..5b2d1aa2256f 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -622,7 +622,7 @@ def do_query(): page_token = query_results.page_token more_pages = page_token is not None if arrow_serialization_options is not None: - more_pages = _arrow_query.has_more_rows(response) + more_pages = _arrow_query.has_more_rows(response, max_results) if more_pages or not query_results.complete: # TODO(swast): Avoid a call to jobs.get in some cases (few diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index 56e7accd5485..8f77faac9165 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -2408,11 +2408,17 @@ def to_arrow( if pyarrow is None: raise ValueError(_NO_PYARROW_ERROR) - self._maybe_warn_max_results(bqstorage_client) - - if not self._should_use_bqstorage(bqstorage_client, create_bqstorage_client): + if self._query_results_format == QueryResultsFormat.ARROW.value: + # ARROW results apply max_results themselves and read with the + # bqstorage_client they're given, so keep it. create_bqstorage_client = False - bqstorage_client = None + else: + self._maybe_warn_max_results(bqstorage_client) + if not self._should_use_bqstorage( + bqstorage_client, create_bqstorage_client + ): + create_bqstorage_client = False + bqstorage_client = None owns_bqstorage_client = False if not bqstorage_client and create_bqstorage_client: @@ -2911,7 +2917,9 @@ def to_dataframe( if dtypes is None: dtypes = {} - self._maybe_warn_max_results(bqstorage_client) + # ARROW results apply max_results with the BigQuery Storage API too. + if self._query_results_format != QueryResultsFormat.ARROW.value: + self._maybe_warn_max_results(bqstorage_client) if ( self._query_results_format != QueryResultsFormat.ARROW.value diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index e838fc8e5282..5f7c2878aaa7 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -753,3 +753,32 @@ def test_query_and_wait_arrow_failed_query(bigquery_client, priority): assert type(arrow_error.value) is type(struct_encoding_error.value) assert "failed at row 1000" in arrow_error.value.message + + +@pytest.mark.parametrize( + ("page_size", "max_results", "expected_rows", "streamed"), + [ + (None, 10, 10, False), + (100, 150, 150, True), + (100, 5000, 1000, True), + ], + ids=["within_first_page", "across_stream", "past_last_row"], +) +def test_query_and_wait_arrow_max_results( + bigquery_client, bqstorage_client, page_size, max_results, expected_rows, streamed +): + read_rows = mock.Mock(wraps=bqstorage_client.read_rows) + + rows = bigquery_client.query_and_wait( + "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x", + job_config=bigquery.QueryJobConfig(use_query_cache=False), + page_size=page_size, + max_results=max_results, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + with mock.patch.object(bqstorage_client, "read_rows", read_rows): + batches = list(rows.to_arrow_iterable(bqstorage_client=bqstorage_client)) + + table = pyarrow.Table.from_batches(batches) + assert table.column("x").to_pylist() == list(range(1, expected_rows + 1)) + assert read_rows.called == streamed diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index 426d55506da3..022ca1d9dbef 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -15,7 +15,9 @@ import base64 import unittest from unittest import mock +import warnings +import google.api_core.exceptions as core_exceptions import packaging.version import pytest @@ -1049,5 +1051,175 @@ def test_query_and_wait_without_arrow_waits_with_page_size(): assert wait_or_cancel.call_args.kwargs["page_size"] == 10 +def _arrow_pages(*row_counts): + """Serialize an Arrow schema and record batches of consecutive x values.""" + pyarrow = pytest.importorskip("pyarrow") + schema = pyarrow.schema([("x", pyarrow.int64())]) + batches, start = [], 1 + for count in row_counts: + values = pyarrow.array(range(start, start + count), type=pyarrow.int64()) + batch = pyarrow.record_batch([values], schema=schema) + batches.append(batch.serialize().to_pybytes()) + start += count + return schema.serialize().to_pybytes(), batches + + +def _stream_response(batch): + return mock.Mock( + arrow_schema=None, + arrow_record_batch=mock.Mock(serialized_record_batch=batch), + ) + + +def _limited_iterator(row_counts, total_rows, max_results): + """A RowIterator whose first page has row_counts[0] rows and whose stream + returns the rest, plus its mock BigQuery Storage client and stream. + """ + schema, batches = _arrow_pages(*row_counts) + reader = mock.MagicMock() + reader.__iter__.return_value = iter( + [_stream_response(batch) for batch in batches[1:]] + ) + bqstorage_client = mock.MagicMock() + bqstorage_client.read_rows.return_value = reader + iterator = RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=(), + max_results=max_results, + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + first_page_response={ + "jobComplete": True, + "totalRows": str(total_rows), + "arrowSchema": {"serializedSchema": base64.b64encode(schema).decode()}, + "arrowRecordBatch": { + "serializedRecordBatch": base64.b64encode(batches[0]).decode() + }, + }, + ) + return iterator, bqstorage_client, reader + + +def _values(batches): + return [value for batch in batches for value in batch.column(0).to_pylist()] + + +def test_download_arrow_max_results_within_first_page_skips_stream(): + iterator, bqstorage_client, _ = _limited_iterator([10, 15], 25, max_results=4) + + batches = list(iterator._download_arrow_from_job_id(bqstorage_client)) + + assert _values(batches) == [1, 2, 3, 4] + bqstorage_client.read_rows.assert_not_called() + + +def test_download_arrow_max_results_across_stream_cancels_it(): + iterator, bqstorage_client, reader = _limited_iterator( + [10, 10, 5], 25, max_results=15 + ) + + batches = list(iterator._download_arrow_from_job_id(bqstorage_client)) + + assert _values(batches) == list(range(1, 16)) + assert bqstorage_client.read_rows.call_args.kwargs["offset"] == 10 + reader.cancel.assert_called_once_with() + + +def test_download_arrow_max_results_past_last_row_reads_everything(): + iterator, bqstorage_client, reader = _limited_iterator( + [10, 10, 5], 25, max_results=100 + ) + + batches = list(iterator._download_arrow_from_job_id(bqstorage_client)) + + assert _values(batches) == list(range(1, 26)) + reader.cancel.assert_not_called() + + +def test_download_arrow_cancels_stream_when_closed_early(): + iterator, bqstorage_client, reader = _limited_iterator( + [10, 10, 5], 25, max_results=None + ) + batches = iterator._download_arrow_from_job_id(bqstorage_client) + + next(batches) # first page + next(batches) # first stream batch + batches.close() + + reader.cancel.assert_called_once_with() + + +def test_download_arrow_raises_stream_errors_after_earlier_batches(): + iterator, bqstorage_client, reader = _limited_iterator( + [10, 10, 5], 25, max_results=None + ) + first_stream_response = next(reader.__iter__.return_value) + + def fail_after_one_response(): + yield first_stream_response + raise core_exceptions.PermissionDenied("stream failed") + + reader.__iter__.return_value = fail_after_one_response() + received = [] + + with pytest.raises(core_exceptions.PermissionDenied): + for batch in iterator._download_arrow_from_job_id(bqstorage_client): + received.append(batch) + + # The caller gets the error after the batches that arrived before it, so + # a partial result never looks complete. + assert _values(received) == list(range(1, 21)) + reader.cancel.assert_called_once_with() + + +def test_query_and_wait_arrow_first_page_reaching_max_results_doesnt_wait(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = _arrow_query_response( + pageRowCount="10", totalRows="25" + ) + + with mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel: + rows = _query_and_wait(client, query_results_format="ARROW", max_results=10) + + wait_or_cancel.assert_not_called() + assert rows.max_results == 10 + assert client._call_api.call_args.kwargs["data"]["maxResults"] == 10 + + +def test_to_arrow_with_arrow_format_and_max_results_keeps_bqstorage_client(): + pyarrow = pytest.importorskip("pyarrow") + bqstorage_client = mock.MagicMock() + iterator = RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=(), + max_results=5, + query_results_format="ARROW", + ) + batch = pyarrow.record_batch([pyarrow.array([1, 2])], names=["x"]) + + with ( + mock.patch.object( + iterator, "to_arrow_iterable", return_value=iter([batch]) + ) as to_arrow_iterable, + warnings.catch_warnings(record=True) as caught, + ): + warnings.simplefilter("always") + table = iterator.to_arrow(bqstorage_client=bqstorage_client) + + to_arrow_iterable.assert_called_once_with( + bqstorage_client=bqstorage_client, timeout=None + ) + assert table.num_rows == 2 + # ARROW results apply max_results while streaming, so there's no warning + # about falling back to the REST API. + assert not [w for w in caught if "max_results" in str(w.message)] + + if __name__ == "__main__": unittest.main() From 18c7c06d8b526d09c1ea66c4ad87e10bf043268a Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 08/10] feat(bigquery): share one BigQuery Storage client per Client 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. --- .../google/cloud/bigquery/_arrow_query.py | 26 ++-- .../google/cloud/bigquery/client.py | 77 ++++++++++ .../google/cloud/bigquery/table.py | 3 +- .../tests/system/test_arrow.py | 66 +++++++++ .../tests/unit/test_client.py | 139 ++++++++++++++++++ .../unit/test_query_results_format_arrow.py | 64 ++++++-- 6 files changed, 351 insertions(+), 24 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index 78ebf9b69953..daf321bf531d 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -176,23 +176,13 @@ def download_arrow_from_job( if pyarrow is None: raise ValueError(_NO_PYARROW_ERROR) - # Step 1: Ensure BigQuery Read API client is available upfront. - if bqstorage_client is None: - if row_iterator.client is None: - raise ValueError("RowIterator client is None.") - bqstorage_client = row_iterator.client._ensure_bqstorage_client() - if bqstorage_client is None: - raise ValueError( - "The google-cloud-bigquery-storage library is required to read Arrow results." - ) - offset = 0 pa_schema = None total_rows = row_iterator.total_rows job_complete = False remaining = row_iterator.max_results - # Step 2: Process initial inline Arrow response from jobs.query if available. + # Step 1: Process initial inline Arrow response from jobs.query if available. if row_iterator._first_page_response: first_page = row_iterator._first_page_response @@ -222,12 +212,22 @@ def download_arrow_from_job( batch, remaining = _slice_batch(batch, remaining) yield batch - # Step 3: Return early if all results were delivered in the first page + # Step 2: Return early if all results were delivered in the first page # response, or if the first page already reached max_results. if (job_complete and offset >= total_rows) or remaining == 0: return - # Step 4: Stream remaining Arrow record batches from the job default stream. + # Step 3: Stream remaining Arrow record batches from the job default stream, + # with the BigQuery client's shared BigQuery Storage client unless the + # caller passed one. Results that fit in the first page never create it. + if bqstorage_client is None: + if row_iterator.client is None: + raise ValueError("RowIterator client is None.") + bqstorage_client = row_iterator.client._get_shared_bqstorage_client() + if bqstorage_client is None: + raise ValueError( + "The google-cloud-bigquery-storage library is required to read Arrow results." + ) project = row_iterator._project or ( row_iterator.client.project if row_iterator.client else None ) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index b8f0e9b99db2..7df2bf1b93df 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -26,6 +26,7 @@ import math import os import tempfile +import threading import typing import uuid import warnings @@ -166,6 +167,35 @@ # https://github.com/googleapis/python-bigquery/issues/438 _MIN_GET_QUERY_RESULTS_TIMEOUT = 120 +# While a stream is open and the connection is idle, the shared BigQuery +# Storage client pings the server after a minute, and closes the connection +# if a ping isn't answered within a minute. A dropped ReadRows stream is +# resumed at the next unread row. +_BQSTORAGE_KEEPALIVE_TIME_MS = 60 * 1000 +_BQSTORAGE_KEEPALIVE_TIMEOUT_MS = 60 * 1000 + + +def _create_bqstorage_transport(**kwargs): + """Create a BigQuery Storage gRPC transport whose channel sends keepalive + pings. + + ``BigQueryReadClient`` calls this with the arguments it passes to its + default transport. + """ + from google.cloud import bigquery_storage + + transport_class = bigquery_storage.BigQueryReadClient.get_transport_class("grpc") + + def create_channel(*args, options=(), **channel_kwargs): + 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) + + return transport_class(channel=create_channel, **kwargs) + + _LOAD_TABLE_FROM_DATAFRAME_DEPRECATED = ( "Loading DataFrames via google-cloud-bigquery is deprecated. " "For direct, optimized loading, please call 'pandas_gbq.to_gbq()' directly." @@ -300,6 +330,14 @@ def __init__( self._default_load_job_config = copy.deepcopy(default_load_job_config) self.default_job_creation_mode = default_job_creation_mode + # ARROW query results share one BigQuery Storage client. It's created + # the first time a query needs a stream, and close() closes it. + self._universe_domain = _get_client_universe(client_options) + self._shared_bqstorage_client: Optional[ + "google.cloud.bigquery_storage.BigQueryReadClient" + ] = None + self._shared_bqstorage_client_lock = threading.Lock() + # Use property setter so validation can run. self.default_query_job_config = default_query_job_config @@ -348,6 +386,9 @@ def default_load_job_config(self, value: LoadJobConfig): def close(self): """Close the underlying transport objects, releasing system resources. + This also closes the BigQuery Storage client that ``ARROW`` query + results share. + .. note:: The client instance can be used for making additional requests even @@ -356,6 +397,11 @@ def close(self): """ self._http._auth_request.session.close() self._http.close() + with self._shared_bqstorage_client_lock: + bqstorage_client = self._shared_bqstorage_client + self._shared_bqstorage_client = None + if bqstorage_client is not None: + bqstorage_client._transport.close() def get_service_account_email( self, @@ -600,6 +646,7 @@ def _ensure_bqstorage_client( client_info: Optional[ "google.api_core.gapic_v1.client_info.ClientInfo" ] = DEFAULT_BQSTORAGE_CLIENT_INFO, + transport: Optional[Callable[..., Any]] = None, ) -> Optional["google.cloud.bigquery_storage.BigQueryReadClient"]: """Create a BigQuery Storage API client using this client's credentials. @@ -613,6 +660,9 @@ def _ensure_bqstorage_client( client_info: The client info used with a new BigQuery Storage client instance if one is created. + transport: + A callable that creates the transport of a new BigQuery Storage + client instance. If ``None``, the default transport is used. Returns: A BigQuery Storage API client. @@ -680,10 +730,37 @@ def _ensure_bqstorage_client( credentials=self._credentials, client_options=client_options, client_info=amended_client_info, + transport=transport, ) return bqstorage_client + def _get_shared_bqstorage_client( + self, + ) -> Optional["google.cloud.bigquery_storage.BigQueryReadClient"]: + """Return the BigQuery Storage client that ``ARROW`` query results share. + + The client is created on first use, with this client's credentials and + universe domain and a gRPC channel that sends keepalive pings. + :meth:`close` closes it, and the next call creates a new one. + + Returns: + A BigQuery Storage API client, or ``None`` if + google-cloud-bigquery-storage is missing or outdated. + """ + with self._shared_bqstorage_client_lock: + if self._shared_bqstorage_client is None: + client_options = None + if self._universe_domain != _DEFAULT_UNIVERSE: + client_options = google.api_core.client_options.ClientOptions( + universe_domain=self._universe_domain + ) + self._shared_bqstorage_client = self._ensure_bqstorage_client( + client_options=client_options, + transport=_create_bqstorage_transport, + ) + return self._shared_bqstorage_client + def _dataset_from_arg(self, dataset) -> Union[Dataset, DatasetReference]: if isinstance(dataset, str): dataset = DatasetReference.from_string( diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index 8f77faac9165..fd7094fd5f1e 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -2323,7 +2323,8 @@ def _download_arrow_from_job_id( Args: bqstorage_client (Optional[bigquery_storage.BigQueryReadClient]): - Client for BigQuery Storage Read API. If None, one will be created. + Client for BigQuery Storage Read API. If None, the BigQuery + Storage client that the BigQuery client shares is used. timeout (Optional[float]): Timeout in seconds for read operations. diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index 5f7c2878aaa7..e156b296a0a3 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -782,3 +782,69 @@ def test_query_and_wait_arrow_max_results( table = pyarrow.Table.from_batches(batches) assert table.column("x").to_pylist() == list(range(1, expected_rows + 1)) assert read_rows.called == streamed + + +def _query_arrow(client, row_count): + return client.query_and_wait( + f"SELECT x FROM UNNEST(GENERATE_ARRAY(1, {row_count})) AS x ORDER BY x", + job_config=bigquery.QueryJobConfig(use_query_cache=False), + page_size=100, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + + +def test_query_and_wait_arrow_queries_share_one_bqstorage_client(): + from google.cloud.bigquery import client as client_module + + with bigquery.Client() as client: + # Wrap the client factory to count how many clients get created. + ensure = mock.Mock(wraps=client._ensure_bqstorage_client) + with mock.patch.object(client, "_ensure_bqstorage_client", ensure): + tables = [ + pyarrow.Table.from_batches( + list(_query_arrow(client, 1000).to_arrow_iterable()) + ), + _query_arrow(client, 1000).to_arrow(), + pyarrow.Table.from_batches( + list(_query_arrow(client, 1000).to_arrow_iterable()) + ), + ] + shared = client._shared_bqstorage_client + + for table in tables: + assert table.column("x").to_pylist() == list(range(1, 1001)) + ensure.assert_called_once_with( + client_options=None, transport=client_module._create_bqstorage_transport + ) + assert shared is not None + # Leaving the with block calls close(), which closes the shared channel. + assert client._shared_bqstorage_client is None + with pytest.raises(ValueError, match="(?i)closed"): + shared.read_rows("projects/p/locations/US/jobs/j/streams/_default") + + +def test_query_and_wait_arrow_ending_a_read_early_cancels_only_its_stream(): + from google.cloud.bigquery_storage_v1 import reader + + # Wrap ReadRowsStream.cancel() to count calls without changing it. + cancel_stream = reader.ReadRowsStream.cancel + with bigquery.Client() as client, mock.patch.object( + reader.ReadRowsStream, "cancel", autospec=True, side_effect=cancel_stream + ) as cancel: + ended_early_rows = _query_arrow(client, 100000) + ended_early = ended_early_rows.to_arrow_iterable() + read_fully = _query_arrow(client, 100000).to_arrow_iterable() + # Two reads reach the stream, whether or not jobs.query returned a + # first page. + next(ended_early) + next(ended_early) + rows = [next(read_fully), next(read_fully)] + + ended_early.close() + rows.extend(read_fully) + + (stream,) = cancel.call_args.args + cancel.assert_called_once_with(stream) + assert stream._name.endswith(f"/jobs/{ended_early_rows.job_id}/streams/_default") + table = pyarrow.Table.from_batches(rows) + assert table.column("x").to_pylist() == list(range(1, 100001)) diff --git a/packages/google-cloud-bigquery/tests/unit/test_client.py b/packages/google-cloud-bigquery/tests/unit/test_client.py index 88629500a007..7f49ddfadc33 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_client.py +++ b/packages/google-cloud-bigquery/tests/unit/test_client.py @@ -1033,6 +1033,128 @@ def test_ensure_bqstorage_client_existing_client_check_fails(self): ] assert matching_warnings, "Obsolete dependency warning not raised." + def test_get_shared_bqstorage_client_creates_one_client(self): + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud.bigquery import client as client_module + + mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + + with mock.patch( + "google.cloud.bigquery_storage.BigQueryReadClient", mock_client + ): + first = client._get_shared_bqstorage_client() + second = client._get_shared_bqstorage_client() + + self.assertIs(first, mock_client.return_value) + self.assertIs(second, first) + mock_client.assert_called_once_with( + credentials=creds, + client_options=None, + client_info=mock.ANY, + transport=client_module._create_bqstorage_transport, + ) + + @pytest.mark.skipif( + packaging.version.parse(getattr(google.api_core, "__version__", "0.0.0")) + < packaging.version.Version("2.15.0"), + reason="universe_domain not supported with google-api-core < 2.15.0", + ) + def test_get_shared_bqstorage_client_uses_universe_domain(self): + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + + mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + creds = _make_credentials() + client = self._make_one( + project=self.PROJECT, + credentials=creds, + client_options={"universe_domain": "foo.com"}, + ) + + with mock.patch( + "google.cloud.bigquery_storage.BigQueryReadClient", mock_client + ): + client._get_shared_bqstorage_client() + + _, kwargs = mock_client.call_args + self.assertEqual(kwargs["client_options"].universe_domain, "foo.com") + + def test_get_shared_bqstorage_client_tries_again_if_unavailable(self): + from google.cloud.bigquery import client as client_module + + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + + with mock.patch.object( + client, "_ensure_bqstorage_client", return_value=None + ) as ensure: + self.assertIsNone(client._get_shared_bqstorage_client()) + self.assertIsNone(client._get_shared_bqstorage_client()) + + expected = mock.call( + client_options=None, transport=client_module._create_bqstorage_transport + ) + self.assertEqual(ensure.call_args_list, [expected, expected]) + + def test_get_shared_bqstorage_client_creates_one_client_across_threads(self): + import threading + import time + + from google.cloud.bigquery import client as client_module + + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + + def slow_create(**kwargs): + time.sleep(0.05) + return mock.Mock() + + results = [] + + def get_client(): + results.append(client._get_shared_bqstorage_client()) + + with mock.patch.object( + client, "_ensure_bqstorage_client", side_effect=slow_create + ) as ensure: + threads = [threading.Thread(target=get_client) for _ in range(5)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + ensure.assert_called_once_with( + client_options=None, transport=client_module._create_bqstorage_transport + ) + self.assertEqual(len(results), 5) + self.assertTrue(all(result is results[0] for result in results)) + + def test_get_shared_bqstorage_client_channel_sends_keepalive_pings(self): + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + import google.auth.credentials + import grpc + + transport_class = bigquery_storage.BigQueryReadClient.get_transport_class( + "grpc" + ) + creds = google.auth.credentials.AnonymousCredentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + channel = grpc.insecure_channel("localhost:1") + + with mock.patch.object( + transport_class, "create_channel", return_value=channel + ) as create_channel: + bqstorage_client = client._get_shared_bqstorage_client() + channel.close() + + self.assertIsInstance(bqstorage_client._transport, transport_class) + options = create_channel.call_args.kwargs["options"] + self.assertIn(("grpc.keepalive_time_ms", 60000), options) + self.assertIn(("grpc.keepalive_timeout_ms", 60000), options) + # The transport's own options are kept. + self.assertIn(("grpc.max_receive_message_length", -1), options) + def test_create_routine_w_minimal_resource(self): from google.cloud.bigquery.routine import Routine, RoutineReference @@ -1653,6 +1775,23 @@ def test_close(self): http.close.assert_called_once() http._auth_request.session.close.assert_called_once() + def test_close_closes_shared_bqstorage_client(self): + creds = _make_credentials() + http = mock.Mock() + http._auth_request.session = mock.Mock() + client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) + first, second = mock.Mock(), mock.Mock() + + with mock.patch.object( + client, "_ensure_bqstorage_client", side_effect=[first, second] + ): + self.assertIs(client._get_shared_bqstorage_client(), first) + client.close() + self.assertIs(client._get_shared_bqstorage_client(), second) + + first._transport.close.assert_called_once_with() + second._transport.close.assert_not_called() + def test_get_model(self): path = "projects/%s/datasets/%s/models/%s" % ( self.PROJECT, diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index 022ca1d9dbef..d2da999e9178 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -211,7 +211,7 @@ def test_row_iterator_to_arrow_iterable_delegates_when_format_is_arrow(self): def test_download_arrow_from_job_id_constructs_stream_and_reads(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage iterator = RowIterator( client=mock_client, @@ -247,7 +247,7 @@ def test_download_arrow_from_job_id_constructs_stream_and_reads(self): def test_download_arrow_from_job_id_with_first_page_response(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage raw_schema_bytes = b"schema_bytes_123" raw_batch_bytes = b"batch_bytes_123" @@ -321,7 +321,7 @@ def test_download_arrow_from_job_id_with_first_page_response(self): def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage raw_schema_bytes = b"schema_bytes_456" first_page_response = { @@ -373,7 +373,7 @@ def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): def test_download_arrow_from_job_id_missing_storage_client_raises_value_error(self): mock_client = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = None + mock_client._get_shared_bqstorage_client.return_value = None iterator = RowIterator( client=mock_client, @@ -397,7 +397,7 @@ def test_download_arrow_from_job_id_missing_storage_client_raises_value_error(se def test_download_arrow_from_job_id_with_schema_and_batch(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage iterator = RowIterator( client=mock_client, @@ -476,14 +476,14 @@ def test_download_arrow_from_job_id_avoids_read_rows_when_all_rows_present(self) batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_first_batch]) - mock_client._ensure_bqstorage_client.assert_called_once() + mock_client._get_shared_bqstorage_client.assert_not_called() batches_second = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches_second, [mock_first_batch]) def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage raw_schema_bytes = b"schema_bytes_789" raw_batch_bytes = b"batch_bytes_789" @@ -520,7 +520,7 @@ def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self): batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_first_batch]) - mock_client._ensure_bqstorage_client.assert_called_once() + mock_client._get_shared_bqstorage_client.assert_called_once_with() expected_stream = "projects/test-proj/locations/US/jobs/test-job-incomplete/streams/_default" mock_bqstorage.read_rows.assert_called_once_with( expected_stream, @@ -576,7 +576,7 @@ def test_job_helpers_query_and_wait_accepts_enums(self): def test_row_iterator_to_dataframe_iterable_when_format_is_arrow(self): pytest.importorskip("pandas") mock_client = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock.MagicMock() + mock_client._get_shared_bqstorage_client.return_value = mock.MagicMock() iterator = RowIterator( client=mock_client, @@ -599,7 +599,7 @@ def test_row_iterator_to_dataframe_iterable_when_format_is_arrow(self): def test_row_iterator_to_dataframe_when_format_is_arrow(self): pytest.importorskip("pandas") mock_client = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock.MagicMock() + mock_client._get_shared_bqstorage_client.return_value = mock.MagicMock() iterator = RowIterator( client=mock_client, @@ -1221,5 +1221,49 @@ def test_to_arrow_with_arrow_format_and_max_results_keeps_bqstorage_client(): assert not [w for w in caught if "max_results" in str(w.message)] +def test_download_arrow_uses_passed_client_and_leaves_it_open(): + iterator, bqstorage_client, _ = _limited_iterator([10, 15], 25, max_results=None) + + batches = list(iterator._download_arrow_from_job_id(bqstorage_client)) + + assert _values(batches) == list(range(1, 26)) + iterator.client._get_shared_bqstorage_client.assert_not_called() + bqstorage_client._transport.close.assert_not_called() + + +def test_download_arrow_needs_a_client_only_to_stream(): + iterator, _, _ = _limited_iterator([10, 15], 25, max_results=None) + iterator.client = None + batches = iterator._download_arrow_from_job_id() + + assert next(batches).num_rows == 10 + with pytest.raises(ValueError, match="RowIterator client is None"): + next(batches) + + +def test_to_arrow_with_arrow_format_reads_with_the_shared_client(): + pyarrow = pytest.importorskip("pyarrow") + client = mock.MagicMock() + iterator = RowIterator( + client=client, + api_request=mock.MagicMock(), + path=None, + schema=(), + query_results_format="ARROW", + ) + batch = pyarrow.record_batch([pyarrow.array([1, 2])], names=["x"]) + + with mock.patch.object( + iterator, "to_arrow_iterable", return_value=iter([batch]) + ) as to_arrow_iterable: + table = iterator.to_arrow() + + assert table.num_rows == 2 + # to_arrow() doesn't create (and then close) a client of its own; the + # stream uses the BigQuery client's shared one. + client._ensure_bqstorage_client.assert_not_called() + assert to_arrow_iterable.call_args.kwargs["bqstorage_client"] is None + + if __name__ == "__main__": unittest.main() From 31be909ea49bc65118324734d0e757fd50586bc8 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 09/10] feat(bigquery): read Arrow results from explicit destinations The job's default read stream rejects queries whose job config sets a destination table ("This API does not support explicitly-specified destination tables"), so query_and_wait() with ARROW results failed for those queries. When the query wrote to an explicit destination, read the table through create_read_session() instead, with data_format=ARROW and max_stream_count=1 so rows stay in order. The session's read options carry the same serialization options as jobs.query. An empty table has no streams and yields nothing. WRITE_APPEND destinations are already rejected by validate_query(). --- .../google/cloud/bigquery/_arrow_query.py | 99 +++++++++++--- .../google/cloud/bigquery/_job_helpers.py | 5 + .../google/cloud/bigquery/table.py | 3 + .../tests/system/test_arrow.py | 39 ++++++ .../unit/test_query_results_format_arrow.py | 124 +++++++++++++++++- 5 files changed, 252 insertions(+), 18 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index daf321bf531d..b89a91f28d84 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -164,6 +164,76 @@ def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]: } +def _open_job_stream( + row_iterator: "RowIterator", + bqstorage_client: "bigquery_storage.BigQueryReadClient", + offset: int, + timeout: Optional[float], +): + """Open the job's default stream, starting at row ``offset``.""" + project = row_iterator._project or ( + row_iterator.client.project if row_iterator.client else None + ) + location = row_iterator._location or ( + row_iterator.client.location if row_iterator.client else None + ) + stream_name = f"projects/{project}/locations/{location}/jobs/{row_iterator._job_id}/streams/_default" + return bqstorage_client.read_rows( + stream_name, + offset=offset, + timeout=timeout, + arrow_serialization_options=_read_rows_options( + row_iterator._arrow_serialization_options + ), + ) + + +def _open_destination_stream( + row_iterator: "RowIterator", + bqstorage_client: "bigquery_storage.BigQueryReadClient", + timeout: Optional[float], + pa_schema: Optional["pyarrow.Schema"], +): + """Open a read session on the query's explicit destination table. + + The default job stream rejects jobs with an explicit destination table + ("This API does not support explicitly-specified destination tables"), so + read the table instead. One stream keeps the rows in order, and the + session's read options carry the same serialization options as + ``jobs.query``. Explicit destinations always go through ``jobs.insert``, + so there's no first page and the stream starts at row 0. + + Returns ``(reader, pa_schema)``, with ``reader`` set to ``None`` if the + table is empty and the session has no streams. + """ + from google.cloud import bigquery_storage + + session = bqstorage_client.create_read_session( + parent=f"projects/{row_iterator._billing_project}", + read_session=bigquery_storage.types.stream.ReadSession( + table=row_iterator._table.to_bqstorage(), + data_format=bigquery_storage.types.stream.DataFormat.ARROW, + read_options=bigquery_storage.types.stream.ReadSession.TableReadOptions( + arrow_serialization_options=_read_rows_options( + row_iterator._arrow_serialization_options + ) + ), + ), + max_stream_count=1, + timeout=timeout, + ) + if pa_schema is None and session.arrow_schema.serialized_schema: + pa_schema = pyarrow.ipc.read_schema( + pyarrow.py_buffer(session.arrow_schema.serialized_schema) + ) + if not session.streams: + return None, pa_schema + return ( + bqstorage_client.read_rows(session.streams[0].name, offset=0, timeout=timeout), + pa_schema, + ) + + def download_arrow_from_job( row_iterator: "RowIterator", bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None, @@ -218,8 +288,9 @@ def download_arrow_from_job( return # Step 3: Stream remaining Arrow record batches from the job default stream, - # with the BigQuery client's shared BigQuery Storage client unless the - # caller passed one. Results that fit in the first page never create it. + # or from the destination table if the job config set one, with the + # BigQuery client's shared BigQuery Storage client unless the caller passed + # one. Results that fit in the first page never create it. if bqstorage_client is None: if row_iterator.client is None: raise ValueError("RowIterator client is None.") @@ -228,21 +299,15 @@ def download_arrow_from_job( raise ValueError( "The google-cloud-bigquery-storage library is required to read Arrow results." ) - project = row_iterator._project or ( - row_iterator.client.project if row_iterator.client else None - ) - location = row_iterator._location or ( - row_iterator.client.location if row_iterator.client else None - ) - stream_name = f"projects/{project}/locations/{location}/jobs/{row_iterator._job_id}/streams/_default" - reader = bqstorage_client.read_rows( - stream_name, - offset=offset, - timeout=timeout, - arrow_serialization_options=_read_rows_options( - row_iterator._arrow_serialization_options - ), - ) + if row_iterator._explicit_destination and row_iterator._table is not None: + reader, pa_schema = _open_destination_stream( + row_iterator, bqstorage_client, timeout, pa_schema + ) + else: + reader = _open_job_stream(row_iterator, bqstorage_client, offset, timeout) + if reader is None: + # An empty destination table has no streams. + return finished = False try: for response in reader: diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py index 5b2d1aa2256f..52e6d3b3c3d9 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -562,6 +562,9 @@ def query_and_wait( max_results=max_results, query_results_format=query_results_format, arrow_serialization_options=arrow_serialization_options, + explicit_destination=( + job_config is not None and job_config.destination is not None + ), callback=callback, ) @@ -737,6 +740,7 @@ def _wait_or_cancel( *, query_results_format: Optional[str] = None, arrow_serialization_options: Optional[Dict[str, str]] = None, + explicit_destination: bool = False, callback: Callable = lambda _: None, ) -> table.RowIterator: """Wait for a job to complete and return the results. @@ -783,6 +787,7 @@ def _wait_or_cancel( ) query_results._query_results_format = query_results_format query_results._arrow_serialization_options = arrow_serialization_options + query_results._explicit_destination = explicit_destination return query_results except Exception: # Attempt to cancel the job since we can't return the results. diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index fd7094fd5f1e..1f6dae103493 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -1983,6 +1983,9 @@ def __init__( self._job_ended = ended self._query_results_format = query_results_format self._arrow_serialization_options = arrow_serialization_options + # Set by query_and_wait() if the query wrote to a destination table + # set in its job config. + self._explicit_destination = False @property @_disallow_when_arrow diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index e156b296a0a3..bfd2f8e6452d 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -848,3 +848,42 @@ def test_query_and_wait_arrow_ending_a_read_early_cancels_only_its_stream(): assert stream._name.endswith(f"/jobs/{ended_early_rows.job_id}/streams/_default") table = pyarrow.Table.from_batches(rows) assert table.column("x").to_pylist() == list(range(1, 100001)) + + +@pytest.mark.parametrize( + "write_disposition", + [None, enums.WriteDisposition.WRITE_TRUNCATE], + ids=["default", "overwrite"], +) +def test_query_and_wait_arrow_explicit_destination( + bigquery_client, bqstorage_client, dataset_id, write_disposition +): + table_name = f"arrow_destination_{write_disposition or 'default'}".lower() + table_id = f"{bigquery_client.project}.{dataset_id}.{table_name}" + job_config = bigquery.QueryJobConfig(destination=table_id) + if write_disposition is not None: + # The rows already in the table are replaced, not read. + bigquery_client.query_and_wait( + f"CREATE OR REPLACE TABLE `{table_id}` AS SELECT -1 AS x" + ) + job_config.write_disposition = write_disposition + # Wrap create_read_session() to check how the table is read. + create_read_session = mock.Mock(wraps=bqstorage_client.create_read_session) + + rows = bigquery_client.query_and_wait( + "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x", + job_config=job_config, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + with mock.patch.object( + bqstorage_client, "create_read_session", create_read_session + ): + batches = list(rows.to_arrow_iterable(bqstorage_client=bqstorage_client)) + + table = pyarrow.Table.from_batches(batches) + assert table.column("x").to_pylist() == list(range(1, 1001)) + kwargs = create_read_session.call_args.kwargs + assert kwargs["read_session"].table == ( + f"projects/{bigquery_client.project}/datasets/{dataset_id}/tables/{table_name}" + ) + assert kwargs["max_stream_count"] == 1 diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index d2da999e9178..df3523e20f06 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -30,7 +30,7 @@ from google.cloud.bigquery.enums import TimestampPrecision from google.cloud.bigquery.enums import WriteDisposition from google.cloud.bigquery.job import QueryJobConfig -from google.cloud.bigquery.table import RowIterator, _EmptyRowIterator +from google.cloud.bigquery.table import RowIterator, TableReference, _EmptyRowIterator from .helpers import make_client @@ -870,6 +870,7 @@ def test_query_and_wait_arrow_jobs_insert_omits_arrow_fields(): max_results=mock.ANY, query_results_format="ARROW", arrow_serialization_options=_DEFAULT_QUERY_OPTIONS, + explicit_destination=True, callback=mock.ANY, ) assert wait_or_cancel.call_args.kwargs["page_size"] == 0 @@ -1265,5 +1266,126 @@ def test_to_arrow_with_arrow_format_reads_with_the_shared_client(): assert to_arrow_iterable.call_args.kwargs["bqstorage_client"] is None +@pytest.mark.parametrize( + ("job_config", "expected"), + [ + (QueryJobConfig(destination="p.d.t"), True), + # jobs.query doesn't accept a priority either, so this also runs + # through jobs.insert, without a destination. + (QueryJobConfig(priority="INTERACTIVE"), False), + ], + ids=["destination", "no_destination"], +) +def test_query_and_wait_arrow_marks_explicit_destination(job_config, expected): + client = mock.MagicMock(spec=Client) + + with ( + mock.patch.object(_job_helpers, "query_jobs_insert") as query_jobs_insert, + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, + ): + _query_and_wait(client, job_config=job_config, query_results_format="ARROW") + + wait_or_cancel.assert_called_once_with( + query_jobs_insert.return_value, + api_timeout=mock.ANY, + wait_timeout=mock.ANY, + retry=mock.ANY, + page_size=mock.ANY, + max_results=mock.ANY, + query_results_format="ARROW", + arrow_serialization_options=mock.ANY, + explicit_destination=expected, + callback=mock.ANY, + ) + + +def test_wait_or_cancel_sets_explicit_destination(): + job = mock.MagicMock(dry_run=False) + job.result.return_value = RowIterator( + client=mock.MagicMock(), api_request=mock.MagicMock(), path=None, schema=() + ) + assert not job.result.return_value._explicit_destination + + rows = _job_helpers._wait_or_cancel( + job, + api_timeout=None, + wait_timeout=None, + retry=None, + page_size=0, + max_results=None, + explicit_destination=True, + ) + + assert rows._explicit_destination + + +def _destination_iterator(row_counts, serialized_schema=None): + """A RowIterator for a query with an explicit destination table, plus its + mock BigQuery Storage client. + """ + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + options = bigquery_storage.types.ArrowSerializationOptions + if "picos_timestamp_precision" not in options.meta.fields: + # query_and_wait() rejects ARROW with this release before it runs. + pytest.skip("google-cloud-bigquery-storage has no timestamp precision") + schema, batches = _arrow_pages(*row_counts) + reader = mock.MagicMock() + reader.__iter__.return_value = iter([_stream_response(batch) for batch in batches]) + stream = mock.Mock() + stream.name = "projects/p/locations/us/sessions/s/streams/0" + session = mock.Mock(streams=[stream] if row_counts else []) + session.arrow_schema.serialized_schema = ( + schema if serialized_schema is None else serialized_schema + ) + bqstorage_client = mock.MagicMock() + bqstorage_client.create_read_session.return_value = session + bqstorage_client.read_rows.return_value = reader + iterator = RowIterator( + client=mock.MagicMock(project="billing-project"), + api_request=mock.MagicMock(), + path=None, + schema=(), + table=TableReference.from_string("p.d.t"), + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + ) + iterator._explicit_destination = True + return iterator, bqstorage_client + + +def test_download_arrow_explicit_destination_reads_a_read_session(): + iterator, bqstorage_client = _destination_iterator([10, 15]) + from google.cloud import bigquery_storage + + batches = list(iterator._download_arrow_from_job_id(bqstorage_client)) + + assert _values(batches) == list(range(1, 26)) + bqstorage_client.create_read_session.assert_called_once_with( + parent="projects/billing-project", + read_session=bigquery_storage.types.stream.ReadSession( + table="projects/p/datasets/d/tables/t", + data_format=bigquery_storage.types.stream.DataFormat.ARROW, + read_options=bigquery_storage.types.stream.ReadSession.TableReadOptions( + arrow_serialization_options=_DEFAULT_READ_OPTIONS + ), + ), + max_stream_count=1, + timeout=None, + ) + # The read session carries the serialization options, so ReadRows doesn't. + bqstorage_client.read_rows.assert_called_once_with( + "projects/p/locations/us/sessions/s/streams/0", offset=0, timeout=None + ) + + +def test_download_arrow_explicit_destination_without_streams_reads_nothing(): + iterator, bqstorage_client = _destination_iterator([], serialized_schema=b"") + + assert list(iterator._download_arrow_from_job_id(bqstorage_client)) == [] + bqstorage_client.read_rows.assert_not_called() + + if __name__ == "__main__": unittest.main() From b501684e9b1c2fe19c5ace8c5c5a664f42536fde Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 10/10] fix(bigquery): keep the Arrow schema for empty results A query that returned no rows with ARROW yielded no batches, so to_arrow() fell back to the iterator's schema, which jobs.query omits for ARROW, and returned a table without columns. jobs.query sends arrowSchema even for zero rows, but the default job stream sends nothing for them. to_arrow_iterable() now yields one empty RecordBatch if no rows arrive, like _EmptyRowIterator already does. It has the Arrow schema from the first page or the read session; otherwise the schema from jobs.getQueryResults, which the jobs.insert path has; otherwise no columns. to_arrow() and pyarrow.Table.from_batches() keep that schema. --- .../google/cloud/bigquery/_arrow_query.py | 35 +++++- .../tests/system/test_arrow.py | 48 ++++++++ .../unit/test_query_results_format_arrow.py | 115 +++++++++++++++++- 3 files changed, 193 insertions(+), 5 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index b89a91f28d84..881fdc0661a6 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -17,7 +17,7 @@ from __future__ import annotations import base64 -from typing import Any, Dict, Iterator, Optional, Tuple, TYPE_CHECKING +from typing import Any, Dict, Iterator, Optional, Sequence, Tuple, TYPE_CHECKING from google.cloud.bigquery import _versions_helpers from google.cloud.bigquery import enums @@ -154,6 +154,29 @@ def _slice_batch( return batch, remaining - batch.num_rows +def _empty_batch( + pa_schema: Optional["pyarrow.Schema"], bq_schema: Sequence[Any] +) -> "pyarrow.RecordBatch": + """Build a batch without rows that keeps the result's schema. + + Use the Arrow schema BigQuery sent if there is one. The default job + stream sends nothing for zero rows, so fall back to the schema from + ``jobs.getQueryResults``, and else to no columns. + """ + if pa_schema is None and bq_schema: + # Imported here: _pandas_helpers imports pandas_gbq, which imports + # this package's table module. + from google.cloud.bigquery import _pandas_helpers + + pa_schema = _pandas_helpers.bq_to_arrow_schema(bq_schema) + if pa_schema is None: + pa_schema = pyarrow.schema([]) + return pyarrow.RecordBatch.from_arrays( + [pyarrow.array([], type=field.type) for field in pa_schema], + schema=pa_schema, + ) + + def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]: """Convert ``jobs.query`` ``arrowSerializationOptions`` to the ``ArrowSerializationOptions`` fields that ``ReadRows`` takes.""" @@ -251,6 +274,9 @@ def download_arrow_from_job( total_rows = row_iterator.total_rows job_complete = False remaining = row_iterator.max_results + # Results without rows still yield one empty batch, so callers get the + # schema. + yielded = False # Step 1: Process initial inline Arrow response from jobs.query if available. if row_iterator._first_page_response: @@ -280,11 +306,14 @@ def download_arrow_from_job( ) offset += batch.num_rows batch, remaining = _slice_batch(batch, remaining) + yielded = True yield batch # Step 2: Return early if all results were delivered in the first page # response, or if the first page already reached max_results. if (job_complete and offset >= total_rows) or remaining == 0: + if not yielded: + yield _empty_batch(pa_schema, row_iterator._schema) return # Step 3: Stream remaining Arrow record batches from the job default stream, @@ -307,6 +336,7 @@ def download_arrow_from_job( reader = _open_job_stream(row_iterator, bqstorage_client, offset, timeout) if reader is None: # An empty destination table has no streams. + yield _empty_batch(pa_schema, row_iterator._schema) return finished = False try: @@ -331,6 +361,7 @@ def download_arrow_from_job( pa_schema, ) batch, remaining = _slice_batch(batch, remaining) + yielded = True yield batch if remaining == 0: break @@ -341,3 +372,5 @@ def download_arrow_from_job( # an error, or when the caller closes the iterator. if not finished: reader.cancel() + if not yielded: + yield _empty_batch(pa_schema, row_iterator._schema) diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index bfd2f8e6452d..2730d169b33b 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -887,3 +887,51 @@ def test_query_and_wait_arrow_explicit_destination( f"projects/{bigquery_client.project}/datasets/{dataset_id}/tables/{table_name}" ) assert kwargs["max_stream_count"] == 1 + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["jobs_query", "jobs_insert"], +) +def test_query_and_wait_arrow_zero_rows_keep_the_schema(bigquery_client, priority): + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + # jobs.query doesn't accept a priority, so this runs through jobs.insert. + job_config.priority = priority + + def query(): + return bigquery_client.query_and_wait( + "SELECT x, CAST(x AS STRING) AS s " + "FROM UNNEST(CAST([] AS ARRAY)) AS x", + job_config=job_config, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + + batches = list(query().to_arrow_iterable()) + + assert len(batches) == 1 + for table in (pyarrow.Table.from_batches(batches), query().to_arrow()): + assert table.num_rows == 0 + assert table.schema.names == ["x", "s"] + assert table.schema.field("x").type == pyarrow.int64() + assert table.schema.field("s").type == pyarrow.string() + + +def test_query_and_wait_arrow_statements_without_results(bigquery_client, dataset_id): + table_id = f"{bigquery_client.project}.{dataset_id}.arrow_statements" + + for statement in ( + f"CREATE OR REPLACE TABLE `{table_id}` AS SELECT 1 AS x", + f"UPDATE `{table_id}` SET x = 2 WHERE TRUE", + ): + rows = bigquery_client.query_and_wait( + statement, query_results_format=enums.QueryResultsFormat.ARROW + ) + (batch,) = rows.to_arrow_iterable() + assert batch.num_rows == 0 + + assert rows.num_dml_affected_rows == 1 + assert [ + row.x for row in bigquery_client.query_and_wait(f"SELECT x FROM `{table_id}`") + ] == [2] diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index df3523e20f06..d602b01af3d9 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -232,7 +232,11 @@ def test_download_arrow_from_job_id_constructs_stream_and_reads(self): with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) - self.assertEqual(batches, []) + # No rows and no schema: one empty batch without columns. + self.assertEqual( + batches, [mock_pyarrow.RecordBatch.from_arrays.return_value] + ) + mock_pyarrow.schema.assert_called_once_with([]) expected_stream_name = ( "projects/test-proj/locations/US/jobs/test-job-456/streams/_default" @@ -1380,12 +1384,115 @@ def test_download_arrow_explicit_destination_reads_a_read_session(): ) -def test_download_arrow_explicit_destination_without_streams_reads_nothing(): - iterator, bqstorage_client = _destination_iterator([], serialized_schema=b"") +@pytest.mark.parametrize( + ("serialized_schema", "expected_names"), + [(None, ["x"]), (b"", [])], + ids=["with_schema", "without_schema"], +) +def test_download_arrow_explicit_destination_without_rows( + serialized_schema, expected_names +): + iterator, bqstorage_client = _destination_iterator( + [], serialized_schema=serialized_schema + ) - assert list(iterator._download_arrow_from_job_id(bqstorage_client)) == [] + (batch,) = iterator._download_arrow_from_job_id(bqstorage_client) + + assert batch.num_rows == 0 + assert batch.schema.names == expected_names bqstorage_client.read_rows.assert_not_called() +def _zero_row_iterator(): + """A RowIterator whose jobs.query response has a schema and no rows.""" + schema, _ = _arrow_pages() + return RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=(), + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + first_page_response={ + "jobComplete": True, + "totalRows": "0", + "arrowSchema": {"serializedSchema": base64.b64encode(schema).decode()}, + # A response without rows carries no serialized batch. + "arrowRecordBatch": {}, + }, + ) + + +def test_download_arrow_zero_rows_inline_keep_the_schema(): + iterator = _zero_row_iterator() + + (batch,) = iterator._download_arrow_from_job_id() + + assert batch.num_rows == 0 + assert batch.schema.names == ["x"] + iterator.client._get_shared_bqstorage_client.assert_not_called() + + +def test_to_arrow_zero_rows_keep_the_schema(): + table = _zero_row_iterator().to_arrow() + + assert table.num_rows == 0 + assert table.schema.names == ["x"] + + +def test_download_arrow_zero_rows_streamed_use_the_query_results_schema(): + pyarrow = pytest.importorskip("pyarrow") + from google.cloud.bigquery.schema import SchemaField + + bqstorage_client = mock.MagicMock() + # The default job stream sends no responses for zero rows. + bqstorage_client.read_rows.return_value = [] + iterator = RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=[SchemaField("x", "INTEGER")], + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + first_page_response={"jobComplete": False}, + ) + + (batch,) = iterator._download_arrow_from_job_id(bqstorage_client) + + assert batch.num_rows == 0 + assert batch.schema.names == ["x"] + assert batch.schema.field("x").type == pyarrow.int64() + + +def test_download_arrow_rows_yield_no_extra_empty_batch(): + schema, batches = _arrow_pages(3) + iterator = RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=(), + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + first_page_response={ + "jobComplete": True, + "totalRows": "3", + "arrowSchema": {"serializedSchema": base64.b64encode(schema).decode()}, + "arrowRecordBatch": { + "serializedRecordBatch": base64.b64encode(batches[0]).decode() + }, + }, + ) + + (batch,) = iterator._download_arrow_from_job_id() + + assert batch.num_rows == 3 + + if __name__ == "__main__": unittest.main()