Skip to content

refactor(bigquery): move Arrow query result reading to _arrow_query #18644

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

Merged
ohmayr merged 1 commit into from
Oct 15, 2026
Merged

refactor(bigquery): move Arrow query result reading to _arrow_query #18644

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
125 changes: 125 additions & 0 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
@@ -0,0 +1,125 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Read query results returned in Apache Arrow format."""

from __future__ import annotations

import base64
from typing import Iterator, Optional, TYPE_CHECKING

try:
import pyarrow # type: ignore
except ImportError:
pyarrow = None # type: ignore[assignment]

if TYPE_CHECKING: # pragma: NO COVER
from google.cloud import bigquery_storage
from google.cloud.bigquery.table import RowIterator

_NO_PYARROW_ERROR = (
"The pyarrow library is not installed, please install "
"pyarrow to use the to_arrow() function."
)


def download_arrow_from_job(
row_iterator: "RowIterator",
bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None,
timeout: Optional[float] = None,
) -> Iterator["pyarrow.RecordBatch"]:
"""Yield Arrow record batches for ``row_iterator``'s ARROW query results.

See :meth:`google.cloud.bigquery.table.RowIterator._download_arrow_from_job_id`.
"""
if pyarrow is None:
raise ValueError(_NO_PYARROW_ERROR)

# Step 1: Ensure BigQuery Read API client is available upfront.
if bqstorage_client is None:
if row_iterator.client is None:
raise ValueError("RowIterator client is None.")
bqstorage_client = row_iterator.client._ensure_bqstorage_client()
if bqstorage_client is None:
raise ValueError(
"The google-cloud-bigquery-storage library is required to read Arrow results."
)

offset = 0
pa_schema = None
total_rows = row_iterator.total_rows
job_complete = False

# Step 2: Process initial inline Arrow response from jobs.query if available.
if row_iterator._first_page_response:
first_page = row_iterator._first_page_response
row_iterator._first_page_response = None

job_complete = bool(first_page.get("jobComplete", False))
if job_complete:
total_rows = int(first_page.get("totalRows", 0))

arrow_schema_json = first_page.get("arrowSchema")
if isinstance(arrow_schema_json, dict):
schema_bytes = arrow_schema_json.get("serializedSchema")
if schema_bytes:
if isinstance(schema_bytes, str):
schema_bytes = base64.b64decode(schema_bytes)
pa_schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes))

arrow_batch_json = first_page.get("arrowRecordBatch")
if isinstance(arrow_batch_json, dict) and pa_schema is not None:
batch_bytes = arrow_batch_json.get("serializedRecordBatch")
if batch_bytes:
if isinstance(batch_bytes, str):
batch_bytes = base64.b64decode(batch_bytes)
batch = pyarrow.ipc.read_record_batch(
pyarrow.py_buffer(batch_bytes),
pa_schema,
)
offset += batch.num_rows
yield batch

# Step 3: Return early if all results were delivered in the first page response.
if job_complete and offset >= total_rows:
return

# Step 4: Stream remaining Arrow record batches from the job default stream.
project = row_iterator._project or (
row_iterator.client.project if row_iterator.client else None
)
location = row_iterator._location or (
row_iterator.client.location if row_iterator.client else None
)
stream_name = f"projects/{project}/locations/{location}/jobs/{row_iterator._job_id}/streams/_default"
Comment on lines +99 to +105

@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

If row_iterator.client is None (which can happen if bqstorage_client is provided externally) or if the project, location, or job ID is missing, constructing the stream_name with None values will result in an invalid resource path and cause a cryptic API error. We should validate that project, location, and job_id are all present before constructing the stream name.

    project = row_iterator._project or (
        row_iterator.client.project if row_iterator.client else None
    )
    location = row_iterator._location or (
        row_iterator.client.location if row_iterator.client else None
    )
    job_id = row_iterator._job_id
    if not project or not location or not job_id:
        raise ValueError(
            "Project, location, and job ID must all be present to stream Arrow results."
        )
    stream_name = f"projects/{project}/locations/{location}/jobs/{job_id}/streams/_default"

@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.

I'd rather keep this PR a pure move with no behavior changes. On this path the project, location and job ID come from the job reference that query_and_wait() sets on the RowIterator, so they're populated in practice. If we want stricter validation it should be its own change.

reader = bqstorage_client.read_rows(stream_name, offset=offset, timeout=timeout)
for response in reader:
if (
response.arrow_schema
and response.arrow_schema.serialized_schema
and pa_schema is None
):
pa_schema = pyarrow.ipc.read_schema(
pyarrow.py_buffer(response.arrow_schema.serialized_schema)
)
if (
response.arrow_record_batch
and response.arrow_record_batch.serialized_record_batch
and pa_schema is not None
):
batch = pyarrow.ipc.read_record_batch(
pyarrow.py_buffer(response.arrow_record_batch.serialized_record_batch),
pa_schema,
)
yield batch
83 changes: 4 additions & 79 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 @@ -16,7 +16,6 @@

from __future__ import absolute_import

import base64
import copy
import datetime
import functools
Expand Down Expand Up @@ -59,6 +58,7 @@
import google.cloud._helpers # type: ignore
from google.api_core.page_iterator import HTTPIterator
from google.cloud.bigquery import (
_arrow_query,
_helpers,
_pandas_helpers,
_string_references,
Expand Down Expand Up @@ -2325,84 +2325,9 @@ def _download_arrow_from_job_id(
Yields:
pyarrow.RecordBatch: Record batches generated from the query result.
"""
if pyarrow is None:
raise ValueError(_NO_PYARROW_ERROR)

