feat(bigquery): support queryResultsFormat and compressionCodec in query_and_wait (#18027)

### 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:
- [x] Make sure to open an issue as a
[bug/issue](https://github.com/googleapis/google-cloud-python/issues)
before writing your code! That way we can discuss the change, evaluate
designs, and agree on the general idea
- [x] Ensure the tests and linter pass
- [x] Code coverage does not decrease (if any source code was changed)
- [x] Appropriate docs were updated (if necessary)

Author: alextolpin

packages/google-cloud-bigquery/google/cloud/bigquery/__init__.py

@@ -59,6 +59,8 @@
 from google.cloud.bigquery.external_config import HivePartitioningOptions
 from google.cloud.bigquery.format_options import AvroOptions
 from google.cloud.bigquery.format_options import ParquetOptions
+from google.cloud.bigquery.enums import QueryResultsCompressionCodec
+from google.cloud.bigquery.enums import QueryResultsFormat
 from google.cloud.bigquery.job.base import SessionInfo
 from google.cloud.bigquery.job import Compression
 from google.cloud.bigquery.job import CopyJob
@@ -221,6 +223,8 @@
     "KeyResultStatementKind",
     "OperationType",
     "QueryPriority",
+    "QueryResultsCompressionCodec",
+    "QueryResultsFormat",
     "RoutineType",
     "SchemaUpdateOption",
     "SourceFormat",

packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py

@@ -430,6 +430,8 @@ def query_and_wait(
     job_retry: Optional[retries.Retry],
     page_size: Optional[int] = None,
     max_results: Optional[int] = None,
+    query_results_format: Optional[str] = None,
+    compression_codec: Optional[str] = None,
     callback: Callable = lambda _: None,
 ) -> table.RowIterator:
     """Run the query, wait for it to finish, and return the results.
@@ -475,6 +477,10 @@ def query_and_wait(
             request. Non-positive values are ignored.
         max_results (Optional[int]):
             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`).
+        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`).
         callback (Callable):
             A callback function used by bigframes to report query progress.
 
@@ -499,6 +505,13 @@ def query_and_wait(
     request_body = _to_query_request(
         query=query, job_config=job_config, location=location, timeout=api_timeout
     )
+    if query_results_format is not None:
+        request_body["queryResultsFormat"] = query_results_format
+    if compression_codec is not None:
+        request_body.setdefault("formatOptions", {})
+        request_body["formatOptions"]["arrowSerializationOptions"] = {
+            "bufferCompression": compression_codec
+        }
 
     # Some API parameters aren't supported by the jobs.query API. In these
     # cases, fallback to a jobs.insert call.
@@ -522,6 +535,7 @@ def query_and_wait(
             retry=retry,
             page_size=page_size,
             max_results=max_results,
+            query_results_format=query_results_format,
             callback=callback,
         )
 
@@ -594,6 +608,7 @@ def do_query():
                 retry=retry,
                 page_size=page_size,
                 max_results=max_results,
+                query_results_format=query_results_format,
                 callback=callback,
             )
 
@@ -633,6 +648,7 @@ def do_query():
             created=query_results.created,
             started=query_results.started,
             ended=query_results.ended,
+            query_results_format=query_results_format,
         )
 
     if job_retry is not None:
@@ -673,6 +689,7 @@ def _supported_by_jobs_query(request_body: Dict[str, Any]) -> bool:
         "jobTimeoutMs",
         "reservation",
         "maxSlots",
+        "queryResultsFormat",
     }
 
     unsupported_keys = request_keys - keys_allowlist
@@ -687,6 +704,7 @@ def _wait_or_cancel(
     page_size: Optional[int],
     max_results: Optional[int],
     *,
+    query_results_format: Optional[str] = None,
     callback: Callable = lambda _: None,
 ) -> table.RowIterator:
     """Wait for a job to complete and return the results.
@@ -731,6 +749,7 @@ def _wait_or_cancel(
                     ended=job.ended,
                 )
             )
+        query_results._query_results_format = query_results_format
         return query_results
     except Exception:
         # Attempt to cancel the job since we can't return the results.

packages/google-cloud-bigquery/google/cloud/bigquery/client.py

@@ -3665,6 +3665,8 @@ def query_and_wait(
         job_retry: retries.Retry = DEFAULT_JOB_RETRY,
         page_size: Optional[int] = None,
         max_results: Optional[int] = None,
+        query_results_format: Optional[str] = None,
+        compression_codec: Optional[str] = None,
     ) -> RowIterator:
         """Run the query, wait for it to finish, and return the results.
 
@@ -3712,6 +3714,10 @@ def query_and_wait(
                 by this parameter.
             max_results (Optional[int]):
                 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`).
+            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`).
 
         Returns:
             google.cloud.bigquery.table.RowIterator:
@@ -3742,6 +3748,8 @@ def query_and_wait(
             job_retry=job_retry,
             page_size=page_size,
             max_results=max_results,
+            query_results_format=query_results_format,
+            compression_codec=compression_codec,
         )
 
     def _query_and_wait_bigframes(
@@ -3757,6 +3765,8 @@ def _query_and_wait_bigframes(
         job_retry: retries.Retry = DEFAULT_JOB_RETRY,
         page_size: Optional[int] = None,
         max_results: Optional[int] = None,
+        query_results_format: Optional[str] = None,
+        compression_codec: Optional[str] = None,
         callback: Callable = lambda _: None,
     ) -> RowIterator:
         """See query_and_wait.
@@ -3789,6 +3799,8 @@ def _query_and_wait_bigframes(
             job_retry=job_retry,
             page_size=page_size,
             max_results=max_results,
+            query_results_format=query_results_format,
+            compression_codec=compression_codec,
             callback=callback,
         )
 

packages/google-cloud-bigquery/google/cloud/bigquery/enums.py

@@ -495,3 +495,20 @@ class TimestampPrecision(enum.Enum):
     """
     For TIMESTAMP type with picosecond precision.
     """
+
+
+class QueryResultsFormat(str, enum.Enum):
+    """[Beta] Format for query results response."""
+
+    ARROW = "ARROW"
+    """Specifies Apache Arrow format for query results."""
+
+
+class QueryResultsCompressionCodec(str, enum.Enum):
+    """[Beta] Compression codec for Arrow query results serialization."""
+
+    LZ4_FRAME = "LZ4_FRAME"
+    """Specifies LZ4_FRAME compression codec."""
+
+    ZSTD = "ZSTD"
+    """Specifies ZSTD compression codec."""