diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index 25738a2e818d..8fa08700a208 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -19,6 +19,10 @@ import base64 from typing import Iterator, Optional, 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: @@ -34,6 +38,52 @@ ) +def validate_query( + query_results_format: Optional[str], + *, + compression_codec: Optional[str] = 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()``. + compression_codec: The ``compression_codec`` passed to + ``query_and_wait()``. + + Raises: + ValueError: + If ``compression_codec`` is set without ``ARROW``. + 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.") + return + + # 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 download_arrow_from_job( row_iterator: "RowIterator", bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None, 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..2168de25de54 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 @@ -501,7 +502,16 @@ def query_and_wait( If ``job_config`` is not an instance of :class:`~google.cloud.bigquery.job.QueryJobConfig` class. + ValueError: + If ``compression_codec`` is set without ``ARROW``. + 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, compression_codec=compression_codec + ) + request_body = _to_query_request( query=query, job_config=job_config, location=location, timeout=api_timeout ) 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..214982dc2a23 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -3746,6 +3746,11 @@ def query_and_wait( If ``job_config`` is not an instance of :class:`~google.cloud.bigquery.job.QueryJobConfig` class. + ValueError: + If ``compression_codec`` is set without ``ARROW``. + 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, 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_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index 1e22381419c7..dd11f8f6f686 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 @@ -16,12 +16,30 @@ import unittest from unittest import mock +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.enums import WriteDisposition +from google.cloud.bigquery.job import QueryJobConfig from google.cloud.bigquery.table import RowIterator, _EmptyRowIterator +from .helpers import make_client + +# 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): def test_supported_by_jobs_query_includes_query_results_format(self): @@ -569,5 +587,121 @@ 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_allows_write_append(check_dependencies): + client = make_client( + default_query_job_config=QueryJobConfig( + destination="p.d.t", write_disposition=WriteDisposition.WRITE_APPEND + ) + ) + client._call_api = mock.Mock(side_effect=RuntimeError("sent")) + + with pytest.raises(RuntimeError, match="sent"): + client.query_and_wait("SELECT 1", query_results_format="ARROW") + + check_dependencies.assert_called_once_with() + client._call_api.assert_called_once() + + +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() + + if __name__ == "__main__": unittest.main()