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.

Author: ohmayr

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:

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,

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(