feat(bigquery): support queryResultsFormat and compressionCodec in query_and_wait

[alextolpin]

Summary of Changes

Adds support for fetching query results in Apache Arrow format directly via query_and_wait() using queryResultsFormat="ARROW" and optional buffer compression (e.g., compression_codec="LZ4_FRAME").

1. query_and_wait & _job_helpers Enhancements:

   • Added query_results_format and compression_codec parameters (with [Beta] docstring annotations) to client.query_and_wait(), client._query_and_wait_bigframes(), and _job_helpers.query_and_wait().
   • Included queryResultsFormat in _job_helpers.keys_allowlist and populated formatOptions.arrowSerializationOptions.bufferCompression in jobs.query REST API request payloads.
   • Refactored _wait_or_cancel() to accept and preserve query_results_format on returned RowIterator instances.

2. Arrow Serialization & Direct Job Stream Reading:

   • Added RowIterator._download_arrow_from_job_id() to stream Arrow record batches directly from projects/{project}/locations/{location}/jobs/{job_id}/streams/_default via the BigQuery Storage Read API.
   • Added logic to decode base64 inline arrowSchema and arrowRecordBatch from the initial jobs.query REST response (_first_page_response), calculate the starting row offset, and resume read_rows(stream_name, offset=offset).
   • Added an optimization to skip calling read_rows() or initializing BigQueryReadClient if jobComplete = True and all rows were returned within the first page response.

3. Safety & Enforcement:

   • Overrode pages, __iter__, and __next__ on RowIterator and _EmptyRowIterator to raise a descriptive ValueError if non-Arrow iteration is attempted when queryResultsFormat="ARROW".

4. Testing:

   • Added comprehensive unit test suite in 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:

• Make sure to open an issue as a bug/issue before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea
• Ensure the tests and linter pass
• Code coverage does not decrease (if any source code was changed)
• Appropriate docs were updated (if necessary)

[gemini-code-assist]

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.

[gemini-code-assist]

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.

Suggested change

	"jobTimeoutMs",
	"reservation",
	"maxSlots",
	"queryResultsFormat",
	}
	"jobTimeoutMs",
	"reservation",
	"maxSlots",
	"queryResultsFormat",
	"formatOptions",
	}

References

1. 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.

[gemini-code-assist]

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.

Suggested change

	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

1. 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.

[gemini-code-assist]

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.

Suggested change

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

[gemini-code-assist]

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

1. 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.