fix(bigquery): send the same Arrow options on jobs.query and ReadRows

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

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.

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

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