diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py index 0decb4f413ba..556318ce440e 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py @@ -57,6 +57,7 @@ def read_rows( timeout=google.api_core.gapic_v1.method.DEFAULT, metadata=(), retry_delay_callback=None, + arrow_serialization_options=None, ): """ Reads rows from the table in the format prescribed by the read @@ -124,6 +125,14 @@ def read_rows( BigQueryReadClient will call retry_delay_callback with the delay duration (in seconds) before it starts sleeping until the next attempt. + arrow_serialization_options (Optional[Union[dict, \ + ~google.cloud.bigquery_storage_v1.types.ArrowSerializationOptions \ + ]]): + Arrow serialization options, such as buffer compression, to + send on every ``ReadRows`` request, including after a + reconnect. Only job default streams + (``projects/{project}/locations/{location}/jobs/{job}/streams/_default``) + accept them. Returns: ~google.cloud.bigquery_storage_v1.reader.ReadRowsStream: @@ -144,6 +153,7 @@ def read_rows( offset, {"retry": retry, "timeout": timeout, "metadata": metadata}, retry_delay_callback=retry_delay_callback, + arrow_serialization_options=arrow_serialization_options, ) stream._reconnect() return stream diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py index 27ab7f55de48..5f49187211e9 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py @@ -86,7 +86,13 @@ class ReadRowsStream(object): """ def __init__( - self, client, name, offset, read_rows_kwargs, retry_delay_callback=None + self, + client, + name, + offset, + read_rows_kwargs, + retry_delay_callback=None, + arrow_serialization_options=None, ): """Construct a ReadRowsStream. @@ -114,6 +120,11 @@ def __init__( ReadRowsStream will call retry_delay_callback with the delay duration (in seconds) before it starts sleeping until the next attempt. + arrow_serialization_options (Optional[Union[dict, \ + ~google.cloud.bigquery_storage_v1.types.ArrowSerializationOptions \ + ]]): + Arrow serialization options to send on every ``ReadRows`` + request, including after a reconnect. Returns: Iterable[ \ @@ -129,6 +140,8 @@ def __init__( self._offset = offset self._read_rows_kwargs = read_rows_kwargs self._retry_delay_callback = retry_delay_callback + self._arrow_serialization_options = arrow_serialization_options + self._cancelled = False self._wrapped = None def __iter__(self): @@ -140,13 +153,13 @@ def __iter__(self): ]: A sequence of row messages. """ - # Infinite loop to reconnect on reconnectable errors while processing - # the row stream. + # Loop to reconnect on reconnectable errors while processing the row + # stream, until the stream ends or is cancelled. if self._wrapped is None: self._reconnect() - while True: + while not self._cancelled: try: for message in self._wrapped: rowcount = message.row_count @@ -165,22 +178,57 @@ def __iter__(self): # Transient error, so reconnect to the stream. pass except Exception as exc: + # cancel() makes the call fail, which isn't an error here. + if self._cancelled: + return if not self._resource_exhausted_exception_is_retryable(exc): raise self._reconnect() + def cancel(self): + """Cancel the ``ReadRows`` call. + + Use this to stop reading before the end of the stream. After + ``cancel()``, iterating the stream yields no more messages, and the + stream doesn't reconnect. The client stays open for other calls. + """ + self._cancelled = True + cancel = getattr(self._wrapped, "cancel", None) + if cancel is not None: + cancel() + def _reconnect(self): """Reconnect to the ReadRows stream using the most recent offset.""" - while True: + while not self._cancelled: try: - self._wrapped = self._client.read_rows( - read_stream=self._name, - offset=self._offset, - **self._read_rows_kwargs, - ) + if self._arrow_serialization_options is None: + wrapped = self._client.read_rows( + read_stream=self._name, + offset=self._offset, + **self._read_rows_kwargs, + ) + else: + # The flattened read_rows() arguments don't include + # arrow_serialization_options, so send a full request. + wrapped = self._client.read_rows( + request={ + "read_stream": self._name, + "offset": self._offset, + "arrow_serialization_options": self._arrow_serialization_options, + }, + **self._read_rows_kwargs, + ) + # Store the new call before checking for cancel(), so a + # concurrent cancel() either cancels it or is seen here. + self._wrapped = wrapped + if self._cancelled: + self.cancel() + return break except Exception as exc: + if self._cancelled: + return if not self._resource_exhausted_exception_is_retryable(exc): raise diff --git a/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py b/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py index e9269bda9d96..a9279c89f05c 100644 --- a/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py +++ b/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py @@ -536,3 +536,67 @@ def test_read_rows_to_dataframe_with_wide_table(client_and_types, project_id): some_rows = next(pages_iter) assert all(len(row["tract_geom"].as_py()) > 0 for row in some_rows) + + +def _arrow_table(responses): + import pyarrow + + schema, batches = None, [] + for response in responses: + if schema is None and response.arrow_schema.serialized_schema: + schema = pyarrow.ipc.read_schema( + pyarrow.py_buffer(response.arrow_schema.serialized_schema) + ) + if response.arrow_record_batch.serialized_record_batch: + batches.append( + pyarrow.ipc.read_record_batch( + pyarrow.py_buffer( + response.arrow_record_batch.serialized_record_batch + ), + schema, + ) + ) + return pyarrow.Table.from_batches(batches, schema=schema) + + +@pytest.fixture(scope="module") +def job_default_stream(bq_client): + job = bq_client.query( + "SELECT x, REPEAT('A', 1000) AS repeated_str " + "FROM UNNEST(GENERATE_ARRAY(1, 20000)) AS x ORDER BY x" + ) + job.result() + return ( + f"projects/{job.project}/locations/{job.location}" + f"/jobs/{job.job_id}/streams/_default" + ) + + +@pytest.mark.parametrize("codec", ["LZ4_FRAME", "ZSTD"]) +def test_read_rows_job_stream_with_arrow_serialization_options( + client_v1, job_default_stream, codec +): + plain = list(client_v1.read_rows(job_default_stream)) + compressed = list( + client_v1.read_rows( + job_default_stream, + arrow_serialization_options={"buffer_compression": codec}, + ) + ) + + def size(responses): + return sum(len(r.arrow_record_batch.serialized_record_batch) for r in responses) + + assert size(compressed) < size(plain) + assert _arrow_table(compressed).equals(_arrow_table(plain)) + + +def test_read_rows_job_stream_cancel(client_v1, job_default_stream): + stream = client_v1.read_rows(job_default_stream) + messages = iter(stream) + first = next(messages) + + stream.cancel() + + assert first.row_count > 0 + assert list(messages) == [] diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py b/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py index abb88ba36edc..bca9647bdffd 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py @@ -153,6 +153,58 @@ def test_read_rows(mock_transport, client_under_test): ) +def test_read_rows_with_arrow_serialization_options(mock_transport, client_under_test): + rpc_callable = mock.Mock() + mock_transport._wrapped_methods[mock_transport.read_rows] = rpc_callable + options = {"buffer_compression": "LZ4_FRAME"} + + client_under_test.read_rows( + "teststream", offset=10, arrow_serialization_options=options + ) + + rpc_callable.assert_called_once_with( + types.ReadRowsRequest( + read_stream="teststream", offset=10, arrow_serialization_options=options + ), + metadata=mock.ANY, + retry=mock.ANY, + timeout=mock.ANY, + ) + + +def test_read_rows_reconnect_resends_arrow_serialization_options( + mock_transport, client_under_test +): + def first_stream(): + yield types.ReadRowsResponse(row_count=5) + raise google.api_core.exceptions.ServiceUnavailable("transient disconnect") + + rpc_callable = mock.Mock( + side_effect=[first_stream(), iter([types.ReadRowsResponse(row_count=3)])] + ) + mock_transport._wrapped_methods[mock_transport.read_rows] = rpc_callable + options = {"buffer_compression": "ZSTD"} + + stream = client_under_test.read_rows( + "teststream", offset=10, arrow_serialization_options=options + ) + + assert [m.row_count for m in stream] == [5, 3] + assert rpc_callable.call_args_list == [ + mock.call( + types.ReadRowsRequest( + read_stream="teststream", + offset=offset, + arrow_serialization_options=options, + ), + metadata=mock.ANY, + retry=mock.ANY, + timeout=mock.ANY, + ) + for offset in (10, 15) + ] + + @pytest.mark.parametrize( "module_under_test", ["google.cloud.bigquery_storage_v1", "google.cloud.bigquery_storage_v1beta2"], diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py index d5a83b2a7cf7..6bc4115f9ce4 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py @@ -717,3 +717,86 @@ def test_to_dataframe_emits_pending_deprecation_warning(mut): assert actual_df == mock_df mock_parser.to_dataframe.assert_called_once_with(mock_message, dtypes=None) + + +def test_cancel_stops_iteration_without_reconnecting( + class_under_test, mock_gapic_client +): + first = types.ReadRowsResponse(row_count=1) + + def pages(): + yield first + raise google.api_core.exceptions.ServiceUnavailable("connection closed") + + wrapped = mock.MagicMock() + wrapped.__iter__.return_value = pages() + mock_gapic_client.read_rows.return_value = wrapped + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + messages = iter(reader) + assert next(messages) is first + reader.cancel() + + assert list(messages) == [] + wrapped.cancel.assert_called_once_with() + mock_gapic_client.read_rows.assert_called_once_with( + read_stream="teststream", offset=0 + ) + + +def test_cancel_ignores_the_cancelled_call_error(class_under_test, mock_gapic_client): + def pages(): + yield types.ReadRowsResponse(row_count=1) + raise google.api_core.exceptions.Cancelled("call cancelled") + + mock_gapic_client.read_rows.return_value = pages() + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + messages = iter(reader) + next(messages) + reader.cancel() + + assert list(messages) == [] + + +def test_cancel_before_iterating_yields_nothing(class_under_test, mock_gapic_client): + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + reader.cancel() + + assert list(reader) == [] + mock_gapic_client.read_rows.assert_not_called() + + +def test_cancel_during_reconnect_cancels_new_stream( + class_under_test, mock_gapic_client +): + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + wrapped = mock.MagicMock() + + def cancel_while_opening(**kwargs): + reader.cancel() + return wrapped + + mock_gapic_client.read_rows.side_effect = cancel_while_opening + + assert list(reader) == [] + wrapped.cancel.assert_called_once_with() + assert reader._wrapped is wrapped + + +def test_cancel_during_reconnect_retry_stops_without_retrying( + class_under_test, mock_gapic_client +): + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + def cancel_and_fail(**kwargs): + reader.cancel() + raise google.api_core.exceptions.Cancelled("call cancelled") + + mock_gapic_client.read_rows.side_effect = cancel_and_fail + + assert list(reader) == [] + mock_gapic_client.read_rows.assert_called_once_with( + read_stream="teststream", offset=0 + ) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py new file mode 100644 index 000000000000..881fdc0661a6 --- /dev/null +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -0,0 +1,376 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Read query results returned in Apache Arrow format.""" + +from __future__ import annotations + +import base64 +from typing import Any, Dict, Iterator, Optional, Sequence, Tuple, TYPE_CHECKING + +from google.cloud.bigquery import _versions_helpers +from google.cloud.bigquery import enums +from google.cloud.bigquery import exceptions + +try: + import pyarrow # type: ignore +except ImportError: + pyarrow = None # type: ignore[assignment] + +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud import bigquery_storage + from google.cloud.bigquery import job + from google.cloud.bigquery.table import RowIterator + +_NO_PYARROW_ERROR = ( + "The pyarrow library is not installed, please install " + "pyarrow to use the to_arrow() function." +) + + +def validate_query( + query_results_format: Optional[str], + job_config: Optional["job.QueryJobConfig"], + *, + compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, +) -> None: + """Raise before sending a query whose ``ARROW`` results can't be read. + + Args: + query_results_format: The ``query_results_format`` passed to + ``query_and_wait()``. + job_config: The job config, already merged with + ``Client.default_query_job_config``. + compression_codec: The ``compression_codec`` passed to + ``query_and_wait()``. + timestamp_precision: The ``timestamp_precision`` passed to + ``query_and_wait()``. + + Raises: + ValueError: + If ``compression_codec`` or ``timestamp_precision`` is set + without ``ARROW``, or if an ``ARROW`` query uses ``WRITE_APPEND``. + google.cloud.bigquery.exceptions.BigQueryError: + If pyarrow or google-cloud-bigquery-storage is missing or too old + to read ``ARROW`` results. + """ + if query_results_format != enums.QueryResultsFormat.ARROW: + if compression_codec is not None: + raise ValueError("compression_codec requires query_results_format=ARROW.") + if timestamp_precision is not None: + raise ValueError("timestamp_precision requires query_results_format=ARROW.") + return + + # Results that don't fit in the jobs.query response are read from the + # destination table, which with WRITE_APPEND also holds older rows. + if ( + job_config is not None + and job_config.write_disposition == enums.WriteDisposition.WRITE_APPEND + ): + raise ValueError( + "query_results_format=ARROW doesn't support WRITE_APPEND, because " + "the results would also include the rows already in the " + "destination table." + ) + + # Check before the query runs, so a query doesn't succeed while its + # results fit in the first page and then fail once they need a stream. + _check_dependencies() + + +def _check_dependencies() -> None: + """Raise if pyarrow or google-cloud-bigquery-storage can't read ``ARROW`` + results.""" + _versions_helpers.PYARROW_VERSIONS.try_import(raise_if_error=True) + _versions_helpers.BQ_STORAGE_VERSIONS.try_import(raise_if_error=True) + if not _versions_helpers.BQ_STORAGE_VERSIONS.supports_arrow_serialization_options: + raise exceptions.LegacyBigQueryStorageError( + "query_results_format=ARROW needs a google-cloud-bigquery-storage " + "release whose BigQueryReadClient.read_rows() accepts " + "arrow_serialization_options, please upgrade it " + "(`pip install --upgrade google-cloud-bigquery-storage`) " + "(version found: " + f"{_versions_helpers.BQ_STORAGE_VERSIONS.installed_version})." + ) + + +def serialization_options( + compression_codec: Optional[str], + timestamp_precision: Optional[enums.TimestampPrecision] = None, +) -> Dict[str, str]: + """Build the ``arrowSerializationOptions`` sent on ``jobs.query``. + + ``ReadRows`` gets the same options. Options the caller didn't set are sent + with their defaults, so every batch uses one encoding, whatever each + endpoint would default to. + """ + return { + "bufferCompression": getattr(compression_codec, "value", compression_codec) + or "COMPRESSION_UNSPECIFIED", + # Picosecond values come back as ISO 8601 strings at full precision, + # so there's no need to offer nanoseconds. + "picosTimestampPrecision": ( + "TIMESTAMP_PRECISION_PICOS" + if timestamp_precision == enums.TimestampPrecision.PICOSECOND + else "TIMESTAMP_PRECISION_MICROS" + ), + } + + +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``, or ``max_results`` if + it's smaller, instead of checking for a ``pageToken``. + """ + 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 _empty_batch( + pa_schema: Optional["pyarrow.Schema"], bq_schema: Sequence[Any] +) -> "pyarrow.RecordBatch": + """Build a batch without rows that keeps the result's schema. + + Use the Arrow schema BigQuery sent if there is one. The default job + stream sends nothing for zero rows, so fall back to the schema from + ``jobs.getQueryResults``, and else to no columns. + """ + if pa_schema is None and bq_schema: + # Imported here: _pandas_helpers imports pandas_gbq, which imports + # this package's table module. + from google.cloud.bigquery import _pandas_helpers + + pa_schema = _pandas_helpers.bq_to_arrow_schema(bq_schema) + if pa_schema is None: + pa_schema = pyarrow.schema([]) + return pyarrow.RecordBatch.from_arrays( + [pyarrow.array([], type=field.type) for field in pa_schema], + schema=pa_schema, + ) + + +def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]: + """Convert ``jobs.query`` ``arrowSerializationOptions`` to the + ``ArrowSerializationOptions`` fields that ``ReadRows`` takes.""" + options = options or serialization_options(None) + return { + "buffer_compression": options["bufferCompression"], + "picos_timestamp_precision": options["picosTimestampPrecision"], + } + + +def _open_job_stream( + row_iterator: "RowIterator", + bqstorage_client: "bigquery_storage.BigQueryReadClient", + offset: int, + timeout: Optional[float], +): + """Open the job's default stream, starting at row ``offset``.""" + project = row_iterator._project or ( + row_iterator.client.project if row_iterator.client else None + ) + location = row_iterator._location or ( + row_iterator.client.location if row_iterator.client else None + ) + stream_name = f"projects/{project}/locations/{location}/jobs/{row_iterator._job_id}/streams/_default" + return bqstorage_client.read_rows( + stream_name, + offset=offset, + timeout=timeout, + arrow_serialization_options=_read_rows_options( + row_iterator._arrow_serialization_options + ), + ) + + +def _open_destination_stream( + row_iterator: "RowIterator", + bqstorage_client: "bigquery_storage.BigQueryReadClient", + timeout: Optional[float], + pa_schema: Optional["pyarrow.Schema"], +): + """Open a read session on the query's explicit destination table. + + The default job stream rejects jobs with an explicit destination table + ("This API does not support explicitly-specified destination tables"), so + read the table instead. One stream keeps the rows in order, and the + session's read options carry the same serialization options as + ``jobs.query``. Explicit destinations always go through ``jobs.insert``, + so there's no first page and the stream starts at row 0. + + Returns ``(reader, pa_schema)``, with ``reader`` set to ``None`` if the + table is empty and the session has no streams. + """ + from google.cloud import bigquery_storage + + session = bqstorage_client.create_read_session( + parent=f"projects/{row_iterator._billing_project}", + read_session=bigquery_storage.types.stream.ReadSession( + table=row_iterator._table.to_bqstorage(), + data_format=bigquery_storage.types.stream.DataFormat.ARROW, + read_options=bigquery_storage.types.stream.ReadSession.TableReadOptions( + arrow_serialization_options=_read_rows_options( + row_iterator._arrow_serialization_options + ) + ), + ), + max_stream_count=1, + timeout=timeout, + ) + if pa_schema is None and session.arrow_schema.serialized_schema: + pa_schema = pyarrow.ipc.read_schema( + pyarrow.py_buffer(session.arrow_schema.serialized_schema) + ) + if not session.streams: + return None, pa_schema + return ( + bqstorage_client.read_rows(session.streams[0].name, offset=0, timeout=timeout), + pa_schema, + ) + + +def download_arrow_from_job( + row_iterator: "RowIterator", + bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None, + timeout: Optional[float] = None, +) -> Iterator["pyarrow.RecordBatch"]: + """Yield Arrow record batches for ``row_iterator``'s ARROW query results. + + See :meth:`google.cloud.bigquery.table.RowIterator._download_arrow_from_job_id`. + """ + if pyarrow is None: + raise ValueError(_NO_PYARROW_ERROR) + + offset = 0 + pa_schema = None + total_rows = row_iterator.total_rows + job_complete = False + remaining = row_iterator.max_results + # Results without rows still yield one empty batch, so callers get the + # schema. + yielded = False + + # Step 1: Process initial inline Arrow response from jobs.query if available. + if row_iterator._first_page_response: + first_page = row_iterator._first_page_response + + job_complete = bool(first_page.get("jobComplete", False)) + if job_complete: + total_rows = int(first_page.get("totalRows", 0)) + + arrow_schema_json = first_page.get("arrowSchema") + if isinstance(arrow_schema_json, dict): + schema_bytes = arrow_schema_json.get("serializedSchema") + if schema_bytes: + if isinstance(schema_bytes, str): + schema_bytes = base64.b64decode(schema_bytes) + pa_schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes)) + + arrow_batch_json = first_page.get("arrowRecordBatch") + if isinstance(arrow_batch_json, dict) and pa_schema is not None: + batch_bytes = arrow_batch_json.get("serializedRecordBatch") + if batch_bytes: + if isinstance(batch_bytes, str): + batch_bytes = base64.b64decode(batch_bytes) + batch = pyarrow.ipc.read_record_batch( + pyarrow.py_buffer(batch_bytes), + pa_schema, + ) + offset += batch.num_rows + batch, remaining = _slice_batch(batch, remaining) + yielded = True + yield batch + + # Step 2: 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: + if not yielded: + yield _empty_batch(pa_schema, row_iterator._schema) + return + + # Step 3: Stream remaining Arrow record batches from the job default stream, + # or from the destination table if the job config set one, with the + # BigQuery client's shared BigQuery Storage client unless the caller passed + # one. Results that fit in the first page never create it. + if bqstorage_client is None: + if row_iterator.client is None: + raise ValueError("RowIterator client is None.") + bqstorage_client = row_iterator.client._get_shared_bqstorage_client() + if bqstorage_client is None: + raise ValueError( + "The google-cloud-bigquery-storage library is required to read Arrow results." + ) + if row_iterator._explicit_destination and row_iterator._table is not None: + reader, pa_schema = _open_destination_stream( + row_iterator, bqstorage_client, timeout, pa_schema + ) + else: + reader = _open_job_stream(row_iterator, bqstorage_client, offset, timeout) + if reader is None: + # An empty destination table has no streams. + yield _empty_batch(pa_schema, row_iterator._schema) + return + 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) + yielded = True + 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() + if not yielded: + yield _empty_batch(pa_schema, row_iterator._schema) 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 a7655f4bddf7..52e6d3b3c3d9 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -49,6 +49,7 @@ import google.api_core.exceptions as core_exceptions from google.api_core import retry as retries +from google.cloud.bigquery import _arrow_query from google.cloud.bigquery import enums from google.cloud.bigquery import job import google.cloud.bigquery.job.query @@ -432,6 +433,7 @@ def query_and_wait( max_results: Optional[int] = None, query_results_format: Optional[str] = None, compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, callback: Callable = lambda _: None, ) -> table.RowIterator: """Run the query, wait for it to finish, and return the results. @@ -481,6 +483,11 @@ def query_and_wait( [Beta] The format for query results (e.g. "ARROW" or :class:`~google.cloud.bigquery.enums.QueryResultsFormat.ARROW`). compression_codec (Optional[Union[str, google.cloud.bigquery.enums.QueryResultsCompressionCodec]]): [Beta] Compression codec for Arrow serialization (e.g. "LZ4_FRAME" or :class:`~google.cloud.bigquery.enums.QueryResultsCompressionCodec.LZ4_FRAME`). + timestamp_precision (Optional[enums.TimestampPrecision]): + [Private Preview] Precision of picosecond ``TIMESTAMP`` columns in + ``ARROW`` results. With ``PICOSECOND``, they arrive at full + precision as ISO 8601 strings. Otherwise, they're truncated to + microseconds. Requires ``query_results_format=ARROW``. callback (Callable): A callback function used by bigframes to report query progress. @@ -501,17 +508,35 @@ def query_and_wait( If ``job_config`` is not an instance of :class:`~google.cloud.bigquery.job.QueryJobConfig` class. + ValueError: + If ``compression_codec`` or ``timestamp_precision`` is set + without ``ARROW``, or if an ``ARROW`` query uses ``WRITE_APPEND``. + google.cloud.bigquery.exceptions.BigQueryError: + If pyarrow or google-cloud-bigquery-storage is missing or too old + to read ``ARROW`` results. The query isn't sent. """ + _arrow_query.validate_query( + query_results_format, + job_config, + compression_codec=compression_codec, + timestamp_precision=timestamp_precision, + ) + request_body = _to_query_request( query=query, job_config=job_config, location=location, timeout=api_timeout ) if query_results_format is not None: request_body["queryResultsFormat"] = query_results_format - if compression_codec is not None: - request_body.setdefault("formatOptions", {}) - request_body["formatOptions"]["arrowSerializationOptions"] = { - "bufferCompression": compression_codec - } + arrow_serialization_options = None + if query_results_format == enums.QueryResultsFormat.ARROW: + arrow_serialization_options = _arrow_query.serialization_options( + compression_codec, timestamp_precision + ) + request_body["arrowSerializationOptions"] = arrow_serialization_options + + # jobs.getQueryResults only returns STRUCT_ENCODING rows, so while waiting + # for an ARROW query, request none of them. + wait_page_size = 0 if arrow_serialization_options is not None else page_size # Some API parameters aren't supported by the jobs.query API. In these # cases, fallback to a jobs.insert call. @@ -533,9 +558,13 @@ def query_and_wait( api_timeout=api_timeout, wait_timeout=wait_timeout, retry=retry, - page_size=page_size, + page_size=wait_page_size, max_results=max_results, query_results_format=query_results_format, + arrow_serialization_options=arrow_serialization_options, + explicit_destination=( + job_config is not None and job_config.destination is not None + ), callback=callback, ) @@ -595,6 +624,8 @@ 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, max_results) if more_pages or not query_results.complete: # TODO(swast): Avoid a call to jobs.get in some cases (few @@ -606,9 +637,10 @@ def do_query(): api_timeout=api_timeout, wait_timeout=wait_timeout, retry=retry, - page_size=page_size, + page_size=wait_page_size, max_results=max_results, query_results_format=query_results_format, + arrow_serialization_options=arrow_serialization_options, callback=callback, ) @@ -649,6 +681,7 @@ def do_query(): started=query_results.started, ended=query_results.ended, query_results_format=query_results_format, + arrow_serialization_options=arrow_serialization_options, ) if job_retry is not None: @@ -690,6 +723,7 @@ def _supported_by_jobs_query(request_body: Dict[str, Any]) -> bool: "reservation", "maxSlots", "queryResultsFormat", + "arrowSerializationOptions", } unsupported_keys = request_keys - keys_allowlist @@ -705,6 +739,8 @@ def _wait_or_cancel( max_results: Optional[int], *, query_results_format: Optional[str] = None, + arrow_serialization_options: Optional[Dict[str, str]] = None, + explicit_destination: bool = False, callback: Callable = lambda _: None, ) -> table.RowIterator: """Wait for a job to complete and return the results. @@ -750,6 +786,8 @@ def _wait_or_cancel( ) ) query_results._query_results_format = query_results_format + query_results._arrow_serialization_options = arrow_serialization_options + query_results._explicit_destination = explicit_destination return query_results except Exception: # Attempt to cancel the job since we can't return the results. diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py index d9afc9781532..1dc1848957fb 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py @@ -128,6 +128,21 @@ def is_read_session_optional(self) -> bool: """ return self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION + @property + def supports_arrow_serialization_options(self) -> bool: + """True if ``BigQueryReadClient.read_rows()`` accepts + ``arrow_serialization_options``, which ``ARROW`` query results need. + + That option was added in the same release as + ``ReadRowsStream.cancel()``, which ``ARROW`` query results also use, so + this checks for ``cancel()`` instead of comparing version numbers. That + also gives the right answer for a build of + google-cloud-bigquery-storage that hasn't been released yet. + """ + from google.cloud.bigquery_storage_v1 import reader + + return hasattr(reader.ReadRowsStream, "cancel") + def try_import(self, raise_if_error: bool = False) -> Any: """Tries to import the bigquery_storage module, and returns results accordingly. It also verifies the module version is recent enough. diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index e1cb5e1bdcdf..7df2bf1b93df 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -26,6 +26,7 @@ import math import os import tempfile +import threading import typing import uuid import warnings @@ -166,6 +167,35 @@ # https://github.com/googleapis/python-bigquery/issues/438 _MIN_GET_QUERY_RESULTS_TIMEOUT = 120 +# While a stream is open and the connection is idle, the shared BigQuery +# Storage client pings the server after a minute, and closes the connection +# if a ping isn't answered within a minute. A dropped ReadRows stream is +# resumed at the next unread row. +_BQSTORAGE_KEEPALIVE_TIME_MS = 60 * 1000 +_BQSTORAGE_KEEPALIVE_TIMEOUT_MS = 60 * 1000 + + +def _create_bqstorage_transport(**kwargs): + """Create a BigQuery Storage gRPC transport whose channel sends keepalive + pings. + + ``BigQueryReadClient`` calls this with the arguments it passes to its + default transport. + """ + from google.cloud import bigquery_storage + + transport_class = bigquery_storage.BigQueryReadClient.get_transport_class("grpc") + + def create_channel(*args, options=(), **channel_kwargs): + options = list(options) + [ + ("grpc.keepalive_time_ms", _BQSTORAGE_KEEPALIVE_TIME_MS), + ("grpc.keepalive_timeout_ms", _BQSTORAGE_KEEPALIVE_TIMEOUT_MS), + ] + return transport_class.create_channel(*args, options=options, **channel_kwargs) + + return transport_class(channel=create_channel, **kwargs) + + _LOAD_TABLE_FROM_DATAFRAME_DEPRECATED = ( "Loading DataFrames via google-cloud-bigquery is deprecated. " "For direct, optimized loading, please call 'pandas_gbq.to_gbq()' directly." @@ -300,6 +330,14 @@ def __init__( self._default_load_job_config = copy.deepcopy(default_load_job_config) self.default_job_creation_mode = default_job_creation_mode + # ARROW query results share one BigQuery Storage client. It's created + # the first time a query needs a stream, and close() closes it. + self._universe_domain = _get_client_universe(client_options) + self._shared_bqstorage_client: Optional[ + "google.cloud.bigquery_storage.BigQueryReadClient" + ] = None + self._shared_bqstorage_client_lock = threading.Lock() + # Use property setter so validation can run. self.default_query_job_config = default_query_job_config @@ -348,6 +386,9 @@ def default_load_job_config(self, value: LoadJobConfig): def close(self): """Close the underlying transport objects, releasing system resources. + This also closes the BigQuery Storage client that ``ARROW`` query + results share. + .. note:: The client instance can be used for making additional requests even @@ -356,6 +397,11 @@ def close(self): """ self._http._auth_request.session.close() self._http.close() + with self._shared_bqstorage_client_lock: + bqstorage_client = self._shared_bqstorage_client + self._shared_bqstorage_client = None + if bqstorage_client is not None: + bqstorage_client._transport.close() def get_service_account_email( self, @@ -600,6 +646,7 @@ def _ensure_bqstorage_client( client_info: Optional[ "google.api_core.gapic_v1.client_info.ClientInfo" ] = DEFAULT_BQSTORAGE_CLIENT_INFO, + transport: Optional[Callable[..., Any]] = None, ) -> Optional["google.cloud.bigquery_storage.BigQueryReadClient"]: """Create a BigQuery Storage API client using this client's credentials. @@ -613,6 +660,9 @@ def _ensure_bqstorage_client( client_info: The client info used with a new BigQuery Storage client instance if one is created. + transport: + A callable that creates the transport of a new BigQuery Storage + client instance. If ``None``, the default transport is used. Returns: A BigQuery Storage API client. @@ -680,10 +730,37 @@ def _ensure_bqstorage_client( credentials=self._credentials, client_options=client_options, client_info=amended_client_info, + transport=transport, ) return bqstorage_client + def _get_shared_bqstorage_client( + self, + ) -> Optional["google.cloud.bigquery_storage.BigQueryReadClient"]: + """Return the BigQuery Storage client that ``ARROW`` query results share. + + The client is created on first use, with this client's credentials and + universe domain and a gRPC channel that sends keepalive pings. + :meth:`close` closes it, and the next call creates a new one. + + Returns: + A BigQuery Storage API client, or ``None`` if + google-cloud-bigquery-storage is missing or outdated. + """ + with self._shared_bqstorage_client_lock: + if self._shared_bqstorage_client is None: + client_options = None + if self._universe_domain != _DEFAULT_UNIVERSE: + client_options = google.api_core.client_options.ClientOptions( + universe_domain=self._universe_domain + ) + self._shared_bqstorage_client = self._ensure_bqstorage_client( + client_options=client_options, + transport=_create_bqstorage_transport, + ) + return self._shared_bqstorage_client + def _dataset_from_arg(self, dataset) -> Union[Dataset, DatasetReference]: if isinstance(dataset, str): dataset = DatasetReference.from_string( @@ -3677,6 +3754,7 @@ def query_and_wait( max_results: Optional[int] = None, query_results_format: Optional[str] = None, compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, ) -> RowIterator: """Run the query, wait for it to finish, and return the results. @@ -3728,6 +3806,11 @@ def query_and_wait( [Beta] The format for query results (e.g. "ARROW" or :class:`~google.cloud.bigquery.enums.QueryResultsFormat.ARROW`). compression_codec (Optional[Union[str, google.cloud.bigquery.enums.QueryResultsCompressionCodec]]): [Beta] Compression codec for Arrow serialization (e.g. "LZ4_FRAME" or :class:`~google.cloud.bigquery.enums.QueryResultsCompressionCodec.LZ4_FRAME`). + timestamp_precision (Optional[enums.TimestampPrecision]): + [Private Preview] Precision of picosecond ``TIMESTAMP`` columns + in ``ARROW`` results. With ``PICOSECOND``, they arrive at full + precision as ISO 8601 strings. Otherwise, they're truncated to + microseconds. Requires ``query_results_format=ARROW``. Returns: google.cloud.bigquery.table.RowIterator: @@ -3746,6 +3829,14 @@ def query_and_wait( If ``job_config`` is not an instance of :class:`~google.cloud.bigquery.job.QueryJobConfig` class. + ValueError: + If ``compression_codec`` or ``timestamp_precision`` is set + without ``ARROW``, or if an ``ARROW`` query uses + ``WRITE_APPEND``, including through + ``default_query_job_config``. + google.cloud.bigquery.exceptions.BigQueryError: + If pyarrow or google-cloud-bigquery-storage is missing or too + old to read ``ARROW`` results. The query isn't sent. """ return self._query_and_wait_bigframes( query, @@ -3760,6 +3851,7 @@ def query_and_wait( max_results=max_results, query_results_format=query_results_format, compression_codec=compression_codec, + timestamp_precision=timestamp_precision, ) def _query_and_wait_bigframes( @@ -3777,6 +3869,7 @@ def _query_and_wait_bigframes( max_results: Optional[int] = None, query_results_format: Optional[str] = None, compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, callback: Callable = lambda _: None, ) -> RowIterator: """See query_and_wait. @@ -3811,6 +3904,7 @@ def _query_and_wait_bigframes( max_results=max_results, query_results_format=query_results_format, compression_codec=compression_codec, + timestamp_precision=timestamp_precision, callback=callback, ) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py index e85e43f6a164..0b7af6293851 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py @@ -1842,9 +1842,13 @@ def is_job_done(): # RowIterator response if there are any rows in it. This prevents us # from stopping the iteration early in the cases where we set # maxResults=0. In that case, we're missing rows and there's no next - # page token. + # page token. ARROW responses from jobs.query carry their rows in + # arrowRecordBatch instead of rows. first_page_response = self._query_results._properties - if "rows" not in first_page_response: + if ( + "rows" not in first_page_response + and "arrowRecordBatch" not in first_page_response + ): first_page_response = None rows = self._client._list_rows_from_query_results( diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index d1d31b14ffd8..1f6dae103493 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -16,7 +16,6 @@ from __future__ import absolute_import -import base64 import copy import datetime import functools @@ -59,6 +58,7 @@ import google.cloud._helpers # type: ignore from google.api_core.page_iterator import HTTPIterator from google.cloud.bigquery import ( + _arrow_query, _helpers, _pandas_helpers, _string_references, @@ -1916,6 +1916,9 @@ class RowIterator(HTTPIterator): If representing query results, the start time of the associated query. ended (Optional[datetime.datetime]): If representing query results, the end time of the associated query. + arrow_serialization_options (Optional[Dict[str, str]]): + The ``arrowSerializationOptions`` sent on ``jobs.query`` for + ``ARROW`` results. ``ReadRows`` requests the same options. """ def __init__( @@ -1944,6 +1947,7 @@ def __init__( started: Optional[datetime.datetime] = None, ended: Optional[datetime.datetime] = None, query_results_format: Optional[str] = None, + arrow_serialization_options: Optional[Dict[str, str]] = None, ): super(RowIterator, self).__init__( client, @@ -1978,6 +1982,10 @@ def __init__( self._job_started = started self._job_ended = ended self._query_results_format = query_results_format + self._arrow_serialization_options = arrow_serialization_options + # Set by query_and_wait() if the query wrote to a destination table + # set in its job config. + self._explicit_destination = False @property @_disallow_when_arrow @@ -2318,91 +2326,17 @@ def _download_arrow_from_job_id( Args: bqstorage_client (Optional[bigquery_storage.BigQueryReadClient]): - Client for BigQuery Storage Read API. If None, one will be created. + Client for BigQuery Storage Read API. If None, the BigQuery + Storage client that the BigQuery client shares is used. timeout (Optional[float]): Timeout in seconds for read operations. Yields: pyarrow.RecordBatch: Record batches generated from the query result. """ - if pyarrow is None: - raise ValueError(_NO_PYARROW_ERROR) - - # Step 1: Ensure BigQuery Read API client is available upfront. - if bqstorage_client is None: - if self.client is None: - raise ValueError("RowIterator client is None.") - bqstorage_client = self.client._ensure_bqstorage_client() - if bqstorage_client is None: - raise ValueError( - "The google-cloud-bigquery-storage library is required to read Arrow results." - ) - - offset = 0 - pa_schema = None - total_rows = self.total_rows - job_complete = False - - # Step 2: Process initial inline Arrow response from jobs.query if available. - if self._first_page_response: - first_page = self._first_page_response - self._first_page_response = None - - job_complete = bool(first_page.get("jobComplete", False)) - if job_complete: - total_rows = int(first_page.get("totalRows", 0)) - - arrow_schema_json = first_page.get("arrowSchema") - if isinstance(arrow_schema_json, dict): - schema_bytes = arrow_schema_json.get("serializedSchema") - if schema_bytes: - if isinstance(schema_bytes, str): - schema_bytes = base64.b64decode(schema_bytes) - pa_schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes)) - - arrow_batch_json = first_page.get("arrowRecordBatch") - if isinstance(arrow_batch_json, dict) and pa_schema is not None: - batch_bytes = arrow_batch_json.get("serializedRecordBatch") - if batch_bytes: - if isinstance(batch_bytes, str): - batch_bytes = base64.b64decode(batch_bytes) - batch = pyarrow.ipc.read_record_batch( - pyarrow.py_buffer(batch_bytes), - pa_schema, - ) - offset += batch.num_rows - yield batch - - # Step 3: Return early if all results were delivered in the first page response. - if job_complete and offset >= total_rows: - return - - # Step 4: Stream remaining Arrow record batches from the job default stream. - project = self._project or (self.client.project if self.client else None) - location = self._location or (self.client.location if self.client else None) - stream_name = f"projects/{project}/locations/{location}/jobs/{self._job_id}/streams/_default" - reader = bqstorage_client.read_rows(stream_name, offset=offset, timeout=timeout) - 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 + return _arrow_query.download_arrow_from_job( + self, bqstorage_client=bqstorage_client, timeout=timeout + ) # If changing the signature of this method, make sure to apply the same # changes to job.QueryJob.to_arrow() @@ -2478,11 +2412,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: @@ -2981,7 +2921,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/noxfile.py b/packages/google-cloud-bigquery/noxfile.py index af201dfb8a05..5f62a16aeaae 100644 --- a/packages/google-cloud-bigquery/noxfile.py +++ b/packages/google-cloud-bigquery/noxfile.py @@ -222,6 +222,14 @@ def mypy(session): """Run type checks with mypy.""" session.install("-e", ".[all]") + # TODO: Remove once google-cloud-bigquery requires a + # google-cloud-bigquery-storage release whose read_rows() accepts + # arrow_serialization_options. Not editable, because mypy can't follow the + # import hook of an editable install. + session.install( + str(CURRENT_DIRECTORY.parent / "google-cloud-bigquery-storage"), + "--no-deps", + ) session.install(MYPY_VERSION) # Just install the dependencies' type info directly, since "mypy --install-types" @@ -288,6 +296,16 @@ def system(session): extras = "[all]" session.install("-e", f".{extras}", "-c", constraints_path) + # TODO: Remove once google-cloud-bigquery requires a + # google-cloud-bigquery-storage release whose read_rows() accepts + # arrow_serialization_options. Until then, query_results_format=ARROW + # needs the in-repo version. + session.install( + "-e", + str(CURRENT_DIRECTORY.parent / "google-cloud-bigquery-storage"), + "--no-deps", + ) + # Test with some broken "extras" in case the user didn't install the extra # directly. For example, pandas-gbq is recommended for pandas features, but # we want to test that we fallback to the previous behavior. For context, diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index 9db96503b00b..2730d169b33b 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -14,11 +14,18 @@ """System tests for Arrow connector.""" +import datetime +import json +import math from typing import Optional +from unittest import mock +import uuid import pyarrow import pytest +from dateutil import relativedelta +from google.api_core import exceptions from google.cloud import bigquery from google.cloud.bigquery import enums @@ -274,3 +281,657 @@ def test_query_and_wait_arrow_format_to_dataframe_iterable(bigquery_client): concat_df = pandas.concat(dfs, ignore_index=True) assert concat_df.columns.tolist() == ["num", "text"] assert concat_df.to_dict(orient="records") == [{"num": 200, "text": "python"}] + + +def test_query_and_wait_arrow_rejects_write_append(bigquery_client, dataset_id): + table_id = f"{bigquery_client.project}.{dataset_id}.arrow_write_append" + bigquery_client.query_and_wait( + f"CREATE OR REPLACE TABLE `{table_id}` AS SELECT 1 AS x" + ) + append = bigquery.QueryJobConfig( + destination=table_id, + write_disposition=enums.WriteDisposition.WRITE_APPEND, + ) + + with pytest.raises(ValueError, match="WRITE_APPEND"): + bigquery_client.query_and_wait( + "SELECT 2 AS x", + job_config=append, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + with bigquery.Client(default_query_job_config=append) as client: + with pytest.raises(ValueError, match="WRITE_APPEND"): + client.query_and_wait( + "SELECT 2 AS x", query_results_format=enums.QueryResultsFormat.ARROW + ) + + # Neither query ran, so the table still holds only its original row. + rows = bigquery_client.query_and_wait(f"SELECT x FROM `{table_id}`") + assert [row.x for row in rows] == [1] + + +@pytest.mark.parametrize( + "codec", + [ + enums.QueryResultsCompressionCodec.LZ4_FRAME, + enums.QueryResultsCompressionCodec.ZSTD, + ], +) +@pytest.mark.parametrize( + "page_size", [None, 100], ids=["first_page_only", "with_stream"] +) +def test_query_and_wait_arrow_compression_matches_uncompressed( + bigquery_client, codec, page_size +): + query = ( + "SELECT x, CAST(x AS STRING) AS s, x / 7 AS f, " + "TIMESTAMP_ADD(TIMESTAMP '2026-01-01', INTERVAL x SECOND) AS ts " + "FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=False) + + def read(**kwargs): + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + page_size=page_size, + query_results_format=enums.QueryResultsFormat.ARROW, + **kwargs, + ) + return pyarrow.Table.from_batches(list(rows.to_arrow_iterable())) + + compressed = read(compression_codec=codec) + + assert compressed.num_rows == 1000 + assert compressed.equals(read()) + + +_MICROSECOND_VALUES = [ + datetime.datetime(2025, 1, day, 0, 0, 0, 123456, tzinfo=datetime.timezone.utc) + for day in (1, 2, 3) +] +_TIMESTAMP_PRECISIONS = pytest.mark.parametrize( + "timestamp_precision", + [ + None, + enums.TimestampPrecision.MICROSECOND, + enums.TimestampPrecision.PICOSECOND, + ], + ids=["default", "microsecond", "picosecond"], +) +_PAGE_SIZES = pytest.mark.parametrize( + "page_size", [None, 1], ids=["first_page_only", "with_stream"] +) + + +def _read_timestamps(bigquery_client, query, timestamp_precision, page_size): + rows = bigquery_client.query_and_wait( + query, + job_config=bigquery.QueryJobConfig(use_query_cache=False), + page_size=page_size, + query_results_format=enums.QueryResultsFormat.ARROW, + timestamp_precision=timestamp_precision, + ) + table = pyarrow.Table.from_batches(list(rows.to_arrow_iterable())) + return table.column("ts").to_pylist() + + +@_PAGE_SIZES +@_TIMESTAMP_PRECISIONS +def test_query_and_wait_arrow_timestamp_precision_microsecond_column( + bigquery_client, timestamp_precision, page_size +): + values = _read_timestamps( + bigquery_client, + "SELECT TIMESTAMP_ADD(TIMESTAMP '2025-01-01 00:00:00.123456 UTC', " + "INTERVAL x DAY) AS ts FROM UNNEST([0, 1, 2]) AS x ORDER BY x", + timestamp_precision, + page_size, + ) + + assert values == _MICROSECOND_VALUES + + +@_PAGE_SIZES +@_TIMESTAMP_PRECISIONS +def test_query_and_wait_arrow_timestamp_precision_picosecond_column( + bigquery_client, timestamp_precision, page_size +): + try: + values = _read_timestamps( + bigquery_client, + "SELECT CAST(FORMAT('2025-01-0%dT00:00:00.123456789012Z', x) " + "AS TIMESTAMP(12)) AS ts FROM UNNEST([1, 2, 3]) AS x ORDER BY x", + timestamp_precision, + page_size, + ) + except exceptions.BadRequest as exc: + # TIMESTAMP(12) is enabled per project. The test project has it, but + # a developer's own project may not. + if "Timestamp precision type parameter is not supported" not in str(exc): + raise + pytest.xfail(f"TIMESTAMP(12) isn't enabled in this project: {exc.message}") + + if timestamp_precision == enums.TimestampPrecision.PICOSECOND: + assert [value[:10] for value in values] == [ + "2025-01-01", + "2025-01-02", + "2025-01-03", + ] + assert all(".123456789012" in value for value in values) + else: + assert values == _MICROSECOND_VALUES + + +@_PAGE_SIZES +@_TIMESTAMP_PRECISIONS +def test_query_and_wait_arrow_timestamp_precision_picosecond_table( + bigquery_client, scalars_table_pico, timestamp_precision, page_size +): + values = _read_timestamps( + bigquery_client, + f"SELECT pico_col AS ts FROM `{scalars_table_pico}` ORDER BY pico_col", + timestamp_precision, + page_size, + ) + + if timestamp_precision == enums.TimestampPrecision.PICOSECOND: + assert [value[:10] for value in values] == [ + "2025-01-01", + "2025-01-02", + "2025-01-03", + ] + assert all(".123456789012" in value for value in values) + else: + assert values == _MICROSECOND_VALUES + + +@pytest.mark.parametrize( + ("page_size", "priority", "expected_offset"), + [ + (100, None, 100), + # jobs.query doesn't accept a priority, so this case runs through + # jobs.insert, which returns no rows inline. + (None, enums.QueryPriority.INTERACTIVE, 0), + ], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_matches_struct_encoding( + bigquery_client, bqstorage_client, page_size, priority, expected_offset +): + query = ( + "SELECT x, CAST(x AS STRING) AS s, MOD(x, 2) = 0 AS even, " + "DATE_ADD(DATE '2026-01-01', INTERVAL x DAY) AS d " + "FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + # Wrap read_rows() to check where the stream starts. + read_rows = mock.Mock(wraps=bqstorage_client.read_rows) + + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + page_size=page_size, + 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)) + + assert pyarrow.Table.from_batches(batches).to_pylist() == expected + read_rows.assert_called_once_with( + f"projects/{bigquery_client.project}/locations/{rows.location}" + f"/jobs/{rows.job_id}/streams/_default", + offset=expected_offset, + timeout=mock.ANY, + arrow_serialization_options=mock.ANY, + ) + + +def _interval_to_relativedelta(value): + return relativedelta.relativedelta( + months=value.months, days=value.days, microseconds=value.nanoseconds // 1000 + ) + + +def _arrow_rows(table): + """Return the rows of an Arrow table as STRUCT_ENCODING returns them. + + Arrow has JSON values as strings and INTERVAL values as MonthDayNano, + while STRUCT_ENCODING rows have parsed JSON and relativedelta values. + """ + converters = {} + for field in table.schema: + if field.metadata == {b"ARROW:extension:name": b"google:sqlType:json"}: + converters[field.name] = json.loads + elif pyarrow.types.is_interval(field.type): + converters[field.name] = _interval_to_relativedelta + rows = table.to_pylist() + for row in rows: + for name, convert in converters.items(): + if row[name] is not None: + row[name] = convert(row[name]) + return rows + + +@pytest.mark.parametrize( + ("page_size", "priority", "expected_offset"), + [ + (1, None, 1), + # jobs.query doesn't accept a priority, so this case runs through + # jobs.insert, which returns no rows inline. + (None, enums.QueryPriority.INTERACTIVE, 0), + ], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_locations( + bigquery_client, + bqstorage_client, + scalars_table_multi_location, + page_size, + priority, + expected_offset, +): + location, table_id = scalars_table_multi_location + query = f"SELECT * FROM `{table_id}` ORDER BY rowindex" + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + read_rows = mock.Mock(wraps=bqstorage_client.read_rows) + + # The query doesn't name a location, so only jobReference has it. + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + page_size=page_size, + 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)) + + assert rows.location == location + assert _arrow_rows(pyarrow.Table.from_batches(batches)) == expected + read_rows.assert_called_once_with( + f"projects/{bigquery_client.project}/locations/{location}" + f"/jobs/{rows.job_id}/streams/_default", + offset=expected_offset, + timeout=mock.ANY, + arrow_serialization_options=mock.ANY, + ) + + +_DATA_TYPES_QUERY = """ +SELECT + i, + IF(i = 2, NULL, i = 1) AS bool_col, + IF(i = 2, NULL, i * 1000000000000) AS int64_col, + IF(i = 2, NULL, i / 3) AS float64_col, + IF(i = 2, NULL, NUMERIC '-123456789.123456789') AS numeric_col, + IF(i = 2, NULL, BIGNUMERIC '1.23456789012345678901234567890123456789') AS bignumeric_col, + IF(i = 2, NULL, FORMAT('row %d', i)) AS string_col, + IF(i = 2, NULL, CAST(FORMAT('row %d', i) AS BYTES)) AS bytes_col, + IF(i = 2, NULL, DATE_ADD(DATE '2026-01-01', INTERVAL i DAY)) AS date_col, + IF(i = 2, NULL, DATETIME '2026-01-01 12:34:56.789012') AS datetime_col, + IF(i = 2, NULL, TIME '12:34:56.789012') AS time_col, + IF(i = 2, NULL, TIMESTAMP '2026-01-01 12:34:56.789012 UTC') AS timestamp_col, + IF(i = 2, NULL, INTERVAL '1-2 3 4:5:6.789012' YEAR TO SECOND) AS interval_col, + IF(i = 2, NULL, ST_GEOGPOINT(-122.0838511, 37.3860517)) AS geography_col, + IF(i = 2, NULL, JSON '{"a": [1, 2.5, "x"], "b": null, "c": {"d": true}}') AS json_col, + IF(i = 2, NULL, RANGE(DATE '2026-01-01', DATE '2026-02-01')) AS range_date_col, + IF(i = 2, NULL, RANGE(DATETIME '2026-01-01 00:00:00', NULL)) AS range_datetime_col, + IF( + i = 2, + NULL, + RANGE(TIMESTAMP '2026-01-01 00:00:00 UTC', TIMESTAMP '2026-01-02 00:00:00 UTC') + ) AS range_timestamp_col, + IF(i = 2, [], [i, i + 1]) AS array_col, + IF(i = 2, NULL, STRUCT(i AS x, [STRUCT('a' AS s, [1.5, 2.5] AS f)] AS nested)) AS struct_col +FROM UNNEST([1, 2, 3]) AS i +ORDER BY i +""" + + +def test_query_and_wait_arrow_data_types(bigquery_client): + job_config = bigquery.QueryJobConfig(use_query_cache=False) + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait( + _DATA_TYPES_QUERY, job_config=job_config + ) + ] + + rows = bigquery_client.query_and_wait( + _DATA_TYPES_QUERY, + job_config=job_config, + # The first row arrives in the first page, the others in the stream. + page_size=1, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + batches = list(rows.to_arrow_iterable()) + + assert batches[0].num_rows == 1 + assert _arrow_rows(pyarrow.Table.from_batches(batches)) == expected + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["jobs_query", "jobs_insert"], +) +def test_query_and_wait_arrow_unnamed_struct_fields(bigquery_client, priority): + query = "SELECT (1, 'a') AS s, STRUCT(2 AS x, 'b') AS p" + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + # jobs.query doesn't accept a priority, so this runs through jobs.insert. + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + + try: + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + # jobs.query fails with a jobInternalError, which job_retry would + # retry until its deadline. + job_retry=None, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + except exceptions.BadRequest as exc: + if not any(error.get("reason") == "jobInternalError" for error in exc.errors): + raise + pytest.xfail(f"jobs.query can't send unnamed STRUCT fields: {exc.message}") + table = pyarrow.Table.from_batches(list(rows.to_arrow_iterable())) + + assert table.to_pylist() == expected + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["first_page_only", "no_rows_inline"], +) +def test_query_and_wait_arrow_extreme_values( + bigquery_client, scalars_extreme_table, priority +): + # BigQuery fails the extreme INTERVAL values with "Out of range conversion + # for microseconds value", since Arrow's interval type can't hold them. + query = ( + f"SELECT * EXCEPT (interval_col) FROM `{scalars_extreme_table}` " + "ORDER BY rowindex" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + expected = [ + dict(row.items()) + for row in bigquery_client.query_and_wait(query, job_config=job_config) + ] + + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + actual = _arrow_rows(pyarrow.Table.from_batches(list(rows.to_arrow_iterable()))) + + # NaN never equals itself, so check it on its own. + assert math.isnan(actual[2].pop("float64_col")) + assert math.isnan(expected[2].pop("float64_col")) + assert actual == expected + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_cached_results(bigquery_client, priority): + # No other query uses this literal, so only the second run can hit the + # cache. + query = ( + f"SELECT x, '{uuid.uuid4().hex}' AS key " + "FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=True) + if priority is not None: + job_config.priority = priority + + def read(): + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + page_size=100, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + table = pyarrow.Table.from_batches(list(rows.to_arrow_iterable())) + return bigquery_client.get_job(rows.job_id, location=rows.location), table + + first_job, first = read() + cached_job, cached = read() + + assert not first_job.cache_hit + assert cached_job.cache_hit + assert cached.num_rows == 1000 + assert cached.equals(first) + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["jobs_query", "jobs_insert"], +) +def test_query_and_wait_arrow_failed_query(bigquery_client, priority): + # The last row fails, so the query fails while it runs. + query = ( + "SELECT IF(x < 1000, x, ERROR(FORMAT('failed at row %d', x))) AS x " + "FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x" + ) + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + job_config.priority = priority + + with pytest.raises(exceptions.GoogleAPICallError) as struct_encoding_error: + bigquery_client.query_and_wait(query, job_config=job_config) + with pytest.raises(exceptions.GoogleAPICallError) as arrow_error: + rows = bigquery_client.query_and_wait( + query, + job_config=job_config, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + list(rows.to_arrow_iterable()) + + 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 + + +def _query_arrow(client, row_count): + return client.query_and_wait( + f"SELECT x FROM UNNEST(GENERATE_ARRAY(1, {row_count})) AS x ORDER BY x", + job_config=bigquery.QueryJobConfig(use_query_cache=False), + page_size=100, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + + +def test_query_and_wait_arrow_queries_share_one_bqstorage_client(): + from google.cloud.bigquery import client as client_module + + with bigquery.Client() as client: + # Wrap the client factory to count how many clients get created. + ensure = mock.Mock(wraps=client._ensure_bqstorage_client) + with mock.patch.object(client, "_ensure_bqstorage_client", ensure): + tables = [ + pyarrow.Table.from_batches( + list(_query_arrow(client, 1000).to_arrow_iterable()) + ), + _query_arrow(client, 1000).to_arrow(), + pyarrow.Table.from_batches( + list(_query_arrow(client, 1000).to_arrow_iterable()) + ), + ] + shared = client._shared_bqstorage_client + + for table in tables: + assert table.column("x").to_pylist() == list(range(1, 1001)) + ensure.assert_called_once_with( + client_options=None, transport=client_module._create_bqstorage_transport + ) + assert shared is not None + # Leaving the with block calls close(), which closes the shared channel. + assert client._shared_bqstorage_client is None + with pytest.raises(ValueError, match="(?i)closed"): + shared.read_rows("projects/p/locations/US/jobs/j/streams/_default") + + +def test_query_and_wait_arrow_ending_a_read_early_cancels_only_its_stream(): + from google.cloud.bigquery_storage_v1 import reader + + # Wrap ReadRowsStream.cancel() to count calls without changing it. + cancel_stream = reader.ReadRowsStream.cancel + with bigquery.Client() as client, mock.patch.object( + reader.ReadRowsStream, "cancel", autospec=True, side_effect=cancel_stream + ) as cancel: + ended_early_rows = _query_arrow(client, 100000) + ended_early = ended_early_rows.to_arrow_iterable() + read_fully = _query_arrow(client, 100000).to_arrow_iterable() + # Two reads reach the stream, whether or not jobs.query returned a + # first page. + next(ended_early) + next(ended_early) + rows = [next(read_fully), next(read_fully)] + + ended_early.close() + rows.extend(read_fully) + + (stream,) = cancel.call_args.args + cancel.assert_called_once_with(stream) + assert stream._name.endswith(f"/jobs/{ended_early_rows.job_id}/streams/_default") + table = pyarrow.Table.from_batches(rows) + assert table.column("x").to_pylist() == list(range(1, 100001)) + + +@pytest.mark.parametrize( + "write_disposition", + [None, enums.WriteDisposition.WRITE_TRUNCATE], + ids=["default", "overwrite"], +) +def test_query_and_wait_arrow_explicit_destination( + bigquery_client, bqstorage_client, dataset_id, write_disposition +): + table_name = f"arrow_destination_{write_disposition or 'default'}".lower() + table_id = f"{bigquery_client.project}.{dataset_id}.{table_name}" + job_config = bigquery.QueryJobConfig(destination=table_id) + if write_disposition is not None: + # The rows already in the table are replaced, not read. + bigquery_client.query_and_wait( + f"CREATE OR REPLACE TABLE `{table_id}` AS SELECT -1 AS x" + ) + job_config.write_disposition = write_disposition + # Wrap create_read_session() to check how the table is read. + create_read_session = mock.Mock(wraps=bqstorage_client.create_read_session) + + rows = bigquery_client.query_and_wait( + "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x", + job_config=job_config, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + with mock.patch.object( + bqstorage_client, "create_read_session", create_read_session + ): + 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, 1001)) + kwargs = create_read_session.call_args.kwargs + assert kwargs["read_session"].table == ( + f"projects/{bigquery_client.project}/datasets/{dataset_id}/tables/{table_name}" + ) + assert kwargs["max_stream_count"] == 1 + + +@pytest.mark.parametrize( + "priority", + [None, enums.QueryPriority.INTERACTIVE], + ids=["jobs_query", "jobs_insert"], +) +def test_query_and_wait_arrow_zero_rows_keep_the_schema(bigquery_client, priority): + job_config = bigquery.QueryJobConfig(use_query_cache=False) + if priority is not None: + # jobs.query doesn't accept a priority, so this runs through jobs.insert. + job_config.priority = priority + + def query(): + return bigquery_client.query_and_wait( + "SELECT x, CAST(x AS STRING) AS s " + "FROM UNNEST(CAST([] AS ARRAY)) AS x", + job_config=job_config, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + + batches = list(query().to_arrow_iterable()) + + assert len(batches) == 1 + for table in (pyarrow.Table.from_batches(batches), query().to_arrow()): + assert table.num_rows == 0 + assert table.schema.names == ["x", "s"] + assert table.schema.field("x").type == pyarrow.int64() + assert table.schema.field("s").type == pyarrow.string() + + +def test_query_and_wait_arrow_statements_without_results(bigquery_client, dataset_id): + table_id = f"{bigquery_client.project}.{dataset_id}.arrow_statements" + + for statement in ( + f"CREATE OR REPLACE TABLE `{table_id}` AS SELECT 1 AS x", + f"UPDATE `{table_id}` SET x = 2 WHERE TRUE", + ): + rows = bigquery_client.query_and_wait( + statement, query_results_format=enums.QueryResultsFormat.ARROW + ) + (batch,) = rows.to_arrow_iterable() + assert batch.num_rows == 0 + + assert rows.num_dml_affected_rows == 1 + assert [ + row.x for row in bigquery_client.query_and_wait(f"SELECT x FROM `{table_id}`") + ] == [2] diff --git a/packages/google-cloud-bigquery/tests/system/test_query.py b/packages/google-cloud-bigquery/tests/system/test_query.py index 69ff08b66073..43ce8d293403 100644 --- a/packages/google-cloud-bigquery/tests/system/test_query.py +++ b/packages/google-cloud-bigquery/tests/system/test_query.py @@ -14,6 +14,7 @@ import datetime import decimal +import json from typing import Tuple import pytest @@ -162,7 +163,7 @@ def test_query_statistics(bigquery_client, query_api_method): assert len(plan) > stages_with_inputs -@pytest.mark.parametrize( +_QUERY_PARAMETER_CASES = pytest.mark.parametrize( ("sql", "expected", "query_parameters"), ( ( @@ -438,6 +439,9 @@ def test_query_statistics(bigquery_client, query_api_method): ), ), ) + + +@_QUERY_PARAMETER_CASES def test_query_parameters( bigquery_client, query_api_method, sql, expected, query_parameters ): @@ -454,6 +458,35 @@ def test_query_parameters( assert rows[0][0] == expected +@_QUERY_PARAMETER_CASES +def test_query_parameters_arrow(bigquery_client, sql, expected, query_parameters): + job_config = bigquery.QueryJobConfig( + query_parameters=query_parameters, use_query_cache=False + ) + try: + rows = bigquery_client.query_and_wait( + sql, + job_config=job_config, + # jobs.query fails unnamed STRUCT fields with a jobInternalError, + # which job_retry would retry until its deadline. + job_retry=None, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + except exceptions.BadRequest as exc: + if not any(error.get("reason") == "jobInternalError" for error in exc.errors): + raise + pytest.xfail(f"jobs.query can't send unnamed STRUCT fields: {exc.message}") + (batch,) = rows.to_arrow_iterable() + assert batch.num_rows == 1 + assert batch.num_columns == 1 + (value,) = batch.column(0).to_pylist() + if batch.schema.field(0).metadata == { + b"ARROW:extension:name": b"google:sqlType:json" + }: + value = json.loads(value) + assert value == expected + + def test_dry_run( bigquery_client: bigquery.Client, query_api_method: str, diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_query.py b/packages/google-cloud-bigquery/tests/unit/job/test_query.py index d6727ae0355e..daa3d4bd4dc7 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_query.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_query.py @@ -1279,6 +1279,42 @@ def test_result_with_done_jobs_query_response_doesnt_call_get_query_results(self timeout=DEFAULT_GET_JOB_TIMEOUT, ) + def test_result_with_done_jobs_query_response_keeps_arrow_first_page(self): + """An ARROW response from jobs.query is kept as the first page, like + rows are, so the stream can continue after it. + """ + job_resource = self._make_resource(started=True, ended=True, location="EU") + conn = make_connection(job_resource) + client = _make_client(self.PROJECT, connection=conn) + query_resource_done = { + "jobComplete": True, + "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID}, + "arrowSchema": {"serializedSchema": "c2NoZW1h"}, + "arrowRecordBatch": {"serializedRecordBatch": "YmF0Y2g="}, + "pageRowCount": "1", + "totalRows": "2", + } + job = google.cloud.bigquery._job_helpers._to_query_job( + client, + "SELECT 1", + request_config=None, + query_response=query_resource_done, + ) + + result = job.result(page_size=0) + + self.assertEqual( + result._first_page_response["arrowRecordBatch"], + query_resource_done["arrowRecordBatch"], + ) + # Only jobs.get, no jobs.getQueryResults. + conn.api_request.assert_called_once_with( + method="GET", + path=f"/projects/{self.PROJECT}/jobs/{self.JOB_ID}", + query_params={"projection": "full"}, + timeout=DEFAULT_GET_JOB_TIMEOUT, + ) + def test_result_with_none_timeout(self): # Verifies that with an intentional None timeout, get job uses None # instead of the default timeout. diff --git a/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py b/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py index aec553bb7b19..2dc8ff8e0a93 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py @@ -200,6 +200,31 @@ def test_bqstorage_is_read_session_optional_false(): assert not bqstorage_versions.is_read_session_optional +@pytest.mark.skipif( + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" +) +def test_bqstorage_supports_arrow_serialization_options_true(): + from google.cloud.bigquery_storage_v1 import reader + + bqstorage_versions = _versions_helpers.BQStorageVersions() + with mock.patch.object(reader.ReadRowsStream, "cancel", create=True): + assert bqstorage_versions.supports_arrow_serialization_options + + +@pytest.mark.skipif( + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" +) +def test_bqstorage_supports_arrow_serialization_options_false(): + from google.cloud.bigquery_storage_v1 import reader + + class ReadRowsStream: + """Stands in for a ReadRowsStream from before cancel() was added.""" + + bqstorage_versions = _versions_helpers.BQStorageVersions() + with mock.patch.object(reader, "ReadRowsStream", ReadRowsStream): + assert not bqstorage_versions.supports_arrow_serialization_options + + @pytest.mark.skipif(pandas is None, reason="pandas is not installed") @pytest.mark.parametrize("version", ["1.1.5", "2.0.0", "2.1.0"]) def test_try_import_raises_no_error_w_recent_pandas(version): diff --git a/packages/google-cloud-bigquery/tests/unit/test_client.py b/packages/google-cloud-bigquery/tests/unit/test_client.py index 88629500a007..7f49ddfadc33 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_client.py +++ b/packages/google-cloud-bigquery/tests/unit/test_client.py @@ -1033,6 +1033,128 @@ def test_ensure_bqstorage_client_existing_client_check_fails(self): ] assert matching_warnings, "Obsolete dependency warning not raised." + def test_get_shared_bqstorage_client_creates_one_client(self): + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud.bigquery import client as client_module + + mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + + with mock.patch( + "google.cloud.bigquery_storage.BigQueryReadClient", mock_client + ): + first = client._get_shared_bqstorage_client() + second = client._get_shared_bqstorage_client() + + self.assertIs(first, mock_client.return_value) + self.assertIs(second, first) + mock_client.assert_called_once_with( + credentials=creds, + client_options=None, + client_info=mock.ANY, + transport=client_module._create_bqstorage_transport, + ) + + @pytest.mark.skipif( + packaging.version.parse(getattr(google.api_core, "__version__", "0.0.0")) + < packaging.version.Version("2.15.0"), + reason="universe_domain not supported with google-api-core < 2.15.0", + ) + def test_get_shared_bqstorage_client_uses_universe_domain(self): + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + + mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + creds = _make_credentials() + client = self._make_one( + project=self.PROJECT, + credentials=creds, + client_options={"universe_domain": "foo.com"}, + ) + + with mock.patch( + "google.cloud.bigquery_storage.BigQueryReadClient", mock_client + ): + client._get_shared_bqstorage_client() + + _, kwargs = mock_client.call_args + self.assertEqual(kwargs["client_options"].universe_domain, "foo.com") + + def test_get_shared_bqstorage_client_tries_again_if_unavailable(self): + from google.cloud.bigquery import client as client_module + + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + + with mock.patch.object( + client, "_ensure_bqstorage_client", return_value=None + ) as ensure: + self.assertIsNone(client._get_shared_bqstorage_client()) + self.assertIsNone(client._get_shared_bqstorage_client()) + + expected = mock.call( + client_options=None, transport=client_module._create_bqstorage_transport + ) + self.assertEqual(ensure.call_args_list, [expected, expected]) + + def test_get_shared_bqstorage_client_creates_one_client_across_threads(self): + import threading + import time + + from google.cloud.bigquery import client as client_module + + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + + def slow_create(**kwargs): + time.sleep(0.05) + return mock.Mock() + + results = [] + + def get_client(): + results.append(client._get_shared_bqstorage_client()) + + with mock.patch.object( + client, "_ensure_bqstorage_client", side_effect=slow_create + ) as ensure: + threads = [threading.Thread(target=get_client) for _ in range(5)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + ensure.assert_called_once_with( + client_options=None, transport=client_module._create_bqstorage_transport + ) + self.assertEqual(len(results), 5) + self.assertTrue(all(result is results[0] for result in results)) + + def test_get_shared_bqstorage_client_channel_sends_keepalive_pings(self): + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + import google.auth.credentials + import grpc + + transport_class = bigquery_storage.BigQueryReadClient.get_transport_class( + "grpc" + ) + creds = google.auth.credentials.AnonymousCredentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + channel = grpc.insecure_channel("localhost:1") + + with mock.patch.object( + transport_class, "create_channel", return_value=channel + ) as create_channel: + bqstorage_client = client._get_shared_bqstorage_client() + channel.close() + + self.assertIsInstance(bqstorage_client._transport, transport_class) + options = create_channel.call_args.kwargs["options"] + self.assertIn(("grpc.keepalive_time_ms", 60000), options) + self.assertIn(("grpc.keepalive_timeout_ms", 60000), options) + # The transport's own options are kept. + self.assertIn(("grpc.max_receive_message_length", -1), options) + def test_create_routine_w_minimal_resource(self): from google.cloud.bigquery.routine import Routine, RoutineReference @@ -1653,6 +1775,23 @@ def test_close(self): http.close.assert_called_once() http._auth_request.session.close.assert_called_once() + def test_close_closes_shared_bqstorage_client(self): + creds = _make_credentials() + http = mock.Mock() + http._auth_request.session = mock.Mock() + client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) + first, second = mock.Mock(), mock.Mock() + + with mock.patch.object( + client, "_ensure_bqstorage_client", side_effect=[first, second] + ): + self.assertIs(client._get_shared_bqstorage_client(), first) + client.close() + self.assertIs(client._get_shared_bqstorage_client(), second) + + first._transport.close.assert_called_once_with() + second._transport.close.assert_not_called() + def test_get_model(self): path = "projects/%s/datasets/%s/models/%s" % ( self.PROJECT, 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 3bc19b4d3ccd..d602b01af3d9 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,12 +15,45 @@ import base64 import unittest from unittest import mock +import warnings +import google.api_core.exceptions as core_exceptions +import packaging.version import pytest +from google.cloud.bigquery import _arrow_query from google.cloud.bigquery import _job_helpers +from google.cloud.bigquery import _versions_helpers +from google.cloud.bigquery import exceptions from google.cloud.bigquery.client import Client -from google.cloud.bigquery.table import RowIterator, _EmptyRowIterator +from google.cloud.bigquery.enums import QueryResultsCompressionCodec +from google.cloud.bigquery.enums import TimestampPrecision +from google.cloud.bigquery.enums import WriteDisposition +from google.cloud.bigquery.job import QueryJobConfig +from google.cloud.bigquery.table import RowIterator, TableReference, _EmptyRowIterator + +from .helpers import make_client + +# The Arrow serialization options that jobs.query and ReadRows get when the +# caller sets none. +_DEFAULT_QUERY_OPTIONS = { + "bufferCompression": "COMPRESSION_UNSPECIFIED", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", +} +_DEFAULT_READ_OPTIONS = { + "buffer_compression": "COMPRESSION_UNSPECIFIED", + "picos_timestamp_precision": "TIMESTAMP_PRECISION_MICROS", +} + +# The real check, saved before the autouse fixture below replaces it. +_check_dependencies = _arrow_query._check_dependencies + + +@pytest.fixture(autouse=True) +def check_dependencies(): + """Skip the ARROW dependency check, which depends on installed packages.""" + with mock.patch.object(_arrow_query, "_check_dependencies") as check: + yield check class TestQueryResultsFormatOption1(unittest.TestCase): @@ -75,12 +108,15 @@ def test_job_helpers_query_and_wait_sets_compression_codec(self): ) call_args = client._call_api.call_args - self.assertIn("formatOptions", call_args.kwargs["data"]) + self.assertNotIn( + "arrowSerializationOptions", call_args.kwargs["data"]["formatOptions"] + ) self.assertEqual( - call_args.kwargs["data"]["formatOptions"]["arrowSerializationOptions"][ - "bufferCompression" - ], - "LZ4_FRAME", + call_args.kwargs["data"]["arrowSerializationOptions"], + { + "bufferCompression": "LZ4_FRAME", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", + }, ) def test_job_helpers_query_and_wait_fallback_preserves_query_results_format(self): @@ -175,7 +211,7 @@ def test_row_iterator_to_arrow_iterable_delegates_when_format_is_arrow(self): def test_download_arrow_from_job_id_constructs_stream_and_reads(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage iterator = RowIterator( client=mock_client, @@ -193,22 +229,29 @@ def test_download_arrow_from_job_id_constructs_stream_and_reads(self): mock_response.arrow_record_batch = None mock_bqstorage.read_rows.return_value = [mock_response] - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) - self.assertEqual(batches, []) + # No rows and no schema: one empty batch without columns. + self.assertEqual( + batches, [mock_pyarrow.RecordBatch.from_arrays.return_value] + ) + mock_pyarrow.schema.assert_called_once_with([]) expected_stream_name = ( "projects/test-proj/locations/US/jobs/test-job-456/streams/_default" ) mock_bqstorage.read_rows.assert_called_once_with( - expected_stream_name, offset=0, timeout=5.0 + expected_stream_name, + offset=0, + timeout=5.0, + arrow_serialization_options=_DEFAULT_READ_OPTIONS, ) def test_download_arrow_from_job_id_with_first_page_response(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage raw_schema_bytes = b"schema_bytes_123" raw_batch_bytes = b"batch_bytes_123" @@ -247,7 +290,7 @@ def test_download_arrow_from_job_id_with_first_page_response(self): mock_bqstorage.read_rows.return_value = [mock_response] - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema" mock_pyarrow.ipc.read_record_batch.side_effect = [ @@ -258,7 +301,8 @@ def test_download_arrow_from_job_id_with_first_page_response(self): batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_first_batch, mock_second_batch]) - self.assertIsNone(iterator._first_page_response) + # The first page is kept for the next read. + self.assertIsNotNone(iterator._first_page_response) mock_pyarrow.ipc.read_schema.assert_called_once_with(raw_schema_bytes) mock_pyarrow.ipc.read_record_batch.assert_has_calls( @@ -272,13 +316,16 @@ def test_download_arrow_from_job_id_with_first_page_response(self): "projects/test-proj/locations/US/jobs/test-job-sync/streams/_default" ) mock_bqstorage.read_rows.assert_called_once_with( - expected_stream_name, offset=10, timeout=5.0 + expected_stream_name, + offset=10, + timeout=5.0, + arrow_serialization_options=_DEFAULT_READ_OPTIONS, ) def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage raw_schema_bytes = b"schema_bytes_456" first_page_response = { @@ -306,7 +353,7 @@ def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): mock_bqstorage.read_rows.return_value = [mock_response] - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema" mock_pyarrow.ipc.read_record_batch.return_value = mock_stream_batch @@ -314,19 +361,23 @@ def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_stream_batch]) - self.assertIsNone(iterator._first_page_response) + # The first page is kept for the next read. + self.assertIsNotNone(iterator._first_page_response) mock_pyarrow.ipc.read_schema.assert_called_once_with(raw_schema_bytes) mock_pyarrow.ipc.read_record_batch.assert_called_once_with( b"stream_batch_bytes", "deserialized_schema" ) expected_stream_name = "projects/test-proj/locations/US/jobs/test-job-schema-only/streams/_default" mock_bqstorage.read_rows.assert_called_once_with( - expected_stream_name, offset=0, timeout=5.0 + expected_stream_name, + offset=0, + timeout=5.0, + arrow_serialization_options=_DEFAULT_READ_OPTIONS, ) def test_download_arrow_from_job_id_missing_storage_client_raises_value_error(self): mock_client = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = None + mock_client._get_shared_bqstorage_client.return_value = None iterator = RowIterator( client=mock_client, @@ -339,7 +390,7 @@ def test_download_arrow_from_job_id_missing_storage_client_raises_value_error(se query_results_format="ARROW", ) - with mock.patch("google.cloud.bigquery.table.pyarrow"): + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow"): with self.assertRaises(ValueError) as ctx: list(iterator._download_arrow_from_job_id()) self.assertIn( @@ -350,7 +401,7 @@ def test_download_arrow_from_job_id_missing_storage_client_raises_value_error(se def test_download_arrow_from_job_id_with_schema_and_batch(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage iterator = RowIterator( client=mock_client, @@ -374,7 +425,7 @@ def test_download_arrow_from_job_id_with_schema_and_batch(self): mock_bqstorage.read_rows.return_value = [mock_response] - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x mock_pyarrow.ipc.read_schema.return_value = "fake_schema" mock_pyarrow.ipc.read_record_batch.return_value = "fake_batch" @@ -422,19 +473,21 @@ def test_download_arrow_from_job_id_avoids_read_rows_when_all_rows_present(self) mock_first_batch = mock.MagicMock() mock_first_batch.num_rows = 10 - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema" mock_pyarrow.ipc.read_record_batch.return_value = mock_first_batch batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_first_batch]) - mock_client._ensure_bqstorage_client.assert_called_once() + mock_client._get_shared_bqstorage_client.assert_not_called() + batches_second = list(iterator._download_arrow_from_job_id(timeout=5.0)) + self.assertEqual(batches_second, [mock_first_batch]) def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage raw_schema_bytes = b"schema_bytes_789" raw_batch_bytes = b"batch_bytes_789" @@ -464,17 +517,20 @@ def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self): mock_first_batch.num_rows = 10 mock_bqstorage.read_rows.return_value = [] - with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow: + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow: mock_pyarrow.py_buffer = lambda x: x mock_pyarrow.ipc.read_schema.return_value = "deserialized_schema" mock_pyarrow.ipc.read_record_batch.return_value = mock_first_batch batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_first_batch]) - mock_client._ensure_bqstorage_client.assert_called_once() + mock_client._get_shared_bqstorage_client.assert_called_once_with() expected_stream = "projects/test-proj/locations/US/jobs/test-job-incomplete/streams/_default" mock_bqstorage.read_rows.assert_called_once_with( - expected_stream, offset=10, timeout=5.0 + expected_stream, + offset=10, + timeout=5.0, + arrow_serialization_options=_DEFAULT_READ_OPTIONS, ) def test_enums_values(self): @@ -517,16 +573,14 @@ def test_job_helpers_query_and_wait_accepts_enums(self): call_args = client._call_api.call_args self.assertEqual(call_args.kwargs["data"]["queryResultsFormat"], "ARROW") self.assertEqual( - call_args.kwargs["data"]["formatOptions"]["arrowSerializationOptions"][ - "bufferCompression" - ], + call_args.kwargs["data"]["arrowSerializationOptions"]["bufferCompression"], "LZ4_FRAME", ) def test_row_iterator_to_dataframe_iterable_when_format_is_arrow(self): pytest.importorskip("pandas") mock_client = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock.MagicMock() + mock_client._get_shared_bqstorage_client.return_value = mock.MagicMock() iterator = RowIterator( client=mock_client, @@ -549,7 +603,7 @@ def test_row_iterator_to_dataframe_iterable_when_format_is_arrow(self): def test_row_iterator_to_dataframe_when_format_is_arrow(self): pytest.importorskip("pandas") mock_client = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock.MagicMock() + mock_client._get_shared_bqstorage_client.return_value = mock.MagicMock() iterator = RowIterator( client=mock_client, @@ -569,5 +623,876 @@ def test_row_iterator_to_dataframe_when_format_is_arrow(self): self.assertEqual(df, "full_df") +def _query_and_wait(client, **kwargs): + return _job_helpers.query_and_wait( + client, + "SELECT 1", + project="p", + location="us", + retry=None, + job_retry=None, + **{"job_config": None, **kwargs}, + ) + + +def test_query_and_wait_compression_codec_requires_arrow(check_dependencies): + client = mock.MagicMock(spec=Client) + + with pytest.raises(ValueError, match="compression_codec requires"): + _query_and_wait(client, compression_codec="LZ4_FRAME") + + client._call_api.assert_not_called() + check_dependencies.assert_not_called() + + +def test_query_and_wait_without_arrow_skips_dependency_check(check_dependencies): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "rows": [], + "schema": {"fields": []}, + } + + _query_and_wait(client) + + check_dependencies.assert_not_called() + + +def test_query_and_wait_arrow_rejects_write_append(): + client = mock.MagicMock(spec=Client) + job_config = QueryJobConfig(write_disposition=WriteDisposition.WRITE_APPEND) + + with pytest.raises(ValueError, match="WRITE_APPEND"): + _query_and_wait(client, job_config=job_config, query_results_format="ARROW") + + client._call_api.assert_not_called() + + +def test_query_and_wait_arrow_rejects_default_write_append(): + client = make_client( + default_query_job_config=QueryJobConfig( + destination="p.d.t", write_disposition=WriteDisposition.WRITE_APPEND + ) + ) + client._call_api = mock.Mock() + + with pytest.raises(ValueError, match="WRITE_APPEND"): + client.query_and_wait("SELECT 1", query_results_format="ARROW") + + client._call_api.assert_not_called() + + +@pytest.mark.parametrize( + "write_disposition", + [None, WriteDisposition.WRITE_TRUNCATE, WriteDisposition.WRITE_EMPTY], +) +def test_validate_query_allows_other_write_dispositions( + check_dependencies, write_disposition +): + job_config = QueryJobConfig(write_disposition=write_disposition) + + _arrow_query.validate_query("ARROW", job_config, compression_codec="ZSTD") + + check_dependencies.assert_called_once_with() + + +def test_query_and_wait_arrow_checks_dependencies_before_sending(check_dependencies): + check_dependencies.side_effect = exceptions.BigQueryStorageNotFoundError("missing") + client = mock.MagicMock(spec=Client) + + with pytest.raises(exceptions.BigQueryStorageNotFoundError): + _query_and_wait(client, query_results_format="ARROW") + + client._call_api.assert_not_called() + + +def test_check_dependencies_requires_pyarrow(): + error = exceptions.LegacyPyarrowError("pyarrow is not installed") + with mock.patch.object( + _versions_helpers.PYARROW_VERSIONS, "try_import", side_effect=error + ) as try_import: + with pytest.raises(exceptions.LegacyPyarrowError): + _check_dependencies() + + try_import.assert_called_once_with(raise_if_error=True) + + +def test_check_dependencies_requires_bigquery_storage(): + error = exceptions.BigQueryStorageNotFoundError("not installed") + with ( + mock.patch.object(_versions_helpers.PYARROW_VERSIONS, "try_import"), + mock.patch.object( + _versions_helpers.BQ_STORAGE_VERSIONS, "try_import", side_effect=error + ) as try_import, + ): + with pytest.raises(exceptions.BigQueryStorageNotFoundError): + _check_dependencies() + + try_import.assert_called_once_with(raise_if_error=True) + + +@pytest.mark.parametrize("supported", [True, False]) +def test_check_dependencies_requires_read_rows_arrow_options(supported): + versions_type = type(_versions_helpers.BQ_STORAGE_VERSIONS) + with ( + mock.patch.object(_versions_helpers.PYARROW_VERSIONS, "try_import"), + mock.patch.object(_versions_helpers.BQ_STORAGE_VERSIONS, "try_import"), + mock.patch.object( + versions_type, + "supports_arrow_serialization_options", + new_callable=mock.PropertyMock, + return_value=supported, + ), + mock.patch.object( + versions_type, + "installed_version", + new_callable=mock.PropertyMock, + return_value=packaging.version.Version("2.42.0"), + ), + ): + if supported: + _check_dependencies() + else: + with pytest.raises( + exceptions.LegacyBigQueryStorageError, + match=r"pip install --upgrade google-cloud-bigquery-storage.*" + r"\(version found: 2\.42\.0\)", + ): + _check_dependencies() + + +def test_query_and_wait_arrow_sends_default_serialization_options(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "schema": {"fields": []}, + } + + row_iterator = _query_and_wait(client, query_results_format="ARROW") + + data = client._call_api.call_args.kwargs["data"] + assert data["arrowSerializationOptions"] == _DEFAULT_QUERY_OPTIONS + assert "arrowSerializationOptions" not in data["formatOptions"] + assert row_iterator._arrow_serialization_options == _DEFAULT_QUERY_OPTIONS + + +def test_query_and_wait_without_arrow_sends_no_serialization_options(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "rows": [], + "schema": {"fields": []}, + } + + row_iterator = _query_and_wait(client) + + assert "arrowSerializationOptions" not in client._call_api.call_args.kwargs["data"] + assert row_iterator._arrow_serialization_options is None + + +def test_query_and_wait_arrow_waits_with_the_same_options(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": False, + } + + with ( + mock.patch.object(_job_helpers, "_to_query_job") as to_query_job, + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, + ): + _query_and_wait(client, query_results_format="ARROW", compression_codec="ZSTD") + + wait_or_cancel.assert_called_once_with( + to_query_job.return_value, + api_timeout=mock.ANY, + wait_timeout=mock.ANY, + retry=mock.ANY, + page_size=mock.ANY, + max_results=mock.ANY, + query_results_format="ARROW", + arrow_serialization_options={ + "bufferCompression": "ZSTD", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", + }, + callback=mock.ANY, + ) + + +def test_wait_or_cancel_sets_arrow_serialization_options(): + job = mock.MagicMock(dry_run=False) + job.result.return_value = RowIterator( + client=mock.MagicMock(), api_request=mock.MagicMock(), path=None, schema=() + ) + + rows = _job_helpers._wait_or_cancel( + job, + api_timeout=None, + wait_timeout=None, + retry=None, + page_size=None, + max_results=None, + query_results_format="ARROW", + arrow_serialization_options=_DEFAULT_QUERY_OPTIONS, + ) + + assert rows._query_results_format == "ARROW" + assert rows._arrow_serialization_options == _DEFAULT_QUERY_OPTIONS + + +def test_query_and_wait_arrow_jobs_insert_omits_arrow_fields(): + client = make_client() + client._call_api = mock.Mock( + return_value={ + "jobReference": {"projectId": "PROJECT", "jobId": "j", "location": "US"}, + "configuration": {"query": {"query": "SELECT 1"}}, + "status": {"state": "RUNNING"}, + } + ) + # jobs.query doesn't accept a destination table, so this uses jobs.insert. + job_config = QueryJobConfig(destination="PROJECT.dataset.table") + + with mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel: + client.query_and_wait( + "SELECT 1", job_config=job_config, query_results_format="ARROW" + ) + + (insert_call,) = client._call_api.call_args_list + assert insert_call.kwargs["path"] == "/projects/PROJECT/jobs" + query_config = insert_call.kwargs["data"]["configuration"]["query"] + assert "queryResultsFormat" not in query_config + assert "arrowSerializationOptions" not in query_config + wait_or_cancel.assert_called_once_with( + mock.ANY, + api_timeout=mock.ANY, + wait_timeout=mock.ANY, + retry=mock.ANY, + page_size=mock.ANY, + max_results=mock.ANY, + query_results_format="ARROW", + arrow_serialization_options=_DEFAULT_QUERY_OPTIONS, + explicit_destination=True, + callback=mock.ANY, + ) + assert wait_or_cancel.call_args.kwargs["page_size"] == 0 + + +def test_supported_by_jobs_query_includes_arrow_serialization_options(): + body = { + "query": "SELECT 1", + "queryResultsFormat": "ARROW", + "arrowSerializationOptions": _DEFAULT_QUERY_OPTIONS, + } + + assert _job_helpers._supported_by_jobs_query(body) + + +@pytest.mark.parametrize("codec", ["LZ4_FRAME", QueryResultsCompressionCodec.LZ4_FRAME]) +def test_serialization_options_accepts_str_and_enum(codec): + assert _arrow_query.serialization_options(codec) == { + "bufferCompression": "LZ4_FRAME", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", + } + + +def test_read_rows_options(): + assert _arrow_query._read_rows_options(None) == _DEFAULT_READ_OPTIONS + assert _arrow_query._read_rows_options( + { + "bufferCompression": "ZSTD", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_PICOS", + } + ) == { + "buffer_compression": "ZSTD", + "picos_timestamp_precision": "TIMESTAMP_PRECISION_PICOS", + } + + +def test_download_arrow_from_job_id_sends_query_options_on_read_rows(): + bqstorage_client = mock.MagicMock() + bqstorage_client.read_rows.return_value = [] + iterator = RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=(), + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + arrow_serialization_options={ + "bufferCompression": "LZ4_FRAME", + "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", + }, + ) + + with mock.patch("google.cloud.bigquery._arrow_query.pyarrow"): + list(iterator._download_arrow_from_job_id(bqstorage_client=bqstorage_client)) + + bqstorage_client.read_rows.assert_called_once_with( + "projects/p/locations/US/jobs/j/streams/_default", + offset=0, + timeout=None, + arrow_serialization_options={ + "buffer_compression": "LZ4_FRAME", + "picos_timestamp_precision": "TIMESTAMP_PRECISION_MICROS", + }, + ) + + +@pytest.mark.parametrize( + "timestamp_precision", + [TimestampPrecision.MICROSECOND, TimestampPrecision.PICOSECOND], +) +def test_query_and_wait_timestamp_precision_requires_arrow( + check_dependencies, timestamp_precision +): + client = mock.MagicMock(spec=Client) + + with pytest.raises(ValueError, match="timestamp_precision requires"): + _query_and_wait(client, timestamp_precision=timestamp_precision) + + client._call_api.assert_not_called() + check_dependencies.assert_not_called() + + +@pytest.mark.parametrize( + ("timestamp_precision", "expected"), + [ + (None, "TIMESTAMP_PRECISION_MICROS"), + (TimestampPrecision.MICROSECOND, "TIMESTAMP_PRECISION_MICROS"), + (TimestampPrecision.PICOSECOND, "TIMESTAMP_PRECISION_PICOS"), + ], +) +def test_client_query_and_wait_arrow_timestamp_precision(timestamp_precision, expected): + client = make_client() + client._call_api = mock.Mock( + return_value={ + "jobReference": {"projectId": "PROJECT", "jobId": "j", "location": "US"}, + "jobComplete": True, + "schema": {"fields": []}, + } + ) + + rows = client.query_and_wait( + "SELECT 1", + query_results_format="ARROW", + timestamp_precision=timestamp_precision, + ) + + data = client._call_api.call_args.kwargs["data"] + assert data["arrowSerializationOptions"]["picosTimestampPrecision"] == expected + assert rows._arrow_serialization_options["picosTimestampPrecision"] == expected + # formatOptions only shapes STRUCT_ENCODING rows, which ARROW doesn't return. + assert "timestampOutputFormat" not in data["formatOptions"] + + +def _arrow_query_response(**fields): + return { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "arrowSchema": {"serializedSchema": "c2NoZW1h"}, + **fields, + } + + +@pytest.mark.parametrize( + "response", + [ + _arrow_query_response(pageRowCount="10", totalRows="25"), + _arrow_query_response(jobComplete=False), + ], + ids=["some_rows_inline", "no_rows_inline"], +) +def test_query_and_wait_arrow_waits_without_fetching_rows(response): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = response + + with ( + mock.patch.object(_job_helpers, "_to_query_job") as to_query_job, + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, + ): + rows = _query_and_wait(client, query_results_format="ARROW", page_size=10) + + assert rows is wait_or_cancel.return_value + assert wait_or_cancel.call_args.args[0] is to_query_job.return_value + assert wait_or_cancel.call_args.kwargs["page_size"] == 0 + # page_size still sets the size of the first page from jobs.query. + assert client._call_api.call_args.kwargs["data"]["maxResults"] == 10 + + +def test_query_and_wait_arrow_all_rows_inline_ignores_page_token(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = _arrow_query_response( + pageRowCount="25", 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") + + wait_or_cancel.assert_not_called() + assert rows._first_page_response["pageRowCount"] == "25" + + +def test_query_and_wait_without_arrow_waits_with_page_size(): + client = mock.MagicMock(spec=Client) + client._call_api.return_value = { + "jobReference": {"projectId": "p", "jobId": "j", "location": "us"}, + "jobComplete": True, + "rows": [], + "schema": {"fields": []}, + "pageToken": "next", + } + + with ( + mock.patch.object(_job_helpers, "_to_query_job"), + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, + ): + _query_and_wait(client, page_size=10) + + 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)] + + +def test_download_arrow_uses_passed_client_and_leaves_it_open(): + iterator, bqstorage_client, _ = _limited_iterator([10, 15], 25, max_results=None) + + batches = list(iterator._download_arrow_from_job_id(bqstorage_client)) + + assert _values(batches) == list(range(1, 26)) + iterator.client._get_shared_bqstorage_client.assert_not_called() + bqstorage_client._transport.close.assert_not_called() + + +def test_download_arrow_needs_a_client_only_to_stream(): + iterator, _, _ = _limited_iterator([10, 15], 25, max_results=None) + iterator.client = None + batches = iterator._download_arrow_from_job_id() + + assert next(batches).num_rows == 10 + with pytest.raises(ValueError, match="RowIterator client is None"): + next(batches) + + +def test_to_arrow_with_arrow_format_reads_with_the_shared_client(): + pyarrow = pytest.importorskip("pyarrow") + client = mock.MagicMock() + iterator = RowIterator( + client=client, + api_request=mock.MagicMock(), + path=None, + schema=(), + 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: + table = iterator.to_arrow() + + assert table.num_rows == 2 + # to_arrow() doesn't create (and then close) a client of its own; the + # stream uses the BigQuery client's shared one. + client._ensure_bqstorage_client.assert_not_called() + assert to_arrow_iterable.call_args.kwargs["bqstorage_client"] is None + + +@pytest.mark.parametrize( + ("job_config", "expected"), + [ + (QueryJobConfig(destination="p.d.t"), True), + # jobs.query doesn't accept a priority either, so this also runs + # through jobs.insert, without a destination. + (QueryJobConfig(priority="INTERACTIVE"), False), + ], + ids=["destination", "no_destination"], +) +def test_query_and_wait_arrow_marks_explicit_destination(job_config, expected): + client = mock.MagicMock(spec=Client) + + with ( + mock.patch.object(_job_helpers, "query_jobs_insert") as query_jobs_insert, + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, + ): + _query_and_wait(client, job_config=job_config, query_results_format="ARROW") + + wait_or_cancel.assert_called_once_with( + query_jobs_insert.return_value, + api_timeout=mock.ANY, + wait_timeout=mock.ANY, + retry=mock.ANY, + page_size=mock.ANY, + max_results=mock.ANY, + query_results_format="ARROW", + arrow_serialization_options=mock.ANY, + explicit_destination=expected, + callback=mock.ANY, + ) + + +def test_wait_or_cancel_sets_explicit_destination(): + job = mock.MagicMock(dry_run=False) + job.result.return_value = RowIterator( + client=mock.MagicMock(), api_request=mock.MagicMock(), path=None, schema=() + ) + assert not job.result.return_value._explicit_destination + + rows = _job_helpers._wait_or_cancel( + job, + api_timeout=None, + wait_timeout=None, + retry=None, + page_size=0, + max_results=None, + explicit_destination=True, + ) + + assert rows._explicit_destination + + +def _destination_iterator(row_counts, serialized_schema=None): + """A RowIterator for a query with an explicit destination table, plus its + mock BigQuery Storage client. + """ + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + options = bigquery_storage.types.ArrowSerializationOptions + if "picos_timestamp_precision" not in options.meta.fields: + # query_and_wait() rejects ARROW with this release before it runs. + pytest.skip("google-cloud-bigquery-storage has no timestamp precision") + schema, batches = _arrow_pages(*row_counts) + reader = mock.MagicMock() + reader.__iter__.return_value = iter([_stream_response(batch) for batch in batches]) + stream = mock.Mock() + stream.name = "projects/p/locations/us/sessions/s/streams/0" + session = mock.Mock(streams=[stream] if row_counts else []) + session.arrow_schema.serialized_schema = ( + schema if serialized_schema is None else serialized_schema + ) + bqstorage_client = mock.MagicMock() + bqstorage_client.create_read_session.return_value = session + bqstorage_client.read_rows.return_value = reader + iterator = RowIterator( + client=mock.MagicMock(project="billing-project"), + api_request=mock.MagicMock(), + path=None, + schema=(), + table=TableReference.from_string("p.d.t"), + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + ) + iterator._explicit_destination = True + return iterator, bqstorage_client + + +def test_download_arrow_explicit_destination_reads_a_read_session(): + iterator, bqstorage_client = _destination_iterator([10, 15]) + from google.cloud import bigquery_storage + + batches = list(iterator._download_arrow_from_job_id(bqstorage_client)) + + assert _values(batches) == list(range(1, 26)) + bqstorage_client.create_read_session.assert_called_once_with( + parent="projects/billing-project", + read_session=bigquery_storage.types.stream.ReadSession( + table="projects/p/datasets/d/tables/t", + data_format=bigquery_storage.types.stream.DataFormat.ARROW, + read_options=bigquery_storage.types.stream.ReadSession.TableReadOptions( + arrow_serialization_options=_DEFAULT_READ_OPTIONS + ), + ), + max_stream_count=1, + timeout=None, + ) + # The read session carries the serialization options, so ReadRows doesn't. + bqstorage_client.read_rows.assert_called_once_with( + "projects/p/locations/us/sessions/s/streams/0", offset=0, timeout=None + ) + + +@pytest.mark.parametrize( + ("serialized_schema", "expected_names"), + [(None, ["x"]), (b"", [])], + ids=["with_schema", "without_schema"], +) +def test_download_arrow_explicit_destination_without_rows( + serialized_schema, expected_names +): + iterator, bqstorage_client = _destination_iterator( + [], serialized_schema=serialized_schema + ) + + (batch,) = iterator._download_arrow_from_job_id(bqstorage_client) + + assert batch.num_rows == 0 + assert batch.schema.names == expected_names + bqstorage_client.read_rows.assert_not_called() + + +def _zero_row_iterator(): + """A RowIterator whose jobs.query response has a schema and no rows.""" + schema, _ = _arrow_pages() + return RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=(), + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + first_page_response={ + "jobComplete": True, + "totalRows": "0", + "arrowSchema": {"serializedSchema": base64.b64encode(schema).decode()}, + # A response without rows carries no serialized batch. + "arrowRecordBatch": {}, + }, + ) + + +def test_download_arrow_zero_rows_inline_keep_the_schema(): + iterator = _zero_row_iterator() + + (batch,) = iterator._download_arrow_from_job_id() + + assert batch.num_rows == 0 + assert batch.schema.names == ["x"] + iterator.client._get_shared_bqstorage_client.assert_not_called() + + +def test_to_arrow_zero_rows_keep_the_schema(): + table = _zero_row_iterator().to_arrow() + + assert table.num_rows == 0 + assert table.schema.names == ["x"] + + +def test_download_arrow_zero_rows_streamed_use_the_query_results_schema(): + pyarrow = pytest.importorskip("pyarrow") + from google.cloud.bigquery.schema import SchemaField + + bqstorage_client = mock.MagicMock() + # The default job stream sends no responses for zero rows. + bqstorage_client.read_rows.return_value = [] + iterator = RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=[SchemaField("x", "INTEGER")], + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + first_page_response={"jobComplete": False}, + ) + + (batch,) = iterator._download_arrow_from_job_id(bqstorage_client) + + assert batch.num_rows == 0 + assert batch.schema.names == ["x"] + assert batch.schema.field("x").type == pyarrow.int64() + + +def test_download_arrow_rows_yield_no_extra_empty_batch(): + schema, batches = _arrow_pages(3) + iterator = RowIterator( + client=mock.MagicMock(), + api_request=mock.MagicMock(), + path=None, + schema=(), + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + first_page_response={ + "jobComplete": True, + "totalRows": "3", + "arrowSchema": {"serializedSchema": base64.b64encode(schema).decode()}, + "arrowRecordBatch": { + "serializedRecordBatch": base64.b64encode(batches[0]).decode() + }, + }, + ) + + (batch,) = iterator._download_arrow_from_job_id() + + assert batch.num_rows == 3 + + if __name__ == "__main__": unittest.main()