feat(bigquery): support Arrow destination tables

BigQuery Storage Read API streams/_default only supports anonymous query
result tables and rejects explicit destination tables with
INVALID_ARGUMENT.

When query_and_wait falls back to jobs.insert with an explicit
destination table, open a single-stream Arrow read session via
create_read_session (passing compression_codec in read_options) and
stream batches over read_rows while preserving 0-row Arrow schemas.

Author: ohmayr

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

@@ -517,7 +517,7 @@ def query_and_wait(
     # Some API parameters aren't supported by the jobs.query API. In these
     # cases, fallback to a jobs.insert call.
     if not _supported_by_jobs_query(request_body):
-        return _wait_or_cancel(
+        rows = _wait_or_cancel(
             query_jobs_insert(
                 client=client,
                 query=query,
@@ -540,6 +540,9 @@ def query_and_wait(
             compression_codec=compression_codec,
             callback=callback,
         )
+        if "destinationTable" in request_body:
+            rows._use_read_session = True
+        return rows
 
     path = _to_query_path(project)
 

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

@@ -2374,10 +2374,46 @@ def _stream_arrow_via_bqstorage(
         remaining = (
             self.max_results - offset if self.max_results is not None 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"
+        use_read_session = (
+            getattr(self, "_use_read_session", False) and self._table is not None
+        )
+        if use_read_session:
+            table_ref = (
+                f"projects/{self._table.project}/datasets/"
+                f"{self._table.dataset_id}/tables/{self._table.table_id}"
+            )
+            read_session: Dict[str, Any] = {
+                "table": table_ref,
+                "data_format": "ARROW",
+            }
+            if self._compression_codec is not None:
+                read_session["read_options"] = {
+                    "arrow_serialization_options": {
+                        "buffer_compression": self._compression_codec
+                    }
+                }
+            session = bqstorage_client.create_read_session(
+                parent=f"projects/{project}",
+                read_session=read_session,
+                max_stream_count=1,
+            )
+            if (
+                session.arrow_schema
+                and session.arrow_schema.serialized_schema
+                and pa_schema is None
+            ):
+                pa_schema = pyarrow.ipc.read_schema(
+                    pyarrow.py_buffer(session.arrow_schema.serialized_schema)
+                )
+                self._arrow_schema = pa_schema
+            if not session.streams:
+                return
+            stream_name = session.streams[0].name
+        else:
+            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"
         read_kwargs: Dict[str, Any] = {"offset": offset, "timeout": timeout}
-        if self._compression_codec is not None:
+        if self._compression_codec is not None and not use_read_session:
             try:
                 reader = bqstorage_client.read_rows(
                     stream_name,

packages/google-cloud-bigquery/tests/system/test_arrow.py

@@ -474,3 +474,44 @@ def test_query_arrow_cached_multi_page(bigquery_client):
     assert len(first_table) == 5000
     assert len(cached_table) == 5000
     assert cached_table.equals(first_table)
+
+
+def test_query_arrow_explicit_destination_table(
+    dataset_client, dataset_id, test_table_name
+):
+    """System test verifying explicit destination tables on QueryJobConfig read via CreateReadSession."""
+    dest_table = dataset_client.dataset(dataset_id).table(test_table_name)
+    job_config = bigquery.QueryJobConfig(
+        destination=dest_table,
+        write_disposition=bigquery.WriteDisposition.WRITE_TRUNCATE,
+    )
+
+    # 1. Full + max_results read with LZ4_FRAME compression on explicit destination table
+    results = dataset_client.query_and_wait(
+        "SELECT num, CONCAT('row_', CAST(num AS STRING)) AS label "
+        "FROM UNNEST(GENERATE_ARRAY(1, 2000)) AS num ORDER BY num",
+        job_config=job_config,
+        query_results_format=enums.QueryResultsFormat.ARROW,
+        compression_codec=enums.QueryResultsCompressionCodec.LZ4_FRAME,
+        max_results=1200,
+    )
+    table = pyarrow.Table.from_batches(results.to_arrow_iterable())
+    assert len(table) == 1200
+    assert table.column_names == ["num", "label"]
+    assert table.column("num").to_pylist() == list(range(1, 1201))
+
+    # 2. Zero-row explicit destination table preserves pyarrow.Schema
+    empty_dest_table = dataset_client.dataset(dataset_id).table(
+        f"{test_table_name}_empty"
+    )
+    empty_results = dataset_client.query_and_wait(
+        "SELECT 1 AS num, 'abc' AS label FROM UNNEST(GENERATE_ARRAY(1, 5)) AS x WHERE x > 100",
+        job_config=bigquery.QueryJobConfig(destination=empty_dest_table),
+        query_results_format=enums.QueryResultsFormat.ARROW,
+    )
+    assert empty_results.total_rows == 0
+    empty_table = empty_results.to_arrow()
+    assert len(empty_table) == 0
+    assert empty_table.column_names == ["num", "label"]
+    assert empty_table.schema.field("num").type == pyarrow.int64()
+    assert empty_table.schema.field("label").type == pyarrow.string()

packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py

@@ -914,6 +914,72 @@ def test_download_arrow_from_job_id_cached_query_with_page_token_reads_stream(se
                 timeout=5.0,
             )
 
+    def test_download_arrow_from_job_id_explicit_destination_uses_create_read_session(
+        self,
+    ):
+        from google.cloud.bigquery.table import TableReference
+
+        mock_client = mock.MagicMock()
+        mock_bqstorage = mock.MagicMock()
+        mock_client._ensure_bqstorage_client.return_value = mock_bqstorage
+
+        mock_stream = mock.MagicMock()
+        mock_stream.name = "projects/test-proj/locations/US/sessions/s1/streams/0"
+        mock_session = mock.MagicMock()
+        mock_session.arrow_schema.serialized_schema = b"session_schema_bytes"
+        mock_session.streams = [mock_stream]
+        mock_bqstorage.create_read_session.return_value = mock_session
+
+        mock_resp = mock.MagicMock()
+        mock_resp.arrow_schema = None
+        mock_resp.arrow_record_batch.serialized_record_batch = b"batch_bytes"
+        mock_bqstorage.read_rows.return_value = [mock_resp]
+
+        dest_table = TableReference.from_string("test-proj.dest_ds.dest_tbl")
+        iterator = RowIterator(
+            client=mock_client,
+            api_request=mock.MagicMock(),
+            path=None,
+            schema=(),
+            table=dest_table,
+            project="test-proj",
+            location="US",
+            job_id="test-job-dest",
+            query_results_format="ARROW",
+            compression_codec="LZ4_FRAME",
+        )
+        iterator._use_read_session = True
+
+        mock_batch = mock.MagicMock()
+        mock_batch.num_rows = 3
+
+        with mock.patch("google.cloud.bigquery.table.pyarrow") as mock_pyarrow:
+            mock_pyarrow.py_buffer = lambda x: x
+            mock_pyarrow.ipc.read_schema.return_value = "session_schema"
+            mock_pyarrow.ipc.read_record_batch.return_value = mock_batch
+
+            batches = list(iterator._download_arrow_from_job_id(timeout=5.0))
+            self.assertEqual(batches, [mock_batch])
+            self.assertEqual(iterator._arrow_schema, "session_schema")
+            mock_bqstorage.create_read_session.assert_called_once_with(
+                parent="projects/test-proj",
+                read_session={
+                    "table": "projects/test-proj/datasets/dest_ds/tables/dest_tbl",
+                    "data_format": "ARROW",
+                    "read_options": {
+                        "arrow_serialization_options": {
+                            "buffer_compression": "LZ4_FRAME"
+                        }
+                    },
+                },
+                max_stream_count=1,
+            )
+            mock_bqstorage.read_rows.assert_called_once_with(
+                "projects/test-proj/locations/US/sessions/s1/streams/0",
+                offset=0,
+                timeout=5.0,
+            )
+
 
 if __name__ == "__main__":
     unittest.main()