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()