import numpy as np import pytest from pandas.errors import DtypeWarning import pandas._testing as tm from pandas.core.arrays import ArrowExtensionArray from pandas.io.parsers.c_parser_wrapper import _concatenate_chunks def test_concatenate_chunks_pyarrow(): # GH#51876 pa = pytest.importorskip("pyarrow") chunks = [ {0: ArrowExtensionArray(pa.array([1.5, 2.5]))}, {0: ArrowExtensionArray(pa.array([1, 2]))}, ] result = _concatenate_chunks(chunks) expected = ArrowExtensionArray(pa.array([1.5, 2.5, 1.0, 2.0])) tm.assert_extension_array_equal(result[0], expected) def test_concatenate_chunks_pyarrow_strings(): # GH#51876 pa = pytest.importorskip("pyarrow") chunks = [ {0: ArrowExtensionArray(pa.array([1.5, 2.5]))}, {0: ArrowExtensionArray(pa.array(["a", "b"]))}, ] with tm.assert_produces_warning(DtypeWarning, match="have mixed types"): result = _concatenate_chunks(chunks) expected = np.concatenate( [np.array([1.5, 2.5], dtype=object), np.array(["a", "b"])] ) tm.assert_numpy_array_equal(result[0], expected)