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