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 Any, 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
@@ -113,14 +113,29 @@ def serialization_options(
     }
 
 
-def has_more_rows(response: Dict[str, Any]) -> bool:
+def has_more_rows(response: Dict[str, Any], max_results: Optional[int]) -> 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``.
+    ``pageRowCount``, so compare that with ``totalRows``, or ``max_results`` if
+    it's smaller, instead of checking for a ``pageToken``.
     """
-    return int(response.get("pageRowCount", 0)) < int(response.get("totalRows", 0))
+    rows_wanted = int(response.get("totalRows", 0))
+    if max_results is not None:
+        rows_wanted = min(rows_wanted, max_results)
+    return int(response.get("pageRowCount", 0)) < rows_wanted
+
+
+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]:
@@ -159,6 +174,7 @@ def download_arrow_from_job(
     pa_schema = None
     total_rows = row_iterator.total_rows
     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:
@@ -187,10 +203,12 @@ 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.
-    if job_complete and offset >= total_rows:
+    # Step 3: Return early if all results were delivered in the first page
+    # response, or if the first page already reached max_results.
+    if (job_complete and offset >= total_rows) or remaining == 0:
         return
 
     # Step 4: Stream remaining Arrow record batches from the job default stream.
@@ -209,22 +227,36 @@ 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
-            and pa_schema is not None
-        ):
-            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
+                and pa_schema is not None
+            ):
+                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

@@ -621,7 +621,7 @@ 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)
+            more_pages = _arrow_query.has_more_rows(response, max_results)
 
         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

@@ -2408,11 +2408,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:
@@ -2911,7 +2917,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

@@ -726,3 +726,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

packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py

@@ -15,7 +15,9 @@
 import base64
 import unittest
 from unittest import mock
+import warnings
 
+import google.api_core.exceptions as core_exceptions
 import packaging.version
 import pytest
 
@@ -1026,5 +1028,175 @@ def test_query_and_wait_without_arrow_waits_with_page_size():
     assert wait_or_cancel.call_args.kwargs["page_size"] == 10
 
 
+def _arrow_pages(*row_counts):
+    """Serialize an Arrow schema and record batches of consecutive x values."""
+    pyarrow = pytest.importorskip("pyarrow")
+    schema = pyarrow.schema([("x", pyarrow.int64())])
+    batches, start = [], 1
+    for count in row_counts:
+        values = pyarrow.array(range(start, start + count), type=pyarrow.int64())
+        batch = pyarrow.record_batch([values], schema=schema)
+        batches.append(batch.serialize().to_pybytes())
+        start += count
+    return schema.serialize().to_pybytes(), batches
+
+
+def _stream_response(batch):
+    return mock.Mock(
+        arrow_schema=None,
+        arrow_record_batch=mock.Mock(serialized_record_batch=batch),
+    )
+
+
+def _limited_iterator(row_counts, total_rows, max_results):
+    """A RowIterator whose first page has row_counts[0] rows and whose stream
+    returns the rest, plus its mock BigQuery Storage client and stream.
+    """
+    schema, batches = _arrow_pages(*row_counts)
+    reader = mock.MagicMock()
+    reader.__iter__.return_value = iter(
+        [_stream_response(batch) for batch in batches[1:]]
+    )
+    bqstorage_client = mock.MagicMock()
+    bqstorage_client.read_rows.return_value = reader
+    iterator = RowIterator(
+        client=mock.MagicMock(),
+        api_request=mock.MagicMock(),
+        path=None,
+        schema=(),
+        max_results=max_results,
+        project="p",
+        location="US",
+        job_id="j",
+        query_results_format="ARROW",
+        first_page_response={
+            "jobComplete": True,
+            "totalRows": str(total_rows),
+            "arrowSchema": {"serializedSchema": base64.b64encode(schema).decode()},
+            "arrowRecordBatch": {
+                "serializedRecordBatch": base64.b64encode(batches[0]).decode()
+            },
+        },
+    )
+    return iterator, bqstorage_client, reader
+
+
+def _values(batches):
+    return [value for batch in batches for value in batch.column(0).to_pylist()]
+
+
+def test_download_arrow_max_results_within_first_page_skips_stream():
+    iterator, bqstorage_client, _ = _limited_iterator([10, 15], 25, max_results=4)
+
+    batches = list(iterator._download_arrow_from_job_id(bqstorage_client))
+
+    assert _values(batches) == [1, 2, 3, 4]
+    bqstorage_client.read_rows.assert_not_called()
+
+
+def test_download_arrow_max_results_across_stream_cancels_it():
+    iterator, bqstorage_client, reader = _limited_iterator(
+        [10, 10, 5], 25, max_results=15
+    )
+
+    batches = list(iterator._download_arrow_from_job_id(bqstorage_client))
+
+    assert _values(batches) == list(range(1, 16))
+    assert bqstorage_client.read_rows.call_args.kwargs["offset"] == 10
+    reader.cancel.assert_called_once_with()
+
+
+def test_download_arrow_max_results_past_last_row_reads_everything():
+    iterator, bqstorage_client, reader = _limited_iterator(
+        [10, 10, 5], 25, max_results=100
+    )
+
+    batches = list(iterator._download_arrow_from_job_id(bqstorage_client))
+
+    assert _values(batches) == list(range(1, 26))
+    reader.cancel.assert_not_called()
+
+
+def test_download_arrow_cancels_stream_when_closed_early():
+    iterator, bqstorage_client, reader = _limited_iterator(
+        [10, 10, 5], 25, max_results=None
+    )
+    batches = iterator._download_arrow_from_job_id(bqstorage_client)
+
+    next(batches)  # first page
+    next(batches)  # first stream batch
+    batches.close()
+
+    reader.cancel.assert_called_once_with()
+
+
+def test_download_arrow_raises_stream_errors_after_earlier_batches():
+    iterator, bqstorage_client, reader = _limited_iterator(
+        [10, 10, 5], 25, max_results=None
+    )
+    first_stream_response = next(reader.__iter__.return_value)
+
+    def fail_after_one_response():
+        yield first_stream_response
+        raise core_exceptions.PermissionDenied("stream failed")
+
+    reader.__iter__.return_value = fail_after_one_response()
+    received = []
+
+    with pytest.raises(core_exceptions.PermissionDenied):
+        for batch in iterator._download_arrow_from_job_id(bqstorage_client):
+            received.append(batch)
+
+    # The caller gets the error after the batches that arrived before it, so
+    # a partial result never looks complete.
+    assert _values(received) == list(range(1, 21))
+    reader.cancel.assert_called_once_with()
+
+
+def test_query_and_wait_arrow_first_page_reaching_max_results_doesnt_wait():
+    client = mock.MagicMock(spec=Client)
+    client._call_api.return_value = _arrow_query_response(
+        pageRowCount="10", totalRows="25"
+    )
+
+    with mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel:
+        rows = _query_and_wait(client, query_results_format="ARROW", max_results=10)
+
+    wait_or_cancel.assert_not_called()
+    assert rows.max_results == 10
+    assert client._call_api.call_args.kwargs["data"]["maxResults"] == 10
+
+
+def test_to_arrow_with_arrow_format_and_max_results_keeps_bqstorage_client():
+    pyarrow = pytest.importorskip("pyarrow")
+    bqstorage_client = mock.MagicMock()
+    iterator = RowIterator(
+        client=mock.MagicMock(),
+        api_request=mock.MagicMock(),
+        path=None,
+        schema=(),
+        max_results=5,
+        query_results_format="ARROW",
+    )
+    batch = pyarrow.record_batch([pyarrow.array([1, 2])], names=["x"])
+
+    with (
+        mock.patch.object(
+            iterator, "to_arrow_iterable", return_value=iter([batch])
+        ) as to_arrow_iterable,
+        warnings.catch_warnings(record=True) as caught,
+    ):
+        warnings.simplefilter("always")
+        table = iterator.to_arrow(bqstorage_client=bqstorage_client)
+
+    to_arrow_iterable.assert_called_once_with(
+        bqstorage_client=bqstorage_client, timeout=None
+    )
+    assert table.num_rows == 2
+    # ARROW results apply max_results while streaming, so there's no warning
+    # about falling back to the REST API.
+    assert not [w for w in caught if "max_results" in str(w.message)]
+
+
 if __name__ == "__main__":
     unittest.main()