Skip to content

Commit aedc66f

Browse files
authored
feat: delegate ReadRowsPage.to_arrow to pandas_gbq.arrow (#17938)
Emits a deprecation warning on `ReadRowsPage.to_arrow()` and delegates Arrow RecordBatch decoding to `pandas_gbq.arrow.from_read_rows_response` when `pandas-gbq` is installed, falling back to internal parsing otherwise. **Changes Made** - `google/cloud/bigquery_storage_v1/reader.py`: Emits `PendingDeprecationWarning` and calls `pandas_gbq.arrow.from_read_rows_response` with stream schema. Falls back to `_stream_parser.to_arrow` if `pandas-gbq` is uninstalled. - `tests/unit/test_reader_v1_arrow.py`: Adds unit tests verifying delegation when `pandas-gbq` is present and fallback behavior when uninstalled. Fixes #<540939659> 🦕
1 parent 97e42b8 commit aedc66f

2 files changed

Lines changed: 69 additions & 0 deletions

File tree

packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import io
1919
import json
2020
import time
21+
import warnings
2122

2223
try:
2324
import fastavro
@@ -572,6 +573,25 @@ def to_arrow(self):
572573
pyarrow.RecordBatch:
573574
Rows from the message, as an Arrow record batch.
574575
"""
576+
warnings.warn(
577+
"Retrieving Arrow record batches directly via google-cloud-bigquery-storage is deprecated. "
578+
"Please use 'pandas_gbq.arrow.from_read_rows_response' or install 'pandas-gbq'.",
579+
PendingDeprecationWarning,
580+
stacklevel=2,
581+
)
582+
try:
583+
import pandas_gbq.arrow # type: ignore[import-not-found]
584+
585+
if hasattr(pandas_gbq.arrow, "from_read_rows_response"):
586+
if hasattr(self._stream_parser, "_parse_arrow_schema"):
587+
self._stream_parser._parse_arrow_schema()
588+
arrow_schema = getattr(self._stream_parser, "_schema", None)
589+
return pandas_gbq.arrow.from_read_rows_response(
590+
self._message, arrow_schema=arrow_schema
591+
)
592+
except ImportError:
593+
pass
594+
575595
return self._stream_parser.to_arrow(self._message)
576596

577597
def to_dataframe(self, dtypes=None):

packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -428,3 +428,52 @@ def test_to_dataframe_mid_stream_failure(mut, class_under_test, mock_gapic_clien
428428

429429
with pytest.raises(RuntimeError, match="Stream format changed mid-stream"):
430430
it.to_dataframe()
431+
432+
433+
def test_to_arrow_delegates_to_pandas_gbq_when_installed(mut):
434+
mock_parser = mock.Mock()
435+
mock_message = mock.Mock()
436+
page = mut.ReadRowsPage(mock_parser, mock_message)
437+
expected_batch = pyarrow.RecordBatch.from_arrays(
438+
[pyarrow.array([100])], names=["id"]
439+
)
440+
441+
mock_arrow_module = mock.Mock()
442+
mock_arrow_module.from_read_rows_response.return_value = expected_batch
443+
444+
mock_pandas_gbq = mock.Mock()
445+
mock_pandas_gbq.arrow = mock_arrow_module
446+
447+
with mock.patch.dict(
448+
"sys.modules",
449+
{"pandas_gbq": mock_pandas_gbq, "pandas_gbq.arrow": mock_arrow_module},
450+
):
451+
with pytest.warns(
452+
PendingDeprecationWarning,
453+
match="google-cloud-bigquery-storage is deprecated",
454+
):
455+
actual_batch = page.to_arrow()
456+
457+
assert actual_batch == expected_batch
458+
mock_arrow_module.from_read_rows_response.assert_called_once_with(
459+
mock_message, arrow_schema=mock_parser._schema
460+
)
461+
462+
463+
def test_to_arrow_falls_back_when_pandas_gbq_uninstalled(mut):
464+
mock_parser = mock.Mock()
465+
mock_message = mock.Mock()
466+
expected_batch = pyarrow.RecordBatch.from_arrays(
467+
[pyarrow.array([200])], names=["id"]
468+
)
469+
mock_parser.to_arrow.return_value = expected_batch
470+
with mock.patch.dict("sys.modules", {"pandas_gbq": None, "pandas_gbq.arrow": None}):
471+
page = mut.ReadRowsPage(mock_parser, mock_message)
472+
with pytest.warns(
473+
PendingDeprecationWarning,
474+
match="google-cloud-bigquery-storage is deprecated",
475+
):
476+
actual_batch = page.to_arrow()
477+
478+
assert actual_batch == expected_batch
479+
mock_parser.to_arrow.assert_called_once_with(mock_message)

0 commit comments

Comments
 (0)