feat(pandas-gbq): add arrow decoder for read rows response

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",
 ]

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 (pa.ArrowException, OSError):
+            pass
+
+    try:
+        reader = pa.ipc.RecordBatchStreamReader(buffer)
+        return reader.read_next_batch()
+    except (pa.ArrowException, OSError):
+        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)
+        return pa.ipc.read_record_batch(msg, arrow_schema)

packages/pandas-gbq/tests/unit/test_arrow.py

@@ -0,0 +1,116 @@
+from unittest import mock
+
+import pandas_gbq.arrow
+import pyarrow as pa
+import pytest
+
+
+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_no_schema_provided_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)
+
+    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_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_serialized_batch_without_schema_raises_value_error():
+    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
+
+    with pytest.raises(
+        ValueError, match="arrow_schema is required to decode a serialized record batch"
+    ):
+        pandas_gbq.arrow.from_read_rows_response(mock_message)
+
+
+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_empty_message_without_schema_returns_empty_batch():
+    mock_message = mock.MagicMock()
+    mock_message.arrow_record_batch.serialized_record_batch = b""
+
+    result_batch = pandas_gbq.arrow.from_read_rows_response(mock_message)
+
+    assert result_batch.num_rows == 0
+    assert result_batch.schema == pa.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):
+        with pytest.raises(ImportError, match="pyarrow is required"):
+            pandas_gbq.arrow.from_read_rows_response(mock_message)