feat(bigquery): support queryResultsFormat and compressionCodec in query_and_wait

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.
@@ -473,8 +475,10 @@ def query_and_wait(
         page_size (Optional[int]):
             The maximum number of rows in each page of results from this
             request. Non-positive values are ignored.
-        max_results (Optional[int]):
-            The maximum total number of rows from this request.
+        query_results_format (Optional[str]):
+            [Beta] The format for query results (e.g. "ARROW").
+        compression_codec (Optional[str]):
+            [Beta] Compression codec for Arrow serialization (e.g. "LZ4_FRAME").
         callback (Callable):
             A callback function used by bigframes to report query progress.
 
@@ -499,6 +503,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 +533,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 +606,7 @@ def do_query():
                 retry=retry,
                 page_size=page_size,
                 max_results=max_results,
+                query_results_format=query_results_format,
                 callback=callback,
             )
 
@@ -633,6 +646,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 +687,7 @@ def _supported_by_jobs_query(request_body: Dict[str, Any]) -> bool:
         "jobTimeoutMs",
         "reservation",
         "maxSlots",
+        "queryResultsFormat",
     }
 
     unsupported_keys = request_keys - keys_allowlist
@@ -687,6 +702,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 +747,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

@@ -3649,6 +3649,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.
 
@@ -3694,8 +3696,10 @@ def query_and_wait(
                 jobs.getQueryResults API calls. Large results downloaded with
                 the BigQuery Storage Read API are intentionally unaffected
                 by this parameter.
-            max_results (Optional[int]):
-                The maximum total number of rows from this request.
+            query_results_format (Optional[str]):
+                [Beta] The format for query results (e.g. "ARROW").
+            compression_codec (Optional[str]):
+                [Beta] Compression codec for Arrow serialization (e.g. "LZ4_FRAME").
 
         Returns:
             google.cloud.bigquery.table.RowIterator:
@@ -3726,6 +3730,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(
@@ -3741,6 +3747,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.
@@ -3773,6 +3781,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/table.py

@@ -16,6 +16,7 @@
 
 from __future__ import absolute_import
 
+import base64
 import copy
 import datetime
 import functools
@@ -1912,6 +1913,7 @@ def __init__(
         created: Optional[datetime.datetime] = None,
         started: Optional[datetime.datetime] = None,
         ended: Optional[datetime.datetime] = None,
+        query_results_format: Optional[str] = None,
     ):
         super(RowIterator, self).__init__(
             client,
@@ -1945,6 +1947,29 @@ def __init__(
         self._job_created = created
         self._job_started = started
         self._job_ended = ended
+        self._query_results_format = query_results_format
+
+    @property
+    def pages(self):
+        if self._query_results_format == "ARROW":
+            raise ValueError(
+                "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead."
+            )
+        return super().pages
+
+    def __iter__(self):
+        if self._query_results_format == "ARROW":
+            raise ValueError(
+                "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead."
+            )
+        return super().__iter__()
+
+    def __next__(self):
+        if self._query_results_format == "ARROW":
+            raise ValueError(
+                "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead."
+            )
+        return super().__next__()
 
     @property
     def _billing_project(self) -> Optional[str]:
@@ -2226,6 +2251,12 @@ def to_arrow_iterable(
 
         .. versionadded:: 2.31.0
         """
+        if self._query_results_format == "ARROW":
+            return self._download_arrow_from_job_id(
+                bqstorage_client=bqstorage_client,
+                timeout=timeout,
+            )
+
         self._maybe_warn_max_results(bqstorage_client)
 
         bqstorage_download = functools.partial(
@@ -2251,6 +2282,90 @@ def to_arrow_iterable(
             bqstorage_client=bqstorage_client,
         )
 
+    def _download_arrow_from_job_id(
+        self,
+        bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None,
+        timeout: Optional[float] = None,
+    ) -> Iterator["pyarrow.RecordBatch"]:
+        if pyarrow is None:
+            raise ValueError(_NO_PYARROW_ERROR)
+
+        offset = 0
+        pa_schema = None
+        total_rows = self.total_rows
+        job_complete = False
+
+        if self._first_page_response:
+            first_page = self._first_page_response
+            self._first_page_response = None
+
+            job_complete = bool(first_page.get("jobComplete", False))
+            if job_complete:
+                total_rows = int(first_page["totalRows"])
+
+            arrow_schema_json = first_page.get("arrowSchema")
+            if isinstance(arrow_schema_json, dict):
+                schema_bytes = arrow_schema_json.get("serializedSchema")
+                if schema_bytes:
+                    if isinstance(schema_bytes, str):
+                        schema_bytes = base64.b64decode(schema_bytes)
+                    pa_schema = pyarrow.ipc.read_schema(
+                        pyarrow.py_buffer(schema_bytes)
+                    )
+
+            arrow_batch_json = first_page.get("arrowRecordBatch")
+            if isinstance(arrow_batch_json, dict) and pa_schema is not None:
+                batch_bytes = arrow_batch_json.get("serializedRecordBatch")
+                if batch_bytes:
+                    if isinstance(batch_bytes, str):
+                        batch_bytes = base64.b64decode(batch_bytes)
+                    batch = pyarrow.ipc.read_record_batch(
+                        pyarrow.py_buffer(batch_bytes),
+                        pa_schema,
+                    )
+                    offset += batch.num_rows
+                    yield batch
+
+        if job_complete and offset >= total_rows:
+            return
+
+        if bqstorage_client is None:
+            if self.client is None:
+                raise ValueError("RowIterator client is None.")
+            bqstorage_client = self.client._ensure_bqstorage_client()
+            if bqstorage_client is None:
+                raise ValueError(
+                    "The google-cloud-bigquery-storage library is required to read Arrow results."
+                )
+
+        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"
+        )
+        reader = bqstorage_client.read_rows(
+            stream_name, offset=offset, timeout=timeout
+        )
+        for response in reader:
+            if (
+                response.arrow_schema
+                and response.arrow_schema.serialized_schema
+                and pa_schema is None
+            ):
+                pa_schema = pyarrow.ipc.read_schema(
+                    pyarrow.py_buffer(response.arrow_schema.serialized_schema)
+                )
+            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 changing the signature of this method, make sure to apply the same
     # changes to job.QueryJob.to_arrow()
     def to_arrow(
@@ -2357,7 +2472,7 @@ def to_arrow(
                 # but mypy cannot infer this correlation. We ignore the union-attr error here.
                 bqstorage_client._transport.close()  # type: ignore[union-attr]
 
-        if record_batches and bqstorage_client is not None:
+        if record_batches and (bqstorage_client is not None or self._query_results_format == "ARROW"):
             return pyarrow.Table.from_batches(record_batches)
         else:
             # No records (not record_batches), use schema based on BigQuery schema
@@ -3036,17 +3151,41 @@ class _EmptyRowIterator(RowIterator):
     """
 
     def __init__(
-        self, client=None, api_request=None, path=None, schema=(), *args, **kwargs
+        self, client=None, api_request=None, path=None, schema=(), *args, query_results_format: Optional[str] = None, **kwargs
     ):
         super().__init__(
             client=client,
             api_request=api_request,
             path=path,
             schema=schema,
+            query_results_format=query_results_format,
             *args,
             **kwargs,
         )
         self._total_rows = 0
+        self._query_results_format = query_results_format
+
+    @property
+    def pages(self):
+        if self._query_results_format == "ARROW":
+            raise ValueError(
+                "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead."
+            )
+        return super().pages
+
+    def __iter__(self):
+        if self._query_results_format == "ARROW":
+            raise ValueError(
+                "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead."
+            )
+        return iter(())
+
+    def __next__(self):
+        if self._query_results_format == "ARROW":
+            raise ValueError(
+                "Cannot iterate over non-arrow results when queryResultsFormat is ARROW. Use to_arrow_iterable() or to_arrow() instead."
+            )
+        raise StopIteration
 
     def to_arrow(
         self,
@@ -3219,11 +3358,10 @@ def to_arrow_iterable(
         Returns:
             An iterator yielding a single empty :class:`~pyarrow.RecordBatch`.
         """
+        if pyarrow is None:
+            raise ValueError(_NO_PYARROW_ERROR)
         return iter((pyarrow.record_batch([]),))
 
-    def __iter__(self):
-        return iter(())
-
 
 class PartitionRange(object):
     """Definition of the ranges for range partitioning.