diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py new file mode 100644 index 000000000000..25738a2e818d --- /dev/null +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -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" + 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 diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index d1d31b14ffd8..14f996b894f2 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -16,7 +16,6 @@ from __future__ import absolute_import -import base64 import copy import datetime import functools @@ -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, @@ -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() 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 3bc19b4d3ccd..1e22381419c7 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 @@ -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, []) @@ -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 = [ @@ -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 @@ -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( @@ -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" @@ -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 @@ -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