Visitar URL original
fix(bigquery): raise when an Arrow batch arrives before its schema (#… · googleapis/google-cloud-python@7944a4c · GitHub
Skip to content

Commit 7944a4c

Browse files
ohmayrohmayr
authored
fix(bigquery): raise when an Arrow batch arrives before its schema (#18791)
Raises an error when an Arrow record batch arrives before its schema. Before, `download_arrow_from_job()` skipped such batches, so rows could go missing from Arrow query results without any error. The skip has been there since Arrow support was added in #18027, so this fixes it on `main` on its own. #18651 (which reviewers flagged it on) will rebase on top of this. - Raises `ValueError` for a batch without a schema in the inline `jobs.query` first page. - Raises `ValueError` for a batch without a schema in the `ReadRows` stream. - Adds unit tests for both cases.
1 parent 5ac7cc6 commit 7944a4c

2 files changed

Lines changed: 62 additions & 2 deletions

File tree

‎packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@
3737
"pyarrow to use the to_arrow() function."
3838
)
3939

40+
_NO_ARROW_SCHEMA_ERROR = "Got an Arrow record batch before its schema."
41+
4042

4143
def validate_query(
4244
query_results_format: Optional[str],
@@ -168,9 +170,11 @@ def download_arrow_from_job(
168170
pa_schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes))
169171

170172
arrow_batch_json = first_page.get("arrowRecordBatch")
171-
if isinstance(arrow_batch_json, dict) and pa_schema is not None:
173+
if isinstance(arrow_batch_json, dict):
172174
batch_bytes = arrow_batch_json.get("serializedRecordBatch")
173175
if batch_bytes:
176+
if pa_schema is None:
177+
raise ValueError(_NO_ARROW_SCHEMA_ERROR)
174178
if isinstance(batch_bytes, str):
175179
batch_bytes = base64.b64decode(batch_bytes)
176180
batch = pyarrow.ipc.read_record_batch(
@@ -214,8 +218,9 @@ def download_arrow_from_job(
214218
if (
215219
response.arrow_record_batch
216220
and response.arrow_record_batch.serialized_record_batch
217-
and pa_schema is not None
218221
):
222+
if pa_schema is None:
223+
raise ValueError(_NO_ARROW_SCHEMA_ERROR)
219224
batch = pyarrow.ipc.read_record_batch(
220225
pyarrow.py_buffer(response.arrow_record_batch.serialized_record_batch),
221226
pa_schema,

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -431,6 +431,61 @@ def test_download_arrow_from_job_id_with_schema_and_batch(self):
431431
b"batch_bytes", "fake_schema"
432432
)
433433

434+
def test_download_arrow_from_job_id_first_page_batch_without_schema_raises(self):
435+
mock_client = mock.MagicMock()
436+
mock_bqstorage = mock.MagicMock()
437+
mock_client._ensure_bqstorage_client.return_value = mock_bqstorage
438+
439+
b64_batch = base64.b64encode(b"batch_bytes").decode("ascii")
440+
iterator = RowIterator(
441+
client=mock_client,
442+
api_request=mock.MagicMock(),
443+
path=None,
444+
schema=(),
445+
project="test-proj",
446+
location="US",
447+
job_id="test-job-no-schema",
448+
query_results_format="ARROW",
449+
first_page_response={
450+
"jobComplete": True,
451+
"arrowRecordBatch": {"serializedRecordBatch": b64_batch},
452+
},
453+
)
454+
455+
with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow:
456+
with self.assertRaisesRegex(ValueError, "before its schema"):
457+
list(iterator._download_arrow_from_job_id(timeout=5.0))
458+
459+
mock_pyarrow.ipc.read_record_batch.assert_not_called()
460+
mock_bqstorage.read_rows.assert_not_called()
461+
462+
def test_download_arrow_from_job_id_read_rows_batch_without_schema_raises(self):
463+
mock_client = mock.MagicMock()
464+
mock_bqstorage = mock.MagicMock()
465+
mock_client._ensure_bqstorage_client.return_value = mock_bqstorage
466+
467+
iterator = RowIterator(
468+
client=mock_client,
469+
api_request=mock.MagicMock(),
470+
path=None,
471+
schema=(),
472+
project="test-proj",
473+
location="US",
474+
job_id="test-job-no-schema",
475+
query_results_format="ARROW",
476+
)
477+
478+
mock_response = mock.MagicMock()
479+
mock_response.arrow_schema = None
480+
mock_response.arrow_record_batch.serialized_record_batch = b"batch_bytes"
481+
mock_bqstorage.read_rows.return_value = [mock_response]
482+
483+
with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow:
484+
with self.assertRaisesRegex(ValueError, "before its schema"):
485+
list(iterator._download_arrow_from_job_id(timeout=5.0))
486+
487+
mock_pyarrow.ipc.read_record_batch.assert_not_called()
488+
434489
def test_empty_row_iterator_to_arrow_iterable_checks_pyarrow(self):
435490
iterator = _EmptyRowIterator(query_results_format="ARROW")
436491
with mock.patch("google.cloud.bigquery.table.pyarrow", None):

0 commit comments

Comments
 (0)