feat(bigquery): skip job creation for ARROW query_and_wait results

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()
 

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]):

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 = {