From fc79738af2b8a2b548e678db1c59938aa2d34d41 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:41:05 +0000 Subject: [PATCH] fix(bigquery): keep the Arrow first page and wait without rows When an ARROW query's results didn't fit in the jobs.query response, QueryJob.result() dropped that response, because it only kept first pages that had STRUCT_ENCODING rows. The stream then reread every row from offset 0. While waiting for the job, query_and_wait() also passed page_size to jobs.getQueryResults, which returned STRUCT_ENCODING rows that ARROW results never use. For ARROW queries, decide whether rows remain by comparing pageRowCount with totalRows instead of checking for a pageToken, wait with page_size=0 so jobs.getQueryResults is called with maxResults=0, and keep a jobs.query response that has arrowRecordBatch as the first page. RowIterator then returns the first page and streams the rest from the next row, also when it's read again. Queries that don't use ARROW are unchanged. --- .../google/cloud/bigquery/_arrow_query.py | 13 +- .../google/cloud/bigquery/_job_helpers.py | 10 +- .../google/cloud/bigquery/job/query.py | 8 +- .../tests/system/test_arrow.py | 314 ++++++++++++++++++ .../tests/system/test_query.py | 35 +- .../tests/unit/job/test_query.py | 36 ++ .../unit/test_query_results_format_arrow.py | 75 ++++- 7 files changed, 482 insertions(+), 9 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index 84c9fc853758..c15891fc064b 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -17,7 +17,7 @@ from __future__ import annotations import base64 -from typing import Dict, Iterator, Optional, TYPE_CHECKING +from typing import Any, Dict, Iterator, Optional, TYPE_CHECKING from google.cloud.bigquery import _versions_helpers from google.cloud.bigquery import enums @@ -113,6 +113,16 @@ def serialization_options( } +def has_more_rows(response: Dict[str, Any]) -> bool: + """Return True if rows remain after an ``ARROW`` ``jobs.query`` response. + + ``ARROW`` responses say how many rows the first page holds in + ``pageRowCount``, so compare that with ``totalRows`` instead of checking + for a ``pageToken``. + """ + return int(response.get("pageRowCount", 0)) < int(response.get("totalRows", 0)) + + def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]: """Convert ``jobs.query`` ``arrowSerializationOptions`` to the ``ArrowSerializationOptions`` fields that ``ReadRows`` takes.""" @@ -153,7 +163,6 @@ def download_arrow_from_job( # 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: 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 8470c19becd2..f5080611a081 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -533,6 +533,10 @@ def query_and_wait( ) request_body["arrowSerializationOptions"] = arrow_serialization_options + # jobs.getQueryResults only returns STRUCT_ENCODING rows, so while waiting + # for an ARROW query, request none of them. + wait_page_size = 0 if arrow_serialization_options is not None else page_size + # Some API parameters aren't supported by the jobs.query API. In these # cases, fallback to a jobs.insert call. if not _supported_by_jobs_query(request_body): @@ -553,7 +557,7 @@ def query_and_wait( api_timeout=api_timeout, wait_timeout=wait_timeout, retry=retry, - page_size=page_size, + page_size=wait_page_size, max_results=max_results, query_results_format=query_results_format, arrow_serialization_options=arrow_serialization_options, @@ -616,6 +620,8 @@ def do_query(): ) page_token = query_results.page_token more_pages = page_token is not None + if arrow_serialization_options is not None: + more_pages = _arrow_query.has_more_rows(response) if more_pages or not query_results.complete: # TODO(swast): Avoid a call to jobs.get in some cases (few @@ -627,7 +633,7 @@ def do_query(): api_timeout=api_timeout, wait_timeout=wait_timeout, retry=retry, - page_size=page_size, + page_size=wait_page_size, max_results=max_results, query_results_format=query_results_format, arrow_serialization_options=arrow_serialization_options, diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py index e85e43f6a164..0b7af6293851 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py @@ -1842,9 +1842,13 @@ def is_job_done(): # RowIterator response if there are any rows in it. This prevents us # from stopping the iteration early in the cases where we set # maxResults=0. In that case, we're missing rows and there's no next - # page token. + # page token. ARROW responses from jobs.query carry their rows in + # arrowRecordBatch instead of rows. first_page_response = self._query_results._properties - if "rows" not in first_page_response: + if ( + "rows" not in first_page_response + and "arrowRecordBatch" not in first_page_response + ): first_page_response = None rows = self._client._list_rows_from_query_results( diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index 4c4bf4fcf583..78fb935124f7 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -15,10 +15,15 @@ """System tests for Arrow connector.""" import datetime +import json +import math from typing import Optional +from unittest import mock +import uuid import pyarrow import pytest +from dateutil import relativedelta from google.api_core import exceptions from google.cloud import bigquery @@ -412,3 +417,312 @@ def test_query_and_wait_arrow_timestamp_precision_picosecond_table( assert all(".123456789012" in value for value in values) else: assert values == _MICROSECOND_VALUES + + +@pytest.mark.parametrize( + ("page_size", "priority", "expected_offset"), + [ + (100, None, 100), + # jobs.query doesn't accept a priority, so this case runs through + # jobs.insert, which returns no rows inline. + (None, enums.QueryPriority.INTERACTIVE, 0), + ], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_matches_struct_encoding( + bigquery_client, bqstorage_client, page_size, priority, expected_offset +): + query = ( + "SELECT x, CAST(x AS STRING) AS s, MOD(x, 2) = 0 AS even, " + "DATE_ADD(DATE '2026-01-01', INTERVAL x DAY) AS d " + "FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + # Wrap read_rows() to check where the stream starts. + read_rows = mock.Mock(wraps=bqstorage_client.read_rows) + + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + page_size=page_size, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + with mock.patch.object(bqstorage_client, "read_rows", read_rows): + batches = list(rows.to_arrow_iterable(bqstorage_client=bqstorage_client)) + + assert pyarrow.Table.from_batches(batches).to_pylist() == expected + read_rows.assert_called_once_with( + f"projects/{bigquery_client.project}/locations/{rows.location}" + f"/jobs/{rows.job_id}/streams/_default", + offset=expected_offset, + timeout=mock.ANY, + arrow_serialization_options=mock.ANY, + ) + + +def _interval_to_relativedelta(value): + return relativedelta.relativedelta( + months=value.months, days=value.days, microseconds=value.nanoseconds // 1000 + ) + + +def _arrow_rows(table): + """Return the rows of an Arrow table as STRUCT_ENCODING returns them. + + Arrow has JSON values as strings and INTERVAL values as MonthDayNano, + while STRUCT_ENCODING rows have parsed JSON and relativedelta values. + """ + converters = {} + for field in table.schema: + if field.metadata == {b"ARROW:extension:name": b"google:sqlType:json"}: + converters[field.name] = json.loads + elif pyarrow.types.is_interval(field.type): + converters[field.name] = _interval_to_relativedelta + rows = table.to_pylist() + for row in rows: + for name, convert in converters.items(): + if row[name] is not None: + row[name] = convert(row[name]) + return rows + + +@pytest.mark.parametrize( + ("page_size", "priority", "expected_offset"), + [ + (1, None, 1), + # jobs.query doesn't accept a priority, so this case runs through + # jobs.insert, which returns no rows inline. + (None, enums.QueryPriority.INTERACTIVE, 0), + ], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_locations( + bigquery_client, + bqstorage_client, + scalars_table_multi_location, + page_size, + priority, + expected_offset, +): + location, table_id = scalars_table_multi_location + query = f"SELECT * FROM `{table_id}` ORDER BY rowindex" + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + read_rows = mock.Mock(wraps=bqstorage_client.read_rows) + + # The query doesn't name a location, so only jobReference has it. + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + page_size=page_size, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + with mock.patch.object(bqstorage_client, "read_rows", read_rows): + batches = list(rows.to_arrow_iterable(bqstorage_client=bqstorage_client)) + + assert rows.location == location + assert _arrow_rows(pyarrow.Table.from_batches(batches)) == expected + read_rows.assert_called_once_with( + f"projects/{bigquery_client.project}/locations/{location}" + f"/jobs/{rows.job_id}/streams/_default", + offset=expected_offset, + timeout=mock.ANY, + arrow_serialization_options=mock.ANY, + ) + + +_DATA_TYPES_QUERY = """ +SELECT + i, + IF(i = 2, NULL, i = 1) AS bool_col, + IF(i = 2, NULL, i * 1000000000000) AS int64_col, + IF(i = 2, NULL, i / 3) AS float64_col, + IF(i = 2, NULL, NUMERIC '-123456789.123456789') AS numeric_col, + IF(i = 2, NULL, BIGNUMERIC '1.23456789012345678901234567890123456789') AS bignumeric_col, + IF(i = 2, NULL, FORMAT('row %d', i)) AS string_col, + IF(i = 2, NULL, CAST(FORMAT('row %d', i) AS BYTES)) AS bytes_col, + IF(i = 2, NULL, DATE_ADD(DATE '2026-01-01', INTERVAL i DAY)) AS date_col, + IF(i = 2, NULL, DATETIME '2026-01-01 12:34:56.789012') AS datetime_col, + IF(i = 2, NULL, TIME '12:34:56.789012') AS time_col, + IF(i = 2, NULL, TIMESTAMP '2026-01-01 12:34:56.789012 UTC') AS timestamp_col, + IF(i = 2, NULL, INTERVAL '1-2 3 4:5:6.789012' YEAR TO SECOND) AS interval_col, + IF(i = 2, NULL, ST_GEOGPOINT(-122.0838511, 37.3860517)) AS geography_col, + IF(i = 2, NULL, JSON '{"a": [1, 2.5, "x"], "b": null, "c": {"d": true}}') AS json_col, + IF(i = 2, NULL, RANGE(DATE '2026-01-01', DATE '2026-02-01')) AS range_date_col, + IF(i = 2, NULL, RANGE(DATETIME '2026-01-01 00:00:00', NULL)) AS range_datetime_col, + IF( + i = 2, + NULL, + RANGE(TIMESTAMP '2026-01-01 00:00:00 UTC', TIMESTAMP '2026-01-02 00:00:00 UTC') + ) AS range_timestamp_col, + IF(i = 2, [], [i, i + 1]) AS array_col, + IF(i = 2, NULL, STRUCT(i AS x, [STRUCT('a' AS s, [1.5, 2.5] AS f)] AS nested)) AS struct_col +FROM UNNEST([1, 2, 3]) AS i +ORDER BY i +""" + + +def test_query_and_wait_arrow_data_types(bigquery_client): + job_config = bigquery.QueryJobConfig(use_query_cache=False) + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait( + _DATA_TYPES_QUERY, job_config=job_config + ) + ] + + rows = bigquery_client.query_and_wait( + _DATA_TYPES_QUERY, + job_config=job_config, + # The first row arrives in the first page, the others in the stream. + page_size=1, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + batches = list(rows.to_arrow_iterable()) + + assert batches[0].num_rows == 1 + assert _arrow_rows(pyarrow.Table.from_batches(batches)) == expected + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["jobs_query", "jobs_insert"], +) +def test_query_and_wait_arrow_unnamed_struct_fields(bigquery_client, priority): + query = "SELECT (1, 'a') AS s, STRUCT(2 AS x, 'b') AS p" + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + # jobs.query doesn't accept a priority, so this runs through jobs.insert. + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + + try: + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + # jobs.query fails with a jobInternalError, which job_retry would + # retry until its deadline. + job_retry=None, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + except exceptions.BadRequest as exc: + if not any(error.get("reason") == "jobInternalError" for error in exc.errors): + raise + pytest.xfail(f"jobs.query can't send unnamed STRUCT fields: {exc.message}") + table = pyarrow.Table.from_batches(list(rows.to_arrow_iterable())) + + assert table.to_pylist() == expected + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["first_page_only", "no_rows_inline"], +) +def test_query_and_wait_arrow_extreme_values( + bigquery_client, scalars_extreme_table, priority +): + # BigQuery fails the extreme INTERVAL values with "Out of range conversion + # for microseconds value", since Arrow's interval type can't hold them. + query = ( + f"SELECT * EXCEPT (interval_col) FROM `{scalars_extreme_table}` " + "ORDER BY rowindex" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + actual = _arrow_rows(pyarrow.Table.from_batches(list(rows.to_arrow_iterable()))) + + # NaN never equals itself, so check it on its own. + assert math.isnan(actual[2].pop("float64_col")) + assert math.isnan(expected[2].pop("float64_col")) + assert actual == expected + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_cached_results(bigquery_client, priority): + # No other query uses this literal, so only the second run can hit the + # cache. + query = ( + f"SELECT x, '{uuid.uuid4().hex}' AS key " + "FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=True) + if priority is not None: + job_config.priority = priority + + def read(): + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + page_size=100, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + table = pyarrow.Table.from_batches(list(rows.to_arrow_iterable())) + return bigquery_client.get_job(rows.job_id, location=rows.location), table + + first_job, first = read() + cached_job, cached = read() + + assert not first_job.cache_hit + assert cached_job.cache_hit + assert cached.num_rows == 1000 + assert cached.equals(first) + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["jobs_query", "jobs_insert"], +) +def test_query_and_wait_arrow_failed_query(bigquery_client, priority): + # The last row fails, so the query fails while it runs. + query = ( + "SELECT IF(x < 1000, x, ERROR(FORMAT('failed at row %d', x))) AS x " + "FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + + with pytest.raises(exceptions.GoogleAPICallError) as struct_encoding_error: + bigquery_client.query_and_wait(query, job_config=job_config) + with pytest.raises(exceptions.GoogleAPICallError) as arrow_error: + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + list(rows.to_arrow_iterable()) + + assert type(arrow_error.value) is type(struct_encoding_error.value) + assert "failed at row 1000" in arrow_error.value.message diff --git a/packages/google-cloud-bigquery/tests/system/test_query.py b/packages/google-cloud-bigquery/tests/system/test_query.py index 69ff08b66073..43ce8d293403 100644 --- a/packages/google-cloud-bigquery/tests/system/test_query.py +++ b/packages/google-cloud-bigquery/tests/system/test_query.py @@ -14,6 +14,7 @@ import datetime import decimal +import json from typing import Tuple import pytest @@ -162,7 +163,7 @@ def test_query_statistics(bigquery_client, query_api_method): assert len(plan) > stages_with_inputs -@pytest.mark.parametrize( +_QUERY_PARAMETER_CASES = pytest.mark.parametrize( ("sql", "expected", "query_parameters"), ( ( @@ -438,6 +439,9 @@ def test_query_statistics(bigquery_client, query_api_method): ), ), ) + + +@_QUERY_PARAMETER_CASES def test_query_parameters( bigquery_client, query_api_method, sql, expected, query_parameters ): @@ -454,6 +458,35 @@ def test_query_parameters( assert rows[0][0] == expected +@_QUERY_PARAMETER_CASES +def test_query_parameters_arrow(bigquery_client, sql, expected, query_parameters): + job_config = bigquery.QueryJobConfig( + query_parameters=query_parameters, use_query_cache=False + ) + try: + rows = bigquery_client.query_and_wait( + sql, + job_config=job_config, + # jobs.query fails unnamed STRUCT fields with a jobInternalError, + # which job_retry would retry until its deadline. + job_retry=None, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + except exceptions.BadRequest as exc: + if not any(error.get("reason") == "jobInternalError" for error in exc.errors): + raise + pytest.xfail(f"jobs.query can't send unnamed STRUCT fields: {exc.message}") + (batch,) = rows.to_arrow_iterable() + assert batch.num_rows == 1 + assert batch.num_columns == 1 + (value,) = batch.column(0).to_pylist() + if batch.schema.field(0).metadata == { + b"ARROW:extension:name": b"google:sqlType:json" + }: + value = json.loads(value) + assert value == expected + + def test_dry_run( bigquery_client: bigquery.Client, query_api_method: str, diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_query.py b/packages/google-cloud-bigquery/tests/unit/job/test_query.py index d6727ae0355e..daa3d4bd4dc7 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_query.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_query.py @@ -1279,6 +1279,42 @@ def test_result_with_done_jobs_query_response_doesnt_call_get_query_results(self timeout=DEFAULT_GET_JOB_TIMEOUT, ) + def test_result_with_done_jobs_query_response_keeps_arrow_first_page(self): + """An ARROW response from jobs.query is kept as the first page, like + rows are, so the stream can continue after it. + """ + job_resource = self._make_resource(started=True, ended=True, location="EU") + conn = make_connection(job_resource) + client = _make_client(self.PROJECT, connection=conn) + query_resource_done = { + "jobComplete": True, + "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID}, + "arrowSchema": {"serializedSchema": "c2NoZW1h"}, + "arrowRecordBatch": {"serializedRecordBatch": "YmF0Y2g="}, + "pageRowCount": "1", + "totalRows": "2", + } + job = google.cloud.bigquery._job_helpers._to_query_job( + client, + "SELECT 1", + request_config=None, + query_response=query_resource_done, + ) + + result = job.result(page_size=0) + + self.assertEqual( + result._first_page_response["arrowRecordBatch"], + query_resource_done["arrowRecordBatch"], + ) + # Only jobs.get, no jobs.getQueryResults. + conn.api_request.assert_called_once_with( + method="GET", + path=f"/projects/{self.PROJECT}/jobs/{self.JOB_ID}", + query_params={"projection": "full"}, + timeout=DEFAULT_GET_JOB_TIMEOUT, + ) + def test_result_with_none_timeout(self): # Verifies that with an intentional None timeout, get job uses None # instead of the default timeout. 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 e81517c0bf19..d8446a3eefc8 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 @@ -295,7 +295,8 @@ def test_download_arrow_from_job_id_with_first_page_response(self): 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) + # The first page is kept for the next read. + self.assertIsNotNone(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( @@ -354,7 +355,8 @@ def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_stream_batch]) - self.assertIsNone(iterator._first_page_response) + # The first page is kept for the next read. + self.assertIsNotNone(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" @@ -473,6 +475,8 @@ def test_download_arrow_from_job_id_avoids_read_rows_when_all_rows_present(self) 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() + batches_second = list(iterator._download_arrow_from_job_id(timeout=5.0)) + self.assertEqual(batches_second, [mock_first_batch]) def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self): mock_client = mock.MagicMock() @@ -843,6 +847,7 @@ def test_query_and_wait_arrow_jobs_insert_omits_arrow_fields(): arrow_serialization_options=_DEFAULT_QUERY_OPTIONS, callback=mock.ANY, ) + assert wait_or_cancel.call_args.kwargs["page_size"] == 0 def test_supported_by_jobs_query_includes_arrow_serialization_options(): @@ -955,5 +960,71 @@ def test_client_query_and_wait_arrow_timestamp_precision(timestamp_precision, ex assert "timestampOutputFormat" not in data["formatOptions"] +def _arrow_query_response(**fields): + return { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "arrowSchema": {"serializedSchema": "c2NoZW1h"}, + **fields, + } + + +@pytest.mark.parametrize( + "response", + [ + _arrow_query_response(pageRowCount="10", totalRows="25"), + _arrow_query_response(jobComplete=False), + ], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_waits_without_fetching_rows(response): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = response + + with ( + mock.patch.object(_job_helpers, "_to_query_job") as to_query_job, + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, + ): + rows = _query_and_wait(client, query_results_format="ARROW", page_size=10) + + assert rows is wait_or_cancel.return_value + assert wait_or_cancel.call_args.args[0] is to_query_job.return_value + assert wait_or_cancel.call_args.kwargs["page_size"] == 0 + # page_size still sets the size of the first page from jobs.query. + assert client._call_api.call_args.kwargs["data"]["maxResults"] == 10 + + +def test_query_and_wait_arrow_all_rows_inline_ignores_page_token(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = _arrow_query_response( + pageRowCount="25", totalRows="25", pageToken="next" + ) + + with mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel: + rows = _query_and_wait(client, query_results_format="ARROW") + + wait_or_cancel.assert_not_called() + assert rows._first_page_response["pageRowCount"] == "25" + + +def test_query_and_wait_without_arrow_waits_with_page_size(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "rows": [], + "schema": {"fields": []}, + "pageToken": "next", + } + + with ( + mock.patch.object(_job_helpers, "_to_query_job"), + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, + ): + _query_and_wait(client, page_size=10) + + assert wait_or_cancel.call_args.kwargs["page_size"] == 10 + + if __name__ == "__main__": unittest.main()