@@ -36,6 +36,27 @@ def test_arrow_roundtrip(data):
36
36
tm .assert_frame_equal (result , df )
37
37
38
38
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
+
39
60
def test_arrow_load_from_zero_chunks (data ):
40
61
# GH-41040
41
62
0 commit comments