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..429b04c591ed --- /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 (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) 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..63d539b8ddf6 --- /dev/null +++ b/packages/pandas-gbq/tests/unit/test_arrow.py @@ -0,0 +1,71 @@ +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_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() + 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)