From 568a0604200c2ffa684a1bbfb66c183bc29ac2de Mon Sep 17 00:00:00 2001
From: ohmayr <omairn@google.com>
Date: Tue, 20 Oct 2026 20:05:38 +0000
Subject: [PATCH] fix(bigquery): raise when an Arrow batch arrives before its
 schema

download_arrow_from_job() skipped record batches that arrived before
any Arrow schema, in both the inline first page and the ReadRows
stream, so rows could go missing without an error. Raise ValueError
instead. Normal responses always send the schema first, so this only
fires on a malformed response.
---
 .../google/cloud/bigquery/_arrow_query.py     |  9 ++-
 .../unit/test_query_results_format_arrow.py   | 55 +++++++++++++++++++
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py
index ca98bc18f39d..bf98b835d82c 100644
--- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py
+++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py
@@ -37,6 +37,8 @@
     "pyarrow to use the to_arrow() function."
 )
 
+_NO_ARROW_SCHEMA_ERROR = "Got an Arrow record batch before its schema."
+
 
 def validate_query(
     query_results_format: Optional[str],
@@ -169,9 +171,11 @@ def download_arrow_from_job(
                 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:
+        if isinstance(arrow_batch_json, dict):
             batch_bytes = arrow_batch_json.get("serializedRecordBatch")
             if batch_bytes:
+                if pa_schema is None:
+                    raise ValueError(_NO_ARROW_SCHEMA_ERROR)
                 if isinstance(batch_bytes, str):
                     batch_bytes = base64.b64decode(batch_bytes)
                 batch = pyarrow.ipc.read_record_batch(
@@ -215,8 +219,9 @@ def download_arrow_from_job(
         if (
             response.arrow_record_batch
             and response.arrow_record_batch.serialized_record_batch
-            and pa_schema is not None
         ):
+            if pa_schema is None:
+                raise ValueError(_NO_ARROW_SCHEMA_ERROR)
             batch = pyarrow.ipc.read_record_batch(
                 pyarrow.py_buffer(response.arrow_record_batch.serialized_record_batch),
                 pa_schema,
diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py
index 1855d854c470..4e407fcfc611 100644
--- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py
+++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py
@@ -431,6 +431,61 @@ def test_download_arrow_from_job_id_with_schema_and_batch(self):
                 b"batch_bytes", "fake_schema"
             )
 
+    def test_download_arrow_from_job_id_first_page_batch_without_schema_raises(self):
+        mock_client = mock.MagicMock()
+        mock_bqstorage = mock.MagicMock()
+        mock_client._ensure_bqstorage_client.return_value = mock_bqstorage
+
+        b64_batch = base64.b64encode(b"batch_bytes").decode("ascii")
+        iterator = RowIterator(
+            client=mock_client,
+            api_request=mock.MagicMock(),
+            path=None,
+            schema=(),
+            project="test-proj",
+            location="US",
+            job_id="test-job-no-schema",
+            query_results_format="ARROW",
+            first_page_response={
+                "jobComplete": True,
+                "arrowRecordBatch": {"serializedRecordBatch": b64_batch},
+            },
+        )
+
+        with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow:
+            with self.assertRaisesRegex(ValueError, "before its schema"):
+                list(iterator._download_arrow_from_job_id(timeout=5.0))
+
+        mock_pyarrow.ipc.read_record_batch.assert_not_called()
+        mock_bqstorage.read_rows.assert_not_called()
+
+    def test_download_arrow_from_job_id_read_rows_batch_without_schema_raises(self):
+        mock_client = mock.MagicMock()
+        mock_bqstorage = mock.MagicMock()
+        mock_client._ensure_bqstorage_client.return_value = mock_bqstorage
+
+        iterator = RowIterator(
+            client=mock_client,
+            api_request=mock.MagicMock(),
+            path=None,
+            schema=(),
+            project="test-proj",
+            location="US",
+            job_id="test-job-no-schema",
+            query_results_format="ARROW",
+        )
+
+        mock_response = mock.MagicMock()
+        mock_response.arrow_schema = None
+        mock_response.arrow_record_batch.serialized_record_batch = b"batch_bytes"
+        mock_bqstorage.read_rows.return_value = [mock_response]
+
+        with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow:
+            with self.assertRaisesRegex(ValueError, "before its schema"):
+                list(iterator._download_arrow_from_job_id(timeout=5.0))
+
+        mock_pyarrow.ipc.read_record_batch.assert_not_called()
+
     def test_empty_row_iterator_to_arrow_iterable_checks_pyarrow(self):
         iterator = _EmptyRowIterator(query_results_format="ARROW")
         with mock.patch("google.cloud.bigquery.table.pyarrow", None):