[ohmayr]

Context & Problem

query_and_wait(..., query_results_format="ARROW") returns 0 rows, with no error, when a query hits the cache and its results are too large to return inline. This is the cause of the intermittent test_query_and_wait_arrow_cached_results[too_large_for_inline] system test failure in #18935.

For these cache hits, the jobs.query response has jobComplete: true and a pageToken, but no rows. BigQuery used to send totalRows: "0" for these responses. It is rolling out a change that leaves totalRows out instead, because it doesn't know the count for cache hits. The test passes against backends that still send "0" and fails against the ones that don't, so it looks flaky during the rollout.

Because of the pageToken, query_and_wait passed the response to QueryJob.result(). That method reads a missing totalRows as a statement without results, such as DDL, and returns an _EmptyRowIterator. The Arrow reader, which would have followed the pageToken to the job's ReadRows stream, never ran.

Live probe of the cached ARROW jobs.query response before this change (18 of 18 runs):

Field	Value
jobComplete	true
pageToken	present
totalRows	missing
Rows read	0 (expected 20000)

Summary of Changes

• _job_helpers.query_and_wait: for ARROW, wait on the job only while jobComplete is false. ARROW rows never come from jobs.getQueryResults. The RowIterator reads the rows after the first page from the job's default ReadRows stream, and it only needs jobComplete and pageToken to decide whether to read it. The default (JSON) format behaves as before, because its RowIterator still pages through jobs.getQueryResults via the job.
• _arrow_query: remove reached_max_results, which nothing calls anymore, because complete ARROW results no longer wait whether or not max_results was reached. Fix the comment that said totalRows is 0 for these cache hits.
• Tests:
  • test_query_and_wait_arrow_waits_without_fetching_rows keeps only the jobComplete: false case. Its other two cases asserted that complete responses with a pageToken wait on the job, which is the routing this change removes.
  • Those two cases move to a new test_query_and_wait_arrow_complete_with_page_token_doesnt_wait, which asserts that they don't wait. It also adds the response that caused the bug: complete, with a pageToken and no totalRows.
  • Fix the stale comment in the system test.

rows.total_rows is None for these cache hits, because BigQuery doesn't send the count. Before this change, the same queries returned no rows at all.

Verification

• Red/green: the 3 new test cases fail without the change and pass with it.
• tests/unit: 2824 passed, 43 skipped.
• flake8 and black==23.7.0 are clean on the changed files.
• tests/system/test_arrow.py against live BigQuery: 48 passed, 6 skipped, 6 xfailed. The xfails are expected failures already in the file: TIMESTAMP(12) isn't enabled in the test project, and jobs.query can't send unnamed STRUCT fields. test_query_and_wait_arrow_cached_results[too_large_for_inline] passes against a backend that leaves out totalRows, the same one that failed it every time before this change.

Fixes #18935

[gemini-code-assist]

Code Review

This pull request removes the reached_max_results helper function and updates the query handling logic for ARROW serialization. Specifically, it ensures that when ARROW serialization options are present, the client does not wait on jobs.getQueryResults since ARROW rows are fetched via the ReadRows stream. It also updates comments and tests to handle cases where BigQuery leaves out totalRows for cached results that are too large to return inline, and adds unit tests to verify these behaviors. There are no review comments to address, so I have no feedback to provide.