From a7db8f7021254c07de243734820be2cc57bcdb33 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 21 Oct 2026 18:57:46 +0000 Subject: [PATCH 1/2] feat(bigquery): skip job creation for ARROW query_and_wait results Callers who pass query_results_format="ARROW" want low latency, so query_and_wait() now sends jobCreationMode=JOB_CREATION_OPTIONAL for them unless the client sets default_job_creation_mode. Single-page results skip creating a job (about 200 ms faster), so RowIterator.job_id is None for them. Multi-page results still create a job. Callers who need a job can set default_job_creation_mode to JOB_CREATION_REQUIRED or use client.query(). Queries without ARROW are unchanged. This matches createQueryArrowStream() in the Node.js client. --- .../google/cloud/bigquery/_job_helpers.py | 4 + .../google/cloud/bigquery/client.py | 11 ++- .../tests/unit/test__job_helpers.py | 79 +++++++++++++++++++ 3 files changed, 93 insertions(+), 1 deletion(-) 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 c06579797a48..83aef7f91cb6 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -572,6 +572,10 @@ def query_and_wait( request_body["maxResults"] = page_size or max_results if client.default_job_creation_mode: request_body["jobCreationMode"] = client.default_job_creation_mode + elif query_results_format == enums.QueryResultsFormat.ARROW: + # ARROW callers want low latency, so let BigQuery skip creating a + # job when it can return every row in the first response. + request_body["jobCreationMode"] = enums.JobCreationMode.JOB_CREATION_OPTIONAL query_sent_factory = QuerySentEventFactory() diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index 9d0e9b1540ef..b35d1997c0c5 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -240,7 +240,9 @@ class Client(ClientWithProject): default_job_creation_mode (Optional[str]): Sets the default job creation mode used by query methods such as query_and_wait(). For lightweight queries, JOB_CREATION_OPTIONAL is - generally recommended. + generally recommended. If unset, query_and_wait() uses + JOB_CREATION_OPTIONAL for ``ARROW`` results and creates a job + otherwise. Raises: google.auth.exceptions.DefaultCredentialsError: @@ -3734,6 +3736,13 @@ def query_and_wait( The maximum total number of rows from this request. query_results_format (Optional[Union[str, google.cloud.bigquery.enums.QueryResultsFormat]]): [Beta] The format for query results (e.g. "ARROW" or :class:`~google.cloud.bigquery.enums.QueryResultsFormat.ARROW`). + + With ``ARROW``, the query uses ``JOB_CREATION_OPTIONAL`` + unless :attr:`default_job_creation_mode` is set. Results that + fit in one page then have no job, so the returned iterator's + ``job_id`` is ``None``. To always create a job, set + ``default_job_creation_mode`` to ``JOB_CREATION_REQUIRED`` or + use :meth:`query`. 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]): diff --git a/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py b/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py index cf38328dd9c1..fe474d65cbc0 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py @@ -537,6 +537,85 @@ def test_query_and_wait_sets_job_creation_mode(): ) +@pytest.mark.parametrize( + ("default_job_creation_mode", "query_results_format", "expected"), + [ + # ARROW skips creating a job by default. + (None, "ARROW", "JOB_CREATION_OPTIONAL"), + # A mode set on the client still wins. + ("JOB_CREATION_REQUIRED", "ARROW", "JOB_CREATION_REQUIRED"), + # Other formats keep sending no mode. + (None, None, None), + ], +) +def test_query_and_wait_job_creation_mode_for_arrow( + default_job_creation_mode, query_results_format, expected +): + client = mock.create_autospec(Client) + client.default_job_creation_mode = default_job_creation_mode + client._call_api.return_value = { + "jobReference": { + "projectId": "response-project", + "jobId": "abc", + "location": "response-location", + }, + "jobComplete": True, + } + _job_helpers.query_and_wait( + client, + query="SELECT 1", + location="request-location", + project="request-project", + job_config=None, + retry=None, + job_retry=None, + page_size=None, + max_results=None, + query_results_format=query_results_format, + ) + + request_body = client._call_api.call_args.kwargs["data"] + assert request_body.get("jobCreationMode") == expected + + +def test_query_and_wait_arrow_with_jobs_insert_creates_a_job(): + # jobs.query doesn't support destination, so this goes through jobs.insert, + # which always creates a job. + client = mock.create_autospec(Client) + client.default_job_creation_mode = None + client._call_api.return_value = { + "jobReference": { + "projectId": "response-project", + "jobId": "abc", + "location": "response-location", + }, + "query": { + "query": "SELECT 1", + }, + "status": {"state": "DONE"}, + "jobComplete": True, + } + job_config = job_query.QueryJobConfig( + destination="dest-project.dest_dset.dest_table", + ) + _job_helpers.query_and_wait( + client, + query="SELECT 1", + location="request-location", + project="request-project", + job_config=job_config, + retry=None, + job_retry=None, + page_size=None, + max_results=None, + query_results_format="ARROW", + ) + + first_call = client._call_api.call_args_list[0] + assert first_call.kwargs["path"] == "/projects/request-project/jobs" + assert "jobCreationMode" not in first_call.kwargs["data"] + + def test_query_and_wait_sets_location(): client = mock.create_autospec(Client) client._call_api.return_value = { From 83ee1f84d7ff76d297fc34d269dfff20bf7b807a Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 21 Oct 2026 20:52:47 +0000 Subject: [PATCH 2/2] test(bigquery): stub the ARROW dependency check in job creation tests The unit_noextras session runs without google-cloud-bigquery-storage, so query_and_wait() raised before building the request. Patch _arrow_query._check_dependencies like the other ARROW tests do, since these tests only check the request. --- .../tests/unit/test__job_helpers.py | 55 +++++++++++-------- 1 file changed, 31 insertions(+), 24 deletions(-) diff --git a/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py b/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py index fe474d65cbc0..242994e4c6be 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py @@ -19,6 +19,7 @@ from google.api_core import retry as retries import pytest +from google.cloud.bigquery import _arrow_query from google.cloud.bigquery import _job_helpers from google.cloud.bigquery import enums from google.cloud.bigquery import retry @@ -561,18 +562,21 @@ def test_query_and_wait_job_creation_mode_for_arrow( }, "jobComplete": True, } - _job_helpers.query_and_wait( - client, - query="SELECT 1", - location="request-location", - project="request-project", - job_config=None, - retry=None, - job_retry=None, - page_size=None, - max_results=None, - query_results_format=query_results_format, - ) + # These tests check the request, not whether pyarrow and + # google-cloud-bigquery-storage are installed. + with mock.patch.object(_arrow_query, "_check_dependencies"): + _job_helpers.query_and_wait( + client, + query="SELECT 1", + location="request-location", + project="request-project", + job_config=None, + retry=None, + job_retry=None, + page_size=None, + max_results=None, + query_results_format=query_results_format, + ) request_body = client._call_api.call_args.kwargs["data"] assert request_body.get("jobCreationMode") == expected @@ -598,18 +602,21 @@ def test_query_and_wait_arrow_with_jobs_insert_creates_a_job(): job_config = job_query.QueryJobConfig( destination="dest-project.dest_dset.dest_table", ) - _job_helpers.query_and_wait( - client, - query="SELECT 1", - location="request-location", - project="request-project", - job_config=job_config, - retry=None, - job_retry=None, - page_size=None, - max_results=None, - query_results_format="ARROW", - ) + # This test checks the request, not whether pyarrow and + # google-cloud-bigquery-storage are installed. + with mock.patch.object(_arrow_query, "_check_dependencies"): + _job_helpers.query_and_wait( + client, + query="SELECT 1", + location="request-location", + project="request-project", + job_config=job_config, + retry=None, + job_retry=None, + page_size=None, + max_results=None, + query_results_format="ARROW", + ) first_call = client._call_api.call_args_list[0] assert first_call.kwargs["path"] == "/projects/request-project/jobs"