From 737d06d963d48861c882dfcb1b8de9ac42a50c6b Mon Sep 17 00:00:00 2001 From: ohmayr Date: Mon, 2 Nov 2026 18:50:48 +0000 Subject: [PATCH 1/2] fix(bigquery): read cached ARROW results too large to return inline For a cached query whose results are too large to return inline, jobs.query returns jobComplete, a pageToken, and no rows. BigQuery no longer sends totalRows for these responses, so routing them through QueryJob.result() returned an _EmptyRowIterator and query_and_wait silently returned 0 rows. ARROW rows never come from jobs.getQueryResults, so wait on the job only while the query runs. The RowIterator follows pageToken to the job's ReadRows stream. Fixes #18935 --- .../google/cloud/bigquery/_arrow_query.py | 12 +-- .../google/cloud/bigquery/_job_helpers.py | 10 ++- .../tests/system/test_arrow.py | 2 +- .../unit/test_query_results_format_arrow.py | 78 +++++++++++++++---- 4 files changed, 74 insertions(+), 28 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 8c0c0649af12..7508e7a0836f 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -114,14 +114,6 @@ def serialization_options( } -def reached_max_results(response: Dict[str, Any], max_results: Optional[int]) -> bool: - """Return True if an ``ARROW`` ``jobs.query`` response's first page - already holds ``max_results`` rows.""" - return ( - max_results is not None and int(response.get("pageRowCount", 0)) >= max_results - ) - - def _slice_batch( batch: "pyarrow.RecordBatch", remaining: Optional[int] ) -> Tuple["pyarrow.RecordBatch", Optional[int]]: @@ -295,8 +287,8 @@ def download_arrow_from_job( # Step 2: Return early if all results were delivered in the first page # response, or if the first page already reached max_results. Only a - # pageToken says rows remain: totalRows is 0 when cached results are too - # large to return inline. + # pageToken says rows remain: BigQuery leaves out totalRows when cached + # results are too large to return inline. if (job_complete and page_token is None) or remaining == 0: if not yielded: yield _empty_batch(pa_schema, row_iterator._schema) 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 c222e0c7fda4..9d0ae908106c 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -627,10 +627,12 @@ def do_query(): ) page_token = query_results.page_token more_pages = page_token is not None - if ( - arrow_serialization_options is not None - and _arrow_query.reached_max_results(response, max_results) - ): + if arrow_serialization_options is not None: + # ARROW rows never come from jobs.getQueryResults: the RowIterator + # follows pageToken to the job's ReadRows stream, so only wait + # while the query runs. QueryJob.result() would also treat a + # missing totalRows, which BigQuery leaves out for cached results + # too large to return inline, as a statement without rows. more_pages = False if more_pages or not query_results.complete: diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index 59b7f82645e5..4c2740837bed 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -669,7 +669,7 @@ def test_query_and_wait_arrow_extreme_values( [ (None, 100, 1000, 0), # The results are too large to return inline, so the cached jobs.query - # response has no rows and totalRows is 0. + # response has no rows and no totalRows. (None, None, 20000, 1024), (enums.QueryPriority.INTERACTIVE, 100, 1000, 0), ], 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 246f398c08b6..387e458a4d08 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 @@ -1083,20 +1083,9 @@ def _arrow_query_response(**fields): } -@pytest.mark.parametrize( - "response", - [ - _arrow_query_response(pageRowCount="10", totalRows="25", pageToken="next"), - _arrow_query_response(jobComplete=False), - # A cache hit too large to return inline has no pageRowCount and - # totalRows is 0, but it still has a pageToken. - _arrow_query_response(totalRows="0", pageToken="next"), - ], - ids=["some_rows_inline", "no_rows_inline", "cache_hit_no_rows_inline"], -) -def test_query_and_wait_arrow_waits_without_fetching_rows(response): +def test_query_and_wait_arrow_waits_without_fetching_rows(): client = mock.MagicMock(spec=Client) - client._call_api.return_value = response + client._call_api.return_value = _arrow_query_response(jobComplete=False) with ( mock.patch.object(_job_helpers, "_to_query_job") as to_query_job, @@ -1111,6 +1100,69 @@ def test_query_and_wait_arrow_waits_without_fetching_rows(response): assert client._call_api.call_args.kwargs["data"]["maxResults"] == 10 +@pytest.mark.parametrize( + "response", + [ + _arrow_query_response(pageRowCount="10", totalRows="25", pageToken="next"), + # A cache hit too large to return inline has no rows, but it still has + # a pageToken. Older BigQuery versions send totalRows "0" and newer + # ones leave it out. + _arrow_query_response(totalRows="0", pageToken="next"), + _arrow_query_response(pageToken="next"), + ], + ids=["some_rows_inline", "cache_hit_total_rows_0", "cache_hit_no_total_rows"], +) +def test_query_and_wait_arrow_complete_with_page_token_doesnt_wait(response): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = response + + 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() + # Not an _EmptyRowIterator: the rows after the first page come from the + # job's stream. + assert type(rows) is RowIterator + assert rows.job_id == "j" + assert rows._first_page_response["pageToken"] == "next" + + +def test_query_and_wait_arrow_cache_hit_without_total_rows_reads_the_stream(): + schema, batches = _arrow_pages(5) + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "cacheHit": True, + "pageToken": "next", + } + reader = mock.MagicMock() + reader.__iter__.return_value = iter( + [ + mock.Mock( + arrow_schema=mock.Mock(serialized_schema=schema), + arrow_record_batch=mock.Mock(serialized_record_batch=batches[0]), + ) + ] + ) + bqstorage_client = mock.MagicMock() + bqstorage_client.read_rows.return_value = reader + + with mock.patch.object( + _job_helpers, + "_wait_or_cancel", + side_effect=AssertionError("a finished query shouldn't wait on its job"), + ): + rows = _query_and_wait(client, query_results_format="ARROW") + received = list(rows.to_arrow_iterable(bqstorage_client=bqstorage_client)) + + assert _values(received) == [1, 2, 3, 4, 5] + assert bqstorage_client.read_rows.call_args.args[0] == ( + "projects/p/locations/us/jobs/j/streams/_default" + ) + assert bqstorage_client.read_rows.call_args.kwargs["offset"] == 0 + + def test_query_and_wait_arrow_all_rows_inline_doesnt_wait(): client = mock.MagicMock(spec=Client) client._call_api.return_value = _arrow_query_response( From d4704fa0190b8247e80bc51cb437dfe8565645f1 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Mon, 2 Nov 2026 18:58:54 +0000 Subject: [PATCH 2/2] test(bigquery): trim ARROW routing tests to the minimal change --- .../unit/test_query_results_format_arrow.py | 53 ++++--------------- 1 file changed, 11 insertions(+), 42 deletions(-) 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 387e458a4d08..e6f12077a0ae 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 @@ -1083,9 +1083,16 @@ def _arrow_query_response(**fields): } -def test_query_and_wait_arrow_waits_without_fetching_rows(): +@pytest.mark.parametrize( + "response", + [ + _arrow_query_response(jobComplete=False), + ], + ids=["no_rows_inline"], +) +def test_query_and_wait_arrow_waits_without_fetching_rows(response): client = mock.MagicMock(spec=Client) - client._call_api.return_value = _arrow_query_response(jobComplete=False) + client._call_api.return_value = response with ( mock.patch.object(_job_helpers, "_to_query_job") as to_query_job, @@ -1120,47 +1127,9 @@ def test_query_and_wait_arrow_complete_with_page_token_doesnt_wait(response): rows = _query_and_wait(client, query_results_format="ARROW") wait_or_cancel.assert_not_called() - # Not an _EmptyRowIterator: the rows after the first page come from the - # job's stream. + # Not an _EmptyRowIterator: the RowIterator reads the rest of the rows + # from the job's ReadRows stream. assert type(rows) is RowIterator - assert rows.job_id == "j" - assert rows._first_page_response["pageToken"] == "next" - - -def test_query_and_wait_arrow_cache_hit_without_total_rows_reads_the_stream(): - schema, batches = _arrow_pages(5) - client = mock.MagicMock(spec=Client) - client._call_api.return_value = { - "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, - "jobComplete": True, - "cacheHit": True, - "pageToken": "next", - } - reader = mock.MagicMock() - reader.__iter__.return_value = iter( - [ - mock.Mock( - arrow_schema=mock.Mock(serialized_schema=schema), - arrow_record_batch=mock.Mock(serialized_record_batch=batches[0]), - ) - ] - ) - bqstorage_client = mock.MagicMock() - bqstorage_client.read_rows.return_value = reader - - with mock.patch.object( - _job_helpers, - "_wait_or_cancel", - side_effect=AssertionError("a finished query shouldn't wait on its job"), - ): - rows = _query_and_wait(client, query_results_format="ARROW") - received = list(rows.to_arrow_iterable(bqstorage_client=bqstorage_client)) - - assert _values(received) == [1, 2, 3, 4, 5] - assert bqstorage_client.read_rows.call_args.args[0] == ( - "projects/p/locations/us/jobs/j/streams/_default" - ) - assert bqstorage_client.read_rows.call_args.kwargs["offset"] == 0 def test_query_and_wait_arrow_all_rows_inline_doesnt_wait():