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 30f89759ee2b..fe4198cfdda4 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -430,6 +430,8 @@ def query_and_wait( job_retry: Optional[retries.Retry], page_size: Optional[int] = None, max_results: Optional[int] = None, + query_results_format: Optional[str] = None, + compression_codec: Optional[str] = None, callback: Callable = lambda _: None, ) -> table.RowIterator: """Run the query, wait for it to finish, and return the results. @@ -473,8 +475,10 @@ def query_and_wait( page_size (Optional[int]): The maximum number of rows in each page of results from this request. Non-positive values are ignored. - max_results (Optional[int]): - The maximum total number of rows from this request. + query_results_format (Optional[str]): + [Beta] The format for query results (e.g. "ARROW"). + compression_codec (Optional[str]): + [Beta] Compression codec for Arrow serialization (e.g. "LZ4_FRAME"). callback (Callable): A callback function used by bigframes to report query progress. @@ -499,6 +503,13 @@ def query_and_wait( request_body = _to_query_request( query=query, job_config=job_config, location=location, timeout=api_timeout ) + 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 + } # Some API parameters aren't supported by the jobs.query API. In these # cases, fallback to a jobs.insert call. @@ -522,6 +533,7 @@ def query_and_wait( retry=retry, page_size=page_size, max_results=max_results, + query_results_format=query_results_format, callback=callback, ) @@ -594,6 +606,7 @@ def do_query(): retry=retry, page_size=page_size, max_results=max_results, + query_results_format=query_results_format, callback=callback, ) @@ -633,6 +646,7 @@ def do_query(): created=query_results.created, started=query_results.started, ended=query_results.ended, + query_results_format=query_results_format, ) if job_retry is not None: @@ -673,6 +687,7 @@ def _supported_by_jobs_query(request_body: Dict[str, Any]) -> bool: "jobTimeoutMs", "reservation", "maxSlots", + "queryResultsFormat", } unsupported_keys = request_keys - keys_allowlist @@ -687,6 +702,7 @@ def _wait_or_cancel( page_size: Optional[int], max_results: Optional[int], *, + query_results_format: Optional[str] = None, callback: Callable = lambda _: None, ) -> table.RowIterator: """Wait for a job to complete and return the results. @@ -731,6 +747,7 @@ def _wait_or_cancel( ended=job.ended, ) ) + query_results._query_results_format = query_results_format 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/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index ce8768b68b30..1ce47e82ca61 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -3649,6 +3649,8 @@ def query_and_wait( job_retry: retries.Retry = DEFAULT_JOB_RETRY, page_size: Optional[int] = None, max_results: Optional[int] = None, + query_results_format: Optional[str] = None, + compression_codec: Optional[str] = None, ) -> RowIterator: """Run the query, wait for it to finish, and return the results. @@ -3694,8 +3696,10 @@ def query_and_wait( jobs.getQueryResults API calls. Large results downloaded with the BigQuery Storage Read API are intentionally unaffected by this parameter. - max_results (Optional[int]): - The maximum total number of rows from this request. + query_results_format (Optional[str]): + [Beta] The format for query results (e.g. "ARROW"). + compression_codec (Optional[str]): + [Beta] Compression codec for Arrow serialization (e.g. "LZ4_FRAME"). Returns: google.cloud.bigquery.table.RowIterator: @@ -3726,6 +3730,8 @@ def query_and_wait( job_retry=job_retry, page_size=page_size, max_results=max_results, + query_results_format=query_results_format, + compression_codec=compression_codec, ) def _query_and_wait_bigframes( @@ -3741,6 +3747,8 @@ def _query_and_wait_bigframes( job_retry: retries.Retry = DEFAULT_JOB_RETRY, page_size: Optional[int] = None, max_results: Optional[int] = None, + query_results_format: Optional[str] = None, + compression_codec: Optional[str] = None, callback: Callable = lambda _: None, ) -> RowIterator: """See query_and_wait. @@ -3773,6 +3781,8 @@ def _query_and_wait_bigframes( job_retry=job_retry, page_size=page_size, max_results=max_results, + query_results_format=query_results_format, + compression_codec=compression_codec, callback=callback, ) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index 870cdcc5d2ab..2ea1c578e4d4 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -16,6 +16,7 @@ from __future__ import absolute_import +import base64 import copy import datetime import functools @@ -1912,6 +1913,7 @@ def __init__( created: Optional[datetime.datetime] = None, started: Optional[datetime.datetime] = None, ended: Optional[datetime.datetime] = None, + query_results_format: Optional[str] = None, ): super(RowIterator, self).__init__( client, @@ -1945,6 +1947,29 @@ def __init__( self._job_created = created self._job_started = started self._job_ended = ended + self._query_results_format = query_results_format + + @property + def pages(self): + if self._query_results_format == "ARROW": + raise ValueError( + "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead." + ) + return super().pages + + def __iter__(self): + if self._query_results_format == "ARROW": + raise ValueError( + "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead." + ) + return super().__iter__() + + def __next__(self): + if self._query_results_format == "ARROW": + raise ValueError( + "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead." + ) + return super().__next__() @property def _billing_project(self) -> Optional[str]: @@ -2226,6 +2251,12 @@ def to_arrow_iterable( .. versionadded:: 2.31.0 """ + if self._query_results_format == "ARROW": + return self._download_arrow_from_job_id( + bqstorage_client=bqstorage_client, + timeout=timeout, + ) + self._maybe_warn_max_results(bqstorage_client) bqstorage_download = functools.partial( @@ -2251,6 +2282,90 @@ def to_arrow_iterable( bqstorage_client=bqstorage_client, ) + def _download_arrow_from_job_id( + self, + bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None, + timeout: Optional[float] = None, + ) -> Iterator["pyarrow.RecordBatch"]: + if pyarrow is None: + raise ValueError(_NO_PYARROW_ERROR) + + offset = 0 + pa_schema = None + total_rows = self.total_rows + job_complete = False + + 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["totalRows"]) + + 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 + + if job_complete and offset >= total_rows: + return + + 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." + ) + + 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 + # If changing the signature of this method, make sure to apply the same # changes to job.QueryJob.to_arrow() def to_arrow( @@ -2357,7 +2472,7 @@ def to_arrow( # but mypy cannot infer this correlation. We ignore the union-attr error here. bqstorage_client._transport.close() # type: ignore[union-attr] - if record_batches and bqstorage_client is not None: + if record_batches and (bqstorage_client is not None or self._query_results_format == "ARROW"): return pyarrow.Table.from_batches(record_batches) else: # No records (not record_batches), use schema based on BigQuery schema @@ -3036,17 +3151,41 @@ class _EmptyRowIterator(RowIterator): """ def __init__( - self, client=None, api_request=None, path=None, schema=(), *args, **kwargs + self, client=None, api_request=None, path=None, schema=(), *args, query_results_format: Optional[str] = None, **kwargs ): super().__init__( client=client, api_request=api_request, path=path, schema=schema, + query_results_format=query_results_format, *args, **kwargs, ) self._total_rows = 0 + self._query_results_format = query_results_format + + @property + def pages(self): + if self._query_results_format == "ARROW": + raise ValueError( + "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead." + ) + return super().pages + + def __iter__(self): + if self._query_results_format == "ARROW": + raise ValueError( + "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead." + ) + return iter(()) + + def __next__(self): + if self._query_results_format == "ARROW": + raise ValueError( + "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead." + ) + raise StopIteration def to_arrow( self, @@ -3219,11 +3358,10 @@ def to_arrow_iterable( Returns: An iterator yielding a single empty :class:`~pyarrow.RecordBatch`. """ + if pyarrow is None: + raise ValueError(_NO_PYARROW_ERROR) return iter((pyarrow.record_batch([]),)) - def __iter__(self): - return iter(()) - class PartitionRange(object): """Definition of the ranges for range partitioning. 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 new file mode 100644 index 000000000000..9ed0c0ff2d60 --- /dev/null +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -0,0 +1,471 @@ +# 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 +# +# https://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. + +import base64 +import unittest +from unittest import mock + +from google.cloud.bigquery import _job_helpers +from google.cloud.bigquery.client import Client +from google.cloud.bigquery.table import RowIterator, _EmptyRowIterator + + +class TestQueryResultsFormatOption1(unittest.TestCase): + def test_supported_by_jobs_query_includes_query_results_format(self): + body = {"query": "SELECT 1", "queryResultsFormat": "ARROW"} + self.assertTrue(_job_helpers._supported_by_jobs_query(body)) + + def test_job_helpers_query_and_wait_sets_request_body(self): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "rows": [], + "schema": {"fields": []}, + } + + row_iterator = _job_helpers.query_and_wait( + client=client, + query="SELECT 1", + project="p", + location="us", + job_config=None, + retry=None, + job_retry=None, + query_results_format="ARROW", + ) + + self.assertEqual(row_iterator._query_results_format, "ARROW") + call_args = client._call_api.call_args + self.assertIn("queryResultsFormat", call_args.kwargs["data"]) + self.assertEqual(call_args.kwargs["data"]["queryResultsFormat"], "ARROW") + + def test_job_helpers_query_and_wait_sets_compression_codec(self): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "rows": [], + "schema": {"fields": []}, + } + + _job_helpers.query_and_wait( + client=client, + query="SELECT 1", + project="p", + location="us", + job_config=None, + retry=None, + job_retry=None, + query_results_format="ARROW", + compression_codec="LZ4_FRAME", + ) + + call_args = client._call_api.call_args + self.assertIn("formatOptions", call_args.kwargs["data"]) + self.assertEqual( + call_args.kwargs["data"]["formatOptions"]["arrowSerializationOptions"][ + "bufferCompression" + ], + "LZ4_FRAME", + ) + + def test_job_helpers_query_and_wait_fallback_preserves_query_results_format(self): + client = mock.MagicMock(spec=Client) + unsupported_body = {"unsupportedKey": "val"} + job_mock = mock.MagicMock() + fake_iterator = mock.MagicMock(spec=RowIterator) + job_mock.result.return_value = fake_iterator + + with mock.patch.object( + _job_helpers, "_to_query_request", return_value=unsupported_body + ), mock.patch.object( + _job_helpers, "query_jobs_insert", return_value=job_mock + ): + res_iterator = _job_helpers.query_and_wait( + client=client, + query="SELECT 1", + project="p", + location="us", + job_config=None, + retry=None, + job_retry=None, + query_results_format="ARROW", + ) + self.assertEqual(res_iterator._query_results_format, "ARROW") + + def test_row_iterator_non_arrow_iteration_raises_value_error(self): + iterator = RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=(), + query_results_format="ARROW", + ) + + with self.assertRaises(ValueError) as ctx: + iter(iterator) + self.assertIn("Cannot iterate over non-arrow results", str(ctx.exception)) + + with self.assertRaises(ValueError) as ctx: + next(iterator) + self.assertIn("Cannot iterate over non-arrow results", str(ctx.exception)) + + with self.assertRaises(ValueError) as ctx: + _ = iterator.pages + self.assertIn("Cannot iterate over non-arrow results", str(ctx.exception)) + + def test_empty_row_iterator_non_arrow_iteration_raises_value_error(self): + iterator = _EmptyRowIterator(query_results_format="ARROW") + + with self.assertRaises(ValueError) as ctx: + iter(iterator) + self.assertIn("Cannot iterate over non-arrow results", str(ctx.exception)) + + with self.assertRaises(ValueError) as ctx: + next(iterator) + self.assertIn("Cannot iterate over non-arrow results", str(ctx.exception)) + + with self.assertRaises(ValueError) as ctx: + _ = iterator.pages + self.assertIn("Cannot iterate over non-arrow results", str(ctx.exception)) + + def test_row_iterator_standard_format_allows_iteration(self): + iterator = RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=(), + query_results_format=None, + ) + self.assertIsNotNone(iterator.pages) + + def test_row_iterator_to_arrow_iterable_delegates_when_format_is_arrow(self): + iterator = RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=(), + project="proj", + location="loc", + job_id="job123", + query_results_format="ARROW", + ) + + with mock.patch.object( + iterator, "_download_arrow_from_job_id", return_value=iter(["batch1", "batch2"]) + ) as mock_download: + res = list(iterator.to_arrow_iterable(timeout=10.0)) + self.assertEqual(res, ["batch1", "batch2"]) + mock_download.assert_called_once_with(bqstorage_client=None, timeout=10.0) + + def test_download_arrow_from_job_id_constructs_stream_and_reads(self): + mock_client = mock.MagicMock() + mock_bqstorage = mock.MagicMock() + mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + + iterator = RowIterator( + client=mock_client, + api_request=mock.MagicMock(), + path=None, + schema=(), + project="test-proj", + location="US", + job_id="test-job-456", + query_results_format="ARROW", + ) + + mock_response = mock.MagicMock() + mock_response.arrow_schema = None + 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: + mock_pyarrow.py_buffer = lambda x: x + batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) + + expected_stream_name = "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) + + def test_download_arrow_from_job_id_with_first_page_response(self): + mock_client = mock.MagicMock() + mock_bqstorage = mock.MagicMock() + mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + + raw_schema_bytes = b"schema_bytes_123" + raw_batch_bytes = b"batch_bytes_123" + b64_schema = base64.b64encode(raw_schema_bytes).decode("ascii") + b64_batch = base64.b64encode(raw_batch_bytes).decode("ascii") + + first_page_response = { + "arrowSchema": {"serializedSchema": b64_schema}, + "arrowRecordBatch": {"serializedRecordBatch": b64_batch}, + } + + iterator = RowIterator( + client=mock_client, + api_request=mock.MagicMock(), + path=None, + schema=(), + project="test-proj", + location="US", + job_id="test-job-sync", + query_results_format="ARROW", + first_page_response=first_page_response, + ) + + mock_first_batch = mock.MagicMock() + mock_first_batch.num_rows = 10 + + mock_second_batch = mock.MagicMock() + mock_second_batch.num_rows = 5 + + # Subsequent ReadRows response chunk without arrow_schema + mock_response = mock.MagicMock() + mock_response.arrow_schema = None + mock_batch_msg = mock.MagicMock() + mock_batch_msg.serialized_record_batch = b"stream_batch_bytes" + mock_response.arrow_record_batch = mock_batch_msg + + mock_bqstorage.read_rows.return_value = [mock_response] + + with mock.patch("google.cloud.bigquery.table.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 = [ + mock_first_batch, + mock_second_batch, + ] + + batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) + + self.assertEqual(batches, [mock_first_batch, mock_second_batch]) + self.assertIsNone(iterator._first_page_response) + + mock_pyarrow.ipc.read_schema.assert_called_once_with(raw_schema_bytes) + mock_pyarrow.ipc.read_record_batch.assert_has_calls( + [ + mock.call(raw_batch_bytes, "deserialized_schema"), + mock.call(b"stream_batch_bytes", "deserialized_schema"), + ] + ) + + expected_stream_name = ( + "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 + ) + + def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): + mock_client = mock.MagicMock() + mock_bqstorage = mock.MagicMock() + mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + + raw_schema_bytes = b"schema_bytes_456" + first_page_response = { + "arrowSchema": {"serializedSchema": raw_schema_bytes}, + } + + iterator = RowIterator( + client=mock_client, + api_request=mock.MagicMock(), + path=None, + schema=(), + project="test-proj", + location="US", + job_id="test-job-schema-only", + query_results_format="ARROW", + first_page_response=first_page_response, + ) + + mock_stream_batch = mock.MagicMock() + mock_response = mock.MagicMock() + mock_response.arrow_schema = None + mock_batch_msg = mock.MagicMock() + mock_batch_msg.serialized_record_batch = b"stream_batch_bytes" + mock_response.arrow_record_batch = mock_batch_msg + + mock_bqstorage.read_rows.return_value = [mock_response] + + with mock.patch("google.cloud.bigquery.table.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 + + batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) + + self.assertEqual(batches, [mock_stream_batch]) + self.assertIsNone(iterator._first_page_response) + mock_pyarrow.ipc.read_schema.assert_called_once_with(raw_schema_bytes) + mock_pyarrow.ipc.read_record_batch.assert_called_once_with( + b"stream_batch_bytes", "deserialized_schema" + ) + 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 + ) + + def test_download_arrow_from_job_id_missing_storage_client_raises_value_error(self): + mock_client = mock.MagicMock() + mock_client._ensure_bqstorage_client.return_value = None + + iterator = RowIterator( + client=mock_client, + api_request=mock.MagicMock(), + path=None, + schema=(), + project="test-proj", + location="US", + job_id="test-job-456", + query_results_format="ARROW", + ) + + with mock.patch("google.cloud.bigquery.table.pyarrow"): + with self.assertRaises(ValueError) as ctx: + list(iterator._download_arrow_from_job_id()) + self.assertIn("The google-cloud-bigquery-storage library is required", str(ctx.exception)) + + def test_download_arrow_from_job_id_with_schema_and_batch(self): + mock_client = mock.MagicMock() + mock_bqstorage = mock.MagicMock() + mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + + iterator = RowIterator( + client=mock_client, + api_request=mock.MagicMock(), + path=None, + schema=(), + project="test-proj", + location="US", + job_id="test-job-789", + query_results_format="ARROW", + ) + + mock_schema_msg = mock.MagicMock() + mock_schema_msg.serialized_schema = b"schema_bytes" + mock_batch_msg = mock.MagicMock() + mock_batch_msg.serialized_record_batch = b"batch_bytes" + + mock_response = mock.MagicMock() + mock_response.arrow_schema = mock_schema_msg + mock_response.arrow_record_batch = mock_batch_msg + + mock_bqstorage.read_rows.return_value = [mock_response] + + with mock.patch("google.cloud.bigquery.table.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" + + batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) + self.assertEqual(batches, ["fake_batch"]) + mock_pyarrow.ipc.read_schema.assert_called_once_with(b"schema_bytes") + mock_pyarrow.ipc.read_record_batch.assert_called_once_with(b"batch_bytes", "fake_schema") + + def test_empty_row_iterator_to_arrow_iterable_checks_pyarrow(self): + iterator = _EmptyRowIterator(query_results_format="ARROW") + with mock.patch("google.cloud.bigquery.table.pyarrow", None): + with self.assertRaises(ValueError) as ctx: + iterator.to_arrow_iterable() + self.assertIn("pyarrow", str(ctx.exception).lower()) + + def test_download_arrow_from_job_id_avoids_read_rows_when_all_rows_present(self): + mock_client = mock.MagicMock() + raw_schema_bytes = b"schema_bytes_789" + raw_batch_bytes = b"batch_bytes_789" + b64_schema = base64.b64encode(raw_schema_bytes).decode("ascii") + b64_batch = base64.b64encode(raw_batch_bytes).decode("ascii") + + first_page_response = { + "jobComplete": True, + "totalRows": "10", + "arrowSchema": {"serializedSchema": b64_schema}, + "arrowRecordBatch": {"serializedRecordBatch": b64_batch}, + } + + iterator = RowIterator( + client=mock_client, + api_request=mock.MagicMock(), + path=None, + schema=(), + project="test-proj", + location="US", + job_id="test-job-complete", + query_results_format="ARROW", + first_page_response=first_page_response, + ) + + mock_first_batch = mock.MagicMock() + mock_first_batch.num_rows = 10 + + with mock.patch("google.cloud.bigquery.table.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 + + batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) + self.assertEqual(batches, [mock_first_batch]) + mock_client._ensure_bqstorage_client.assert_not_called() + + def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self): + mock_client = mock.MagicMock() + mock_bqstorage = mock.MagicMock() + mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + + raw_schema_bytes = b"schema_bytes_789" + raw_batch_bytes = b"batch_bytes_789" + b64_schema = base64.b64encode(raw_schema_bytes).decode("ascii") + b64_batch = base64.b64encode(raw_batch_bytes).decode("ascii") + + first_page_response = { + "jobComplete": False, + "totalRows": "10", + "arrowSchema": {"serializedSchema": b64_schema}, + "arrowRecordBatch": {"serializedRecordBatch": b64_batch}, + } + + iterator = RowIterator( + client=mock_client, + api_request=mock.MagicMock(), + path=None, + schema=(), + project="test-proj", + location="US", + job_id="test-job-incomplete", + query_results_format="ARROW", + first_page_response=first_page_response, + ) + + mock_first_batch = mock.MagicMock() + mock_first_batch.num_rows = 10 + mock_bqstorage.read_rows.return_value = [] + + with mock.patch("google.cloud.bigquery.table.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 + + batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) + self.assertEqual(batches, [mock_first_batch]) + 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) + + +if __name__ == "__main__": + unittest.main() +