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..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 @@ -537,6 +538,91 @@ 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, + } + # 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 + + +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", + ) + # 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" + 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 = {