From 1b5717eaf5a1e89a6f61286bbafc10392dad97c4 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH] fix(bigquery): apply max_results to the whole Arrow result With ARROW, max_results only limited the first page from jobs.query. Every remaining row was then streamed, so a query with max_results=150 and a 100-row first page returned the whole result. to_arrow() and to_dataframe() also warned that max_results made them fall back to REST, which isn't true for ARROW, and to_arrow() then dropped the bqstorage_client the caller passed. RowIterator now counts rows across the first page and the stream, slices the batch that reaches max_results, and stops. If the first page already reaches the limit, it doesn't open a stream, and query_and_wait() doesn't wait for rows it won't read. A stream that ends early, at the limit, on an error, or because the caller closed the iterator, is cancelled with ReadRowsStream.cancel(), and errors are raised after the batches that arrived before them. For ARROW results, to_arrow() and to_dataframe() skip the warning and keep the caller's bqstorage_client. --- .../google/cloud/bigquery/_arrow_query.py | 84 ++++++--- .../google/cloud/bigquery/_job_helpers.py | 2 +- .../google/cloud/bigquery/table.py | 18 +- .../tests/system/test_arrow.py | 29 +++ .../unit/test_query_results_format_arrow.py | 172 ++++++++++++++++++ 5 files changed, 273 insertions(+), 32 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index c15891fc064b..f627e4eb619b 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/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() diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py index f5080611a081..a916245891f4 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/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 diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index 369140f47573..f730914b16b9 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -2396,11 +2396,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: @@ -2899,7 +2905,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 diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index 78fb935124f7..cc3ed5710b9a 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/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 diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index d8446a3eefc8..81d3e6abbf30 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/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()