From 74a74d8f1a33accd76b4d3a777ef466711fdbeb4 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH 1/3] feat(bigquery): validate Arrow queries before sending them query_and_wait() with query_results_format=ARROW only found a missing pyarrow or google-cloud-bigquery-storage when results were read, after the query had already run. A query could work while its results fit in the jobs.query response and then fail once they needed a stream. It also accepted WRITE_APPEND, whose results can't be read on their own, and compression_codec without ARROW, where it does nothing. Check all three before the request is built. compression_codec without ARROW and an ARROW query with WRITE_APPEND (including through default_query_job_config) raise ValueError. An ARROW query checks pyarrow and google-cloud-bigquery-storage with the existing version helpers, and also checks that BigQueryReadClient.read_rows() accepts arrow_serialization_options. That check looks for ReadRowsStream.cancel(), which shipped in the same change, rather than a version number, so it's correct for the unreleased in-repo build that the system tests install until the next storage release. --- .../google/cloud/bigquery/_arrow_query.py | 67 ++++++++ .../google/cloud/bigquery/_job_helpers.py | 11 ++ .../cloud/bigquery/_versions_helpers.py | 15 ++ .../google/cloud/bigquery/client.py | 7 + packages/google-cloud-bigquery/noxfile.py | 10 ++ .../tests/system/test_arrow.py | 27 +++ .../tests/unit/test__versions_helpers.py | 25 +++ .../unit/test_query_results_format_arrow.py | 157 ++++++++++++++++++ 8 files changed, 319 insertions(+) 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..618ff031e9b6 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: @@ -26,6 +30,7 @@ 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 = ( @@ -34,6 +39,68 @@ ) +def validate_query( + query_results_format: Optional[str], + job_config: Optional["job.QueryJobConfig"], + *, + 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()``. + job_config: The job config, already merged with + ``Client.default_query_job_config``. + compression_codec: The ``compression_codec`` passed to + ``query_and_wait()``. + + Raises: + ValueError: + If ``compression_codec`` 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.") + 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 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..a952eff60f8c 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,17 @@ 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``, 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 + ) + 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..1322a7538a40 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -3746,6 +3746,13 @@ 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``, 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, diff --git a/packages/google-cloud-bigquery/noxfile.py b/packages/google-cloud-bigquery/noxfile.py index af201dfb8a05..2cd266371c7c 100644 --- a/packages/google-cloud-bigquery/noxfile.py +++ b/packages/google-cloud-bigquery/noxfile.py @@ -288,6 +288,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..f9d4ade03f18 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -274,3 +274,30 @@ 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] 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..fda8477ad3c8 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,144 @@ 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() + + if __name__ == "__main__": unittest.main() From 50a389d553a2dd4aeebbec9b2b306fc973e3fd7e Mon Sep 17 00:00:00 2001 From: ohmayr Date: Fri, 16 Oct 2026 00:02:32 +0000 Subject: [PATCH 2/3] fix(bigquery): allow WRITE_APPEND for Arrow queries BigQuery returns the whole destination table for WRITE_APPEND in every results format, so ARROW results already match STRUCT_ENCODING. The check also rejected WRITE_APPEND without a destination table, which works. Remove the check and validate_query's job_config parameter, which only the check used. --- .../google/cloud/bigquery/_arrow_query.py | 19 +----------- .../google/cloud/bigquery/_job_helpers.py | 5 ++- .../google/cloud/bigquery/client.py | 4 +-- .../tests/system/test_arrow.py | 27 ---------------- .../unit/test_query_results_format_arrow.py | 31 +++---------------- 5 files changed, 8 insertions(+), 78 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index 618ff031e9b6..8fa08700a208 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -30,7 +30,6 @@ 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 = ( @@ -41,7 +40,6 @@ def validate_query( query_results_format: Optional[str], - job_config: Optional["job.QueryJobConfig"], *, compression_codec: Optional[str] = None, ) -> None: @@ -50,15 +48,12 @@ def validate_query( 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()``. Raises: ValueError: - If ``compression_codec`` is set without ``ARROW``, or if an - ``ARROW`` query uses ``WRITE_APPEND``. + 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. @@ -68,18 +63,6 @@ def validate_query( raise ValueError("compression_codec 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() 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 a952eff60f8c..2168de25de54 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -503,14 +503,13 @@ def query_and_wait( :class:`~google.cloud.bigquery.job.QueryJobConfig` class. ValueError: - If ``compression_codec`` is set without ``ARROW``, or if an - ``ARROW`` query uses ``WRITE_APPEND``. + 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, job_config, compression_codec=compression_codec + query_results_format, compression_codec=compression_codec ) request_body = _to_query_request( diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index 1322a7538a40..214982dc2a23 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -3747,9 +3747,7 @@ def query_and_wait( :class:`~google.cloud.bigquery.job.QueryJobConfig` class. ValueError: - If ``compression_codec`` is set without ``ARROW``, or if an - ``ARROW`` query uses ``WRITE_APPEND``, including through - ``default_query_job_config``. + 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. diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index f9d4ade03f18..9db96503b00b 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -274,30 +274,3 @@ 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] 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 fda8477ad3c8..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 @@ -623,42 +623,19 @@ def test_query_and_wait_without_arrow_skips_dependency_check(check_dependencies) 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(): +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() + client._call_api = mock.Mock(side_effect=RuntimeError("sent")) - with pytest.raises(ValueError, match="WRITE_APPEND"): + with pytest.raises(RuntimeError, match="sent"): 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() + client._call_api.assert_called_once() def test_query_and_wait_arrow_checks_dependencies_before_sending(check_dependencies): From 6ff1885dc43cdbefe5d765d94c4fbd7e4938bca8 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Fri, 16 Oct 2026 01:15:20 +0000 Subject: [PATCH 3/3] ci(bigquery): install released google-cloud-bigquery-storage in system tests google-cloud-bigquery-storage 2.43.0 has read_rows() with arrow_serialization_options and ReadRowsStream.cancel(), which ARROW query results need. The system session no longer has to install the in-repo copy. --- packages/google-cloud-bigquery/noxfile.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/packages/google-cloud-bigquery/noxfile.py b/packages/google-cloud-bigquery/noxfile.py index 2cd266371c7c..af201dfb8a05 100644 --- a/packages/google-cloud-bigquery/noxfile.py +++ b/packages/google-cloud-bigquery/noxfile.py @@ -288,16 +288,6 @@ 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,