fix(bigquery): apply max_results to the whole Arrow result

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, Tuple, TYPE_CHECKING
 
 from google.cloud.bigquery import _versions_helpers
 from google.cloud.bigquery import enums
@@ -114,6 +114,26 @@ 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]]:
+    """Cut ``batch`` to at most ``remaining`` rows, and return it with the
+    number of rows still wanted after it (``None`` for no limit)."""
+    if remaining is None:
+        return batch, None
+    if batch.num_rows > remaining:
+        batch = batch.slice(0, remaining)
+    return batch, remaining - batch.num_rows
+
+
 def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]:
     """Convert ``jobs.query`` ``arrowSerializationOptions`` to the
     ``ArrowSerializationOptions`` fields that ``ReadRows`` takes."""
@@ -150,6 +170,7 @@ def download_arrow_from_job(
     pa_schema = None
     page_token = None
     job_complete = False
+    remaining = row_iterator.max_results
 
     # Step 2: Process initial inline Arrow response from jobs.query if available.
     if row_iterator._first_page_response:
@@ -182,12 +203,14 @@ def download_arrow_from_job(
                     pa_schema,
                 )
                 offset += batch.num_rows
+                batch, remaining = _slice_batch(batch, remaining)
                 yield batch
 
     # Step 3: Return early if all results were delivered in the first page
-    # response. Only a pageToken says rows remain: totalRows is 0 when cached
-    # results are too large to return inline.
-    if job_complete and page_token is None:
+    # 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.
+    if (job_complete and page_token is None) or remaining == 0:
         return
 
     # Step 4: Stream remaining Arrow record batches from the job default stream.
@@ -206,23 +229,37 @@ def download_arrow_from_job(
             row_iterator._arrow_serialization_options
         ),
     )
-    for response in reader:
-        if (
-            response.arrow_schema
-            and response.arrow_schema.serialized_schema
-            and pa_schema is None
-        ):
-            pa_schema = pyarrow.ipc.read_schema(
-                pyarrow.py_buffer(response.arrow_schema.serialized_schema)
-            )
-        if (
-            response.arrow_record_batch
-            and response.arrow_record_batch.serialized_record_batch
-        ):
-            if pa_schema is None:
-                raise ValueError(_NO_ARROW_SCHEMA_ERROR)
-            batch = pyarrow.ipc.read_record_batch(
-                pyarrow.py_buffer(response.arrow_record_batch.serialized_record_batch),
-                pa_schema,
-            )
-            yield batch
+    finished = False
+    try:
+        for response in reader:
+            if (
+                response.arrow_schema
+                and response.arrow_schema.serialized_schema
+                and pa_schema is None
+            ):
+                pa_schema = pyarrow.ipc.read_schema(
+                    pyarrow.py_buffer(response.arrow_schema.serialized_schema)
+                )
+            if (
+                response.arrow_record_batch
+                and response.arrow_record_batch.serialized_record_batch
+            ):
+                if pa_schema is None:
+                    raise ValueError(_NO_ARROW_SCHEMA_ERROR)
+                batch = pyarrow.ipc.read_record_batch(
+                    pyarrow.py_buffer(
+                        response.arrow_record_batch.serialized_record_batch
+                    ),
+                    pa_schema,
+                )
+                batch, remaining = _slice_batch(batch, remaining)
+                yield batch
+                if remaining == 0:
+                    break
+        else:
+            finished = True
+    finally:
+        # Stop the stream if reading ends before it does: at max_results, on
+        # an error, or when the caller closes the iterator.
+        if not finished:
+            reader.cancel()

packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py

@@ -620,6 +620,11 @@ 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)
+        ):
+            more_pages = False
 
         if more_pages or not query_results.complete:
             # TODO(swast): Avoid a call to jobs.get in some cases (few

packages/google-cloud-bigquery/google/cloud/bigquery/table.py

@@ -2397,11 +2397,17 @@ def to_arrow(
         if pyarrow is None:
             raise ValueError(_NO_PYARROW_ERROR)
 
-        self._maybe_warn_max_results(bqstorage_client)
-
-        if not self._should_use_bqstorage(bqstorage_client, create_bqstorage_client):
+        if self._query_results_format == QueryResultsFormat.ARROW.value:
+            # ARROW results apply max_results themselves and read with the
+            # bqstorage_client they're given, so keep it.
             create_bqstorage_client = False
-            bqstorage_client = None
+        else:
+            self._maybe_warn_max_results(bqstorage_client)
+            if not self._should_use_bqstorage(
+                bqstorage_client, create_bqstorage_client
+            ):
+                create_bqstorage_client = False
+                bqstorage_client = None
 
         owns_bqstorage_client = False
         if not bqstorage_client and create_bqstorage_client:
@@ -2900,7 +2906,9 @@ def to_dataframe(
         if dtypes is None:
             dtypes = {}
 
-        self._maybe_warn_max_results(bqstorage_client)
+        # ARROW results apply max_results with the BigQuery Storage API too.
+        if self._query_results_format != QueryResultsFormat.ARROW.value:
+            self._maybe_warn_max_results(bqstorage_client)
 
         if (
             self._query_results_format != QueryResultsFormat.ARROW.value

packages/google-cloud-bigquery/tests/system/test_arrow.py

@@ -734,3 +734,32 @@ def test_query_and_wait_arrow_failed_query(bigquery_client, priority):
 
     assert type(arrow_error.value) is type(struct_encoding_error.value)
     assert "failed at row 1000" in arrow_error.value.message
+
+
+@pytest.mark.parametrize(
+    ("page_size", "max_results", "expected_rows", "streamed"),
+    [
+        (None, 10, 10, False),
+        (100, 150, 150, True),
+        (100, 5000, 1000, True),
+    ],
+    ids=["within_first_page", "across_stream", "past_last_row"],
+)
+def test_query_and_wait_arrow_max_results(
+    bigquery_client, bqstorage_client, page_size, max_results, expected_rows, streamed
+):
+    read_rows = mock.Mock(wraps=bqstorage_client.read_rows)
+
+    rows = bigquery_client.query_and_wait(
+        "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x",
+        job_config=bigquery.QueryJobConfig(use_query_cache=False),
+        page_size=page_size,
+        max_results=max_results,
+        query_results_format=enums.QueryResultsFormat.ARROW,
+    )
+    with mock.patch.object(bqstorage_client, "read_rows", read_rows):
+        batches = list(rows.to_arrow_iterable(bqstorage_client=bqstorage_client))
+
+    table = pyarrow.Table.from_batches(batches)
+    assert table.column("x").to_pylist() == list(range(1, expected_rows + 1))
+    assert read_rows.called == streamed