feat(pandas-gbq): add arrow decoder for read rows response - #17958
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new arrow integration submodule in pandas-gbq to decode ReadRowsResponse protobuf messages into pyarrow.RecordBatch objects, along with corresponding unit tests. The review feedback points out that pyarrow.Message does not have a schema attribute, which would cause decoding to fail or fall back incorrectly. Additionally, it is recommended to catch specific pyarrow.ArrowException errors instead of using broad except Exception blocks.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
| mock_message = mock.MagicMock() | ||
|
|
||
| with mock.patch.object(pandas_gbq.arrow, "pa", None): | ||
| import pytest |
There was a problem hiding this comment.
nit: let's import pytest at the top of the file.
There was a problem hiding this comment.
Done, I move import pytest to the top-level imports.
| import pyarrow as pa | ||
|
|
||
| import pandas_gbq.arrow | ||
|
|
There was a problem hiding this comment.
Could we add a test case where the arrow schema is not provided?
There was a problem hiding this comment.
Sure, test case is added.
3 participants
Adds
pandas_gbq.arrow.from_read_rows_responseto decode BigQuery Storage Read APIReadRowsResponseprotobuf messages intopyarrow.RecordBatchinstances.Changes
pandas_gbq/arrow.py: Implementsfrom_read_rows_response(message, arrow_schema=None)to deserialize IPC stream and message buffers with empty-batch fallback.pandas_gbq/__init__.py: Exposesarrowin package__all__.tests/unit/test_arrow.py: Adds unit tests for valid batch parsing, empty message handling, and missingpyarrowerror handling.Fixes #<540939659> 🦕