Repository navigation
feat(bigquery): skip job creation for ARROW query_and_wait results - #18799
Conversation
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.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the query_and_wait method to automatically set the jobCreationMode to JOB_CREATION_OPTIONAL when the query results format is set to ARROW and no default mode is specified on the client. It also updates the corresponding docstrings and adds unit tests to verify this behavior. Feedback is provided regarding the handling of the query_results_format Enum parameter, advising that type annotations and input handling for raw string representations should be consistent across all related API methods rather than handled selectively.
| 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: |
There was a problem hiding this comment.
When supporting raw string representations for an Enum parameter like query_results_format, we must ensure that type annotations and input handling are consistent across all related API methods. We should not selectively support or normalize raw string inputs in only a subset of methods. If we want to support string inputs (including case-insensitive comparisons), this handling should be applied uniformly across all related methods rather than as a localized change here.
References
- When a parameter accepts an Enum, ensure its type annotations and input handling are consistent across all related API methods. Do not selectively support raw string representations in only a subset of these methods.
There was a problem hiding this comment.
This doesn't add any string handling. It's the same == enums.QueryResultsFormat.ARROW check query_and_wait() already uses a few lines up, and QueryResultsFormat is a str enum, so it behaves the same as the existing check.
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.
Verify that query_and_wait() sets RowIterator.job_id to None and populates query_id when the jobs.query response omits jobReference under JOB_CREATION_OPTIONAL.
ohmayr
added
the
kokoro:force-run
yoshi-kokoro
removed
the
kokoro:force-run
Before,
query_and_wait(query_results_format="ARROW")always created a job unless the caller setclient.default_job_creation_mode, so small Arrow queries paid about 200 ms for a job they didn't need. Arrow callers want low latency, so ARROW queries now default toJOB_CREATION_OPTIONAL, which is whatcreateQueryArrowStream()in the Node.js client already does.This changes behavior for existing ARROW callers (Beta since 3.44.0): results that fit in the first response now come back with
RowIterator.job_idset toNone.locationandquery_idare still set. To always get a job, setdefault_job_creation_mode="JOB_CREATION_REQUIRED"on the client or useclient.query().query_and_wait()sendsjobCreationMode=JOB_CREATION_OPTIONALwhenquery_results_formatisARROWand the client doesn't setdefault_job_creation_mode. A client-level mode still wins.jobs.insert(for example, with a destination table) also still create a job.ARROWare unchanged.query_results_formatanddefault_job_creation_modedocstrings describe the new default and how to opt out.jobs.insertfallback. Ran the unit suite andtests/system/test_arrow.pyagainst BigQuery. A live single-page ARROW query returnedjob_id=None, while a multi-page ARROW query, a non-ARROW query and an ARROW query withJOB_CREATION_REQUIREDon the client each returned a job.