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..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 @@ -1086,13 +1086,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"], + ids=["no_rows_inline"], ) def test_query_and_wait_arrow_waits_without_fetching_rows(response): client = mock.MagicMock(spec=Client) @@ -1111,6 +1107,31 @@ 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 RowIterator reads the rest of the rows + # from the job's ReadRows stream. + assert type(rows) is RowIterator + + def test_query_and_wait_arrow_all_rows_inline_doesnt_wait(): client = mock.MagicMock(spec=Client) client._call_api.return_value = _arrow_query_response(