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