Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ohmayr wants to merge 1 commit into from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 33 additions & 2 deletions packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"],
}
Comment on lines +118 to +125

@gemini-code-assist gemini-code-assist Bot Oct 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using direct dictionary lookups like options["bufferCompression"] can raise a KeyError if the dictionary is passed with missing keys (for example, if a user instantiates RowIterator with custom options). It is safer and more robust to use .get() with fallback defaults, which also simplifies the fallback logic when options is None.

Suggested change
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 _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 {}
return {
"buffer_compression": options.get("bufferCompression", "COMPRESSION_UNSPECIFIED"),
"picos_timestamp_precision": options.get("picosTimestampPrecision", "TIMESTAMP_PRECISION_MICROS"),
}

@ohmayr ohmayr Oct 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These options always come from serialization_options(), which sets both keys, so a missing key would be a bug I'd want to see as a KeyError rather than paper over with defaults. RowIterator's constructor is internal to query_and_wait() here.



def download_arrow_from_job(
row_iterator: "RowIterator",
bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None,
Expand Down Expand Up @@ -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
Expand Down
17 changes: 12 additions & 5 deletions packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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,
)

Expand Down Expand Up @@ -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,
)

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
5 changes: 5 additions & 0 deletions packages/google-cloud-bigquery/google/cloud/bigquery/table.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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__(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
8 changes: 8 additions & 0 deletions packages/google-cloud-bigquery/noxfile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
36 changes: 36 additions & 0 deletions packages/google-cloud-bigquery/tests/system/test_arrow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Loading
Loading