fix(bigquery): read cached ARROW results that are too large to return inline

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)

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:

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),
     ],

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(