# Step 1: Ensure BigQuery Read API client is available upfront.
if bqstorage_client is None:
if self.client is None:
raise ValueError("RowIterator client is None.")
bqstorage_client = self.client._ensure_bqstorage_client()
if bqstorage_client is None:
raise ValueError(
"The google-cloud-bigquery-storage library is required to read Arrow results."
)

offset = 0
pa_schema = None
total_rows = self.total_rows
job_complete = False

# Step 2: Process initial inline Arrow response from jobs.query if available.
if self._first_page_response:
first_page = self._first_page_response
self._first_page_response = None

job_complete = bool(first_page.get("jobComplete", False))
if job_complete:
total_rows = int(first_page.get("totalRows", 0))

arrow_schema_json = first_page.get("arrowSchema")
if isinstance(arrow_schema_json, dict):
schema_bytes = arrow_schema_json.get("serializedSchema")
if schema_bytes:
if isinstance(schema_bytes, str):
schema_bytes = base64.b64decode(schema_bytes)
pa_schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes))

arrow_batch_json = first_page.get("arrowRecordBatch")
if isinstance(arrow_batch_json, dict) and pa_schema is not None:
batch_bytes = arrow_batch_json.get("serializedRecordBatch")
if batch_bytes:
if isinstance(batch_bytes, str):
batch_bytes = base64.b64decode(batch_bytes)
batch = pyarrow.ipc.read_record_batch(
pyarrow.py_buffer(batch_bytes),
pa_schema,
)
offset += batch.num_rows
yield batch

# Step 3: Return early if all results were delivered in the first page response.
if job_complete and offset >= total_rows:
return

# Step 4: Stream remaining Arrow record batches from the job default stream.
project = self._project or (self.client.project if self.client else None)
location = self._location or (self.client.location if self.client else None)
stream_name = f"projects/{project}/locations/{location}/jobs/{self._job_id}/streams/_default"
reader = bqstorage_client.read_rows(stream_name, offset=offset, timeout=timeout)
for response in reader:
if (
response.arrow_schema
and response.arrow_schema.serialized_schema
and pa_schema is None
):
pa_schema = pyarrow.ipc.read_schema(
pyarrow.py_buffer(response.arrow_schema.serialized_schema)
)
if (
response.arrow_record_batch
and response.arrow_record_batch.serialized_record_batch
and pa_schema is not None
):
batch = pyarrow.ipc.read_record_batch(
pyarrow.py_buffer(
response.arrow_record_batch.serialized_record_batch
),
pa_schema,
)
yield batch
return _arrow_query.download_arrow_from_job(
self, bqstorage_client=bqstorage_client, timeout=timeout
)

# If changing the signature of this method, make sure to apply the same
# changes to job.QueryJob.to_arrow()
Expand Down
14 changes: 7 additions & 7 deletions packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,7 +193,7 @@ def test_download_arrow_from_job_id_constructs_stream_and_reads(self):
mock_response.arrow_record_batch = None
mock_bqstorage.read_rows.return_value = [mock_response]

with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow:
with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow:
mock_pyarrow.py_buffer = lambda x: x
batches = list(iterator._download_arrow_from_job_id(timeout=5.0))
self.assertEqual(batches, [])
Expand Down Expand Up @@ -247,7 +247,7 @@ def test_download_arrow_from_job_id_with_first_page_response(self):

mock_bqstorage.read_rows.return_value = [mock_response]

with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow:
with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow:
mock_pyarrow.py_buffer = lambda x: x
mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema"
mock_pyarrow.ipc.read_record_batch.side_effect = [
Expand Down Expand Up @@ -306,7 +306,7 @@ def test_download_arrow_from_job_id_with_first_page_response_schema_only(self):

mock_bqstorage.read_rows.return_value = [mock_response]

with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow:
with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow:
mock_pyarrow.py_buffer = lambda x: x
mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema"
mock_pyarrow.ipc.read_record_batch.return_value = mock_stream_batch
Expand Down Expand Up @@ -339,7 +339,7 @@ def test_download_arrow_from_job_id_missing_storage_client_raises_value_error(se
query_results_format="ARROW",
)

with mock.patch("google.cloud.bigquery.table.pyarrow"):
with mock.patch("google.cloud.bigquery._arrow_query.pyarrow"):
with self.assertRaises(ValueError) as ctx:
list(iterator._download_arrow_from_job_id())
self.assertIn(
Expand Down Expand Up @@ -374,7 +374,7 @@ def test_download_arrow_from_job_id_with_schema_and_batch(self):

mock_bqstorage.read_rows.return_value = [mock_response]

with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow:
with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow:
mock_pyarrow.py_buffer = lambda x: x
mock_pyarrow.ipc.read_schema.return_value = "fake_schema"
mock_pyarrow.ipc.read_record_batch.return_value = "fake_batch"
Expand Down Expand Up @@ -422,7 +422,7 @@ def test_download_arrow_from_job_id_avoids_read_rows_when_all_rows_present(self)
mock_first_batch = mock.MagicMock()
mock_first_batch.num_rows = 10

with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow:
with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow:
mock_pyarrow.py_buffer = lambda x: x
mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema"
mock_pyarrow.ipc.read_record_batch.return_value = mock_first_batch
Expand Down Expand Up @@ -464,7 +464,7 @@ def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self):
mock_first_batch.num_rows = 10
mock_bqstorage.read_rows.return_value = []

with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow:
with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow:
mock_pyarrow.py_buffer = lambda x: x
mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema"
mock_pyarrow.ipc.read_record_batch.return_value = mock_first_batch
Expand Down
Loading