From 16af7f274287ee0f6854843aa69285d26019c722 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 1/2] 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 8fa08700a208..4a235ed66f85 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 @@ -84,6 +84,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, @@ -153,7 +177,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 2168de25de54..9a458374f83d 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -517,11 +517,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. @@ -546,6 +547,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, ) @@ -619,6 +621,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, ) @@ -659,6 +662,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: @@ -700,6 +704,7 @@ def _supported_by_jobs_query(request_body: Dict[str, Any]) -> bool: "reservation", "maxSlots", "queryResultsFormat", + "arrowSerializationOptions", } unsupported_keys = request_keys - keys_allowlist @@ -715,6 +720,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. @@ -760,6 +766,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 453fc6678646..369140f47573 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -1910,6 +1910,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__( @@ -1938,6 +1941,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, @@ -1972,6 +1976,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 88dc5bd2daf1..7563b2bf3d6b 100644 --- a/packages/google-cloud-bigquery/noxfile.py +++ b/packages/google-cloud-bigquery/noxfile.py @@ -213,6 +213,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 9db96503b00b..b6e4c4a1d079 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -274,3 +274,39 @@ 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"}] + + +@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 dd11f8f6f686..cbf443e330be 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", ) @@ -703,5 +728,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 eed63c7773c3e14cef5666913ffa0bc6b3232389 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Fri, 16 Oct 2026 01:15:33 +0000 Subject: [PATCH 2/2] ci(bigquery): install released google-cloud-bigquery-storage for mypy google-cloud-bigquery-storage 2.43.0 has the read_rows() arrow_serialization_options argument, so mypy can type-check against the released package instead of the in-repo copy. --- packages/google-cloud-bigquery/noxfile.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/packages/google-cloud-bigquery/noxfile.py b/packages/google-cloud-bigquery/noxfile.py index 7563b2bf3d6b..88dc5bd2daf1 100644 --- a/packages/google-cloud-bigquery/noxfile.py +++ b/packages/google-cloud-bigquery/noxfile.py @@ -213,14 +213,6 @@ 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"