fix(bigquery): raise when an Arrow batch arrives before its schema

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,

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