Skip to content

Commit 7a26910

Browse files
authored
TST: check compatibility with pyarrow types_mapper parameter (#44369)
1 parent 8149da0 commit 7a26910

File tree

1 file changed

+21
-0
lines changed

1 file changed

+21
-0
lines changed

pandas/tests/arrays/masked/test_arrow_compat.py

+21
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,27 @@ def test_arrow_roundtrip(data):
3636
tm.assert_frame_equal(result, df)
3737

3838

39+
def test_dataframe_from_arrow_types_mapper():
40+
def types_mapper(arrow_type):
41+
if pa.types.is_boolean(arrow_type):
42+
return pd.BooleanDtype()
43+
elif pa.types.is_integer(arrow_type):
44+
return pd.Int64Dtype()
45+
46+
bools_array = pa.array([True, None, False], type=pa.bool_())
47+
ints_array = pa.array([1, None, 2], type=pa.int64())
48+
small_ints_array = pa.array([-1, 0, 7], type=pa.int8())
49+
record_batch = pa.RecordBatch.from_arrays(
50+
[bools_array, ints_array, small_ints_array], ["bools", "ints", "small_ints"]
51+
)
52+
result = record_batch.to_pandas(types_mapper=types_mapper)
53+
bools = pd.Series([True, None, False], dtype="boolean")
54+
ints = pd.Series([1, None, 2], dtype="Int64")
55+
small_ints = pd.Series([-1, 0, 7], dtype="Int64")
56+
expected = pd.DataFrame({"bools": bools, "ints": ints, "small_ints": small_ints})
57+
tm.assert_frame_equal(result, expected)
58+
59+
3960
def test_arrow_load_from_zero_chunks(data):
4061
# GH-41040
4162

0 commit comments

Comments
 (0)