feat(bigquery): support queryResultsFormat and compressionCodec in query_and_wait - #18027
feat(bigquery): support queryResultsFormat and compressionCodec in query_and_wait#18027alextolpin wants to merge 5 commits into
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces support for the Arrow query results format and compression codecs in query_and_wait. It adds the query_results_format and compression_codec parameters, prevents standard iteration on RowIterator when the format is Arrow, and implements _download_arrow_from_job_id to retrieve Arrow results via the BigQuery Storage Read API. The review feedback highlights several important improvements: adding formatOptions to the _supported_by_jobs_query allowlist to prevent unnecessary fallbacks to jobs.insert, raising an error instead of silently skipping record batches when the schema is missing, safely retrieving totalRows to avoid potential KeyErrors, and validating key identifiers before constructing the stream name to prevent cryptic API errors.
| "jobTimeoutMs", | ||
| "reservation", | ||
| "maxSlots", | ||
| "queryResultsFormat", | ||
| } |
There was a problem hiding this comment.
The formatOptions key is not included in the keys_allowlist for _supported_by_jobs_query. When compression_codec is specified, formatOptions is added to the request body. Since it is missing from the allowlist, any query utilizing a compression codec will unnecessarily fallback to the slower jobs.insert path instead of using the optimized jobs.query API. Adding formatOptions to the allowlist ensures the fast path is preserved.
| "jobTimeoutMs", | |
| "reservation", | |
| "maxSlots", | |
| "queryResultsFormat", | |
| } | |
| "jobTimeoutMs", | |
| "reservation", | |
| "maxSlots", | |
| "queryResultsFormat", | |
| "formatOptions", | |
| } |
References
- For performance-critical code paths executed on every request, validate and benchmark any proposed readability simplifications to ensure they do not degrade performance or eliminate fast-path optimizations.
| if ( | ||
| response.arrow_record_batch | ||
| and response.arrow_record_batch.serialized_record_batch | ||
| and pa_schema is not None | ||
| ): | ||
| batch = pyarrow.ipc.read_record_batch( | ||
| pyarrow.py_buffer(response.arrow_record_batch.serialized_record_batch), | ||
| pa_schema, | ||
| ) | ||
| yield batch |
There was a problem hiding this comment.
If response.arrow_record_batch is present but pa_schema is None, the record batch will be silently skipped, leading to silent data loss and incomplete results. Raising a ValueError when the schema is missing is much safer and prevents silent failures.
| if ( | |
| response.arrow_record_batch | |
| and response.arrow_record_batch.serialized_record_batch | |
| and pa_schema is not None | |
| ): | |
| batch = pyarrow.ipc.read_record_batch( | |
| pyarrow.py_buffer(response.arrow_record_batch.serialized_record_batch), | |
| pa_schema, | |
| ) | |
| yield batch | |
| if ( | |
| response.arrow_record_batch | |
| and response.arrow_record_batch.serialized_record_batch | |
| ): | |
| if pa_schema is None: | |
| raise ValueError("Arrow schema is missing; cannot deserialize record batch.") | |
| batch = pyarrow.ipc.read_record_batch( | |
| pyarrow.py_buffer(response.arrow_record_batch.serialized_record_batch), | |
| pa_schema, | |
| ) | |
| yield batch |
References
- For data streams assumed to be of a single format, if an unexpected state or format change is detected mid-stream, it is preferable to raise an exception to make the unexpected state explicit.
| job_complete = bool(first_page.get("jobComplete", False)) | ||
| if job_complete: | ||
| total_rows = int(first_page["totalRows"]) |
There was a problem hiding this comment.
If the query is a DDL/DML statement (which does not return rows) or if totalRows is missing from the response for any other reason, accessing first_page["totalRows"] directly will raise a KeyError. Using .get("totalRows", 0) is safer and prevents potential crashes.
| job_complete = bool(first_page.get("jobComplete", False)) | |
| if job_complete: | |
| total_rows = int(first_page["totalRows"]) | |
| job_complete = bool(first_page.get("jobComplete", False)) | |
| if job_complete: | |
| total_rows = int(first_page.get("totalRows", 0)) |
| project = self._project or (self.client.project if self.client else None) | ||
| location = self._location or (self.client.location if self.client else None) | ||
| stream_name = ( | ||
| f"projects/{project}/locations/{location}/jobs/{self._job_id}/streams/_default" | ||
| ) |
There was a problem hiding this comment.
If project, location, or self._job_id is None, the constructed stream_name will contain literal "None" values (e.g., projects/None/locations/None/...), leading to cryptic API errors. Adding explicit validation checks ensures a clear, local error is raised instead.
project = self._project or (self.client.project if self.client else None)
location = self._location or (self.client.location if self.client else None)
if not project:
raise ValueError("Project is required to read Arrow results.")
if not location:
raise ValueError("Location is required to read Arrow results.")
if not self._job_id:
raise ValueError("Job ID is required to read Arrow results.")
stream_name = (
f"projects/{project}/locations/{location}/jobs/{self._job_id}/streams/_default"
)References
- When a function receives parameters of an unsupported type, it should raise an error instead of silently returning empty values to ensure fail-fast behavior.
Summary of Changes
Adds support for fetching query results in Apache Arrow format directly via
query_and_wait()usingqueryResultsFormat="ARROW"and optional buffer compression (e.g.,compression_codec="LZ4_FRAME").query_and_wait&_job_helpersEnhancements:query_results_formatandcompression_codecparameters (with[Beta]docstring annotations) toclient.query_and_wait(),client._query_and_wait_bigframes(), and_job_helpers.query_and_wait().queryResultsFormatin_job_helpers.keys_allowlistand populatedformatOptions.arrowSerializationOptions.bufferCompressioninjobs.queryREST API request payloads._wait_or_cancel()to accept and preservequery_results_formaton returnedRowIteratorinstances.Arrow Serialization & Direct Job Stream Reading:
RowIterator._download_arrow_from_job_id()to stream Arrow record batches directly fromprojects/{project}/locations/{location}/jobs/{job_id}/streams/_defaultvia the BigQuery Storage Read API.arrowSchemaandarrowRecordBatchfrom the initialjobs.queryREST response (_first_page_response), calculate the starting rowoffset, and resumeread_rows(stream_name, offset=offset).read_rows()or initializingBigQueryReadClientifjobComplete = Trueand all rows were returned within the first page response.Safety & Enforcement:
pages,__iter__, and__next__onRowIteratorand_EmptyRowIteratorto raise a descriptiveValueErrorif non-Arrow iteration is attempted whenqueryResultsFormat="ARROW".Testing:
tests/unit/test_query_results_format_arrow.py(16 passing tests) covering request body formatting, parameter propagation, base64 payload decoding, offset calculation, stream URI construction, and Storage client skipping when all rows are present in the first page.Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: