From 63d04b5cba04d6a5c0ab8d59c7ddb61d1f2294c0 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH] 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()