From f9d547aaa7baf2c888679a907df2c64d1ec86fd6 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 30 Jul 2026 21:23:15 +0000 Subject: [PATCH 1/3] feat(pandas-gbq): add arrow decoder for read rows response --- packages/pandas-gbq/pandas_gbq/__init__.py | 2 + packages/pandas-gbq/pandas_gbq/arrow.py | 48 ++++++++++++++++++ packages/pandas-gbq/tests/unit/test_arrow.py | 51 ++++++++++++++++++++ 3 files changed, 101 insertions(+) create mode 100644 packages/pandas-gbq/pandas_gbq/arrow.py create mode 100644 packages/pandas-gbq/tests/unit/test_arrow.py diff --git a/packages/pandas-gbq/pandas_gbq/__init__.py b/packages/pandas-gbq/pandas_gbq/__init__.py index a8a54179fd5c..b68c84f18f94 100644 --- a/packages/pandas-gbq/pandas_gbq/__init__.py +++ b/packages/pandas-gbq/pandas_gbq/__init__.py @@ -6,6 +6,7 @@ import sys import warnings +from pandas_gbq import arrow from pandas_gbq import version as pandas_gbq_version from pandas_gbq.contexts import Context, context from pandas_gbq.core.sample import sample @@ -32,4 +33,5 @@ "Context", "context", "sample", + "arrow", ] diff --git a/packages/pandas-gbq/pandas_gbq/arrow.py b/packages/pandas-gbq/pandas_gbq/arrow.py new file mode 100644 index 000000000000..a5e762e5dd4b --- /dev/null +++ b/packages/pandas-gbq/pandas_gbq/arrow.py @@ -0,0 +1,48 @@ +"""Arrow integration submodule for pandas-gbq.""" + +from typing import Any, Optional + +try: + import pyarrow as pa +except ImportError: + pa = None # type: ignore[assignment] + + +def from_read_rows_response( + message: Any, + arrow_schema: Optional[Any] = None, +) -> Any: + """Decodes a ReadRowsResponse protobuf message into a pyarrow.RecordBatch.""" + if pa is None: + raise ImportError( + "pyarrow is required to use 'from_read_rows_response'. " + "Please install pyarrow to use this function." + ) + + if ( + not hasattr(message, "arrow_record_batch") + or not message.arrow_record_batch.serialized_record_batch + ): + empty_schema = arrow_schema or pa.schema([]) + return pa.RecordBatch.from_pylist([], schema=empty_schema) + + serialized_batch = message.arrow_record_batch.serialized_record_batch + buffer = pa.py_buffer(serialized_batch) + + if arrow_schema is not None: + try: + return pa.ipc.read_record_batch(buffer, arrow_schema) + except Exception: + pass + + try: + reader = pa.ipc.RecordBatchStreamReader(buffer) + return reader.read_next_batch() + except Exception: + msg = pa.ipc.read_message(buffer) + batch_schema = ( + arrow_schema + if arrow_schema is not None + else getattr(msg, "schema", pa.schema([])) + ) + return pa.ipc.read_record_batch(msg, batch_schema) diff --git a/packages/pandas-gbq/tests/unit/test_arrow.py b/packages/pandas-gbq/tests/unit/test_arrow.py new file mode 100644 index 000000000000..b18c9c4aee30 --- /dev/null +++ b/packages/pandas-gbq/tests/unit/test_arrow.py @@ -0,0 +1,51 @@ +from unittest import mock + +import pyarrow as pa + +import pandas_gbq.arrow + + +def test_from_read_rows_response_valid_message_returns_record_batch(): + schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2]), pa.array(["alice", "bob"])], schema=schema + ) + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, schema) as writer: + writer.write_batch(batch) + serialized_bytes = sink.getvalue().to_pybytes() + + mock_message = mock.MagicMock() + mock_message.arrow_record_batch.serialized_record_batch = serialized_bytes + + result_batch = pandas_gbq.arrow.from_read_rows_response( + mock_message, arrow_schema=schema + ) + + assert result_batch.num_rows == 2 + assert result_batch.schema.names == ["id", "name"] + assert result_batch.column(0).to_pylist() == [1, 2] + assert result_batch.column(1).to_pylist() == ["alice", "bob"] + + +def test_from_read_rows_response_empty_message_returns_empty_batch(): + schema = pa.schema([("val", pa.float64())]) + mock_message = mock.MagicMock() + mock_message.arrow_record_batch.serialized_record_batch = b"" + + result_batch = pandas_gbq.arrow.from_read_rows_response( + mock_message, arrow_schema=schema + ) + + assert result_batch.num_rows == 0 + assert result_batch.schema == schema + + +def test_from_read_rows_response_uninstalled_pyarrow_raises_import_error(): + mock_message = mock.MagicMock() + + with mock.patch.object(pandas_gbq.arrow, "pa", None): + import pytest + + with pytest.raises(ImportError, match="pyarrow is required"): + pandas_gbq.arrow.from_read_rows_response(mock_message) From 622c0074bcb8bf977cf5f254c0968f2beddeb5d9 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 30 Jul 2026 14:32:04 -0700 Subject: [PATCH 2/3] Update packages/pandas-gbq/pandas_gbq/arrow.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/pandas-gbq/pandas_gbq/arrow.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/packages/pandas-gbq/pandas_gbq/arrow.py b/packages/pandas-gbq/pandas_gbq/arrow.py index a5e762e5dd4b..358b53886dbd 100644 --- a/packages/pandas-gbq/pandas_gbq/arrow.py +++ b/packages/pandas-gbq/pandas_gbq/arrow.py @@ -32,17 +32,17 @@ def from_read_rows_response( if arrow_schema is not None: try: return pa.ipc.read_record_batch(buffer, arrow_schema) - except Exception: + except pa.ArrowException: pass try: reader = pa.ipc.RecordBatchStreamReader(buffer) return reader.read_next_batch() - except Exception: + except pa.ArrowException: + if arrow_schema is None: + raise ValueError( + "arrow_schema is required to decode a serialized record batch message " + "when it is not formatted as an Arrow IPC stream." + ) msg = pa.ipc.read_message(buffer) - batch_schema = ( - arrow_schema - if arrow_schema is not None - else getattr(msg, "schema", pa.schema([])) - ) - return pa.ipc.read_record_batch(msg, batch_schema) + return pa.ipc.read_record_batch(msg, arrow_schema) From 108fca3fa17e1542aadd01750efa834e942f656f Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 30 Jul 2026 21:35:52 +0000 Subject: [PATCH 3/3] fix(pandas-gbq): catch OSError in arrow IPC deserialization --- packages/pandas-gbq/pandas_gbq/arrow.py | 4 ++-- packages/pandas-gbq/tests/unit/test_arrow.py | 20 ++++++++++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/packages/pandas-gbq/pandas_gbq/arrow.py b/packages/pandas-gbq/pandas_gbq/arrow.py index 358b53886dbd..429b04c591ed 100644 --- a/packages/pandas-gbq/pandas_gbq/arrow.py +++ b/packages/pandas-gbq/pandas_gbq/arrow.py @@ -32,13 +32,13 @@ def from_read_rows_response( if arrow_schema is not None: try: return pa.ipc.read_record_batch(buffer, arrow_schema) - except pa.ArrowException: + except (pa.ArrowException, OSError): pass try: reader = pa.ipc.RecordBatchStreamReader(buffer) return reader.read_next_batch() - except pa.ArrowException: + except (pa.ArrowException, OSError): if arrow_schema is None: raise ValueError( "arrow_schema is required to decode a serialized record batch message " diff --git a/packages/pandas-gbq/tests/unit/test_arrow.py b/packages/pandas-gbq/tests/unit/test_arrow.py index b18c9c4aee30..63d539b8ddf6 100644 --- a/packages/pandas-gbq/tests/unit/test_arrow.py +++ b/packages/pandas-gbq/tests/unit/test_arrow.py @@ -28,6 +28,26 @@ def test_from_read_rows_response_valid_message_returns_record_batch(): assert result_batch.column(1).to_pylist() == ["alice", "bob"] +def test_from_read_rows_response_serialized_record_batch_returns_record_batch(): + schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + batch = pa.RecordBatch.from_arrays( + [pa.array([10, 20]), pa.array(["carol", "dave"])], schema=schema + ) + serialized_bytes = batch.serialize().to_pybytes() + + mock_message = mock.MagicMock() + mock_message.arrow_record_batch.serialized_record_batch = serialized_bytes + + result_batch = pandas_gbq.arrow.from_read_rows_response( + mock_message, arrow_schema=schema + ) + + assert result_batch.num_rows == 2 + assert result_batch.schema.names == ["id", "name"] + assert result_batch.column(0).to_pylist() == [10, 20] + assert result_batch.column(1).to_pylist() == ["carol", "dave"] + + def test_from_read_rows_response_empty_message_returns_empty_batch(): schema = pa.schema([("val", pa.float64())]) mock_message = mock.MagicMock()