From f43d1957cc483e39e82b3b1838019d7b7260207b Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 1/2] 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 | 83 ++++++--- .../google/cloud/bigquery/_job_helpers.py | 5 + .../google/cloud/bigquery/table.py | 18 +- .../tests/system/test_arrow.py | 29 +++ .../unit/test_query_results_format_arrow.py | 176 ++++++++++++++++++ 5 files changed, 283 insertions(+), 28 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 ca98bc18f39d..32f95a76db45 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 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,6 +113,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.""" @@ -149,6 +169,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: @@ -179,12 +200,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. @@ -203,22 +226,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 a601526ecaa4..c06579797a48 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/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 diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index dcaba702dadf..5de18251d678 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/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 diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index f8ece332c6e8..1c9bf6bff0a2 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/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 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 1855d854c470..2ce08893a98a 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 @@ -1078,6 +1080,180 @@ 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 + first_page_response = { + "jobComplete": True, + "totalRows": str(total_rows), + "arrowSchema": {"serializedSchema": base64.b64encode(schema).decode()}, + "arrowRecordBatch": { + "serializedRecordBatch": base64.b64encode(batches[0]).decode() + }, + } + # BigQuery sends a pageToken whenever rows remain after the first page. + if row_counts[0] < total_rows: + first_page_response["pageToken"] = "next" + 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=first_page_response, + ) + 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", pageToken="next" + ) + + 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)] + + def test_to_arrow_iterable_reads_inline_result_twice(): # With JOB_CREATION_OPTIONAL and every row inline there's no job to # stream from, so the first page has to stay on the iterator for the From cc938a0cbb72b6c8e512cc358f6d9bad1af71f1e Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 20 Oct 2026 22:24:01 +0000 Subject: [PATCH 2/2] test(bigquery): mock a cancellable reader in no-schema test The ReadRows loop now cancels the stream in a finally block when reading stops early. The no-schema test returned a plain list from read_rows(), so cancel() raised AttributeError and hid the ValueError. Return a mock reader instead and check that the stream is cancelled. --- .../tests/unit/test_query_results_format_arrow.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 753d70ea6e12..c6be6a862e1b 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 @@ -480,13 +480,16 @@ def test_download_arrow_from_job_id_read_rows_batch_without_schema_raises(self): mock_response = mock.MagicMock() mock_response.arrow_schema = None mock_response.arrow_record_batch.serialized_record_batch = b"batch_bytes" - mock_bqstorage.read_rows.return_value = [mock_response] + reader = mock.MagicMock() + reader.__iter__.return_value = iter([mock_response]) + mock_bqstorage.read_rows.return_value = reader with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: with self.assertRaisesRegex(ValueError, "before its schema"): list(iterator._download_arrow_from_job_id(timeout=5.0)) mock_pyarrow.ipc.read_record_batch.assert_not_called() + reader.cancel.assert_called_once_with() def test_empty_row_iterator_to_arrow_iterable_checks_pyarrow(self): iterator = _EmptyRowIterator(query_results_format="ARROW")