@@ -2832,9 +2832,14 @@ def test_pivot_margins_with_none_index(self):
28322832 def test_pivot_with_pyarrow_categorical (self ):
28332833 # GH#53051
28342834
2835- # Create dataframe with categorical colum
2835+ pytest .importorskip ("pyarrow" )
2836+
2837+ # Create dataframe with categorical column
28362838 df = (
2837- pd .DataFrame ([("A" , 1 ), ("B" , 2 ), ("C" , 3 )], columns = ["string_column" , "number_column" ])
2839+ DataFrame (
2840+ [("A" , 1 ), ("B" , 2 ), ("C" , 3 )],
2841+ columns = ["string_column" , "number_column" ],
2842+ )
28382843 .astype ({"string_column" : "string" , "number_column" : "float32" })
28392844 .astype ({"string_column" : "category" , "number_column" : "float32" })
28402845 )
@@ -2845,25 +2850,18 @@ def test_pivot_with_pyarrow_categorical(self):
28452850 buffer .seek (0 ) # Reset buffer position
28462851 df = pd .read_parquet (buffer , dtype_backend = "pyarrow" )
28472852
2848-
28492853 # Check that pivot works
28502854 df = df .pivot (columns = ["string_column" ], values = ["number_column" ])
28512855
28522856 # Assert that values of result are correct to prevent silent failure
2853- multi_index = pd .MultiIndex .from_arrays (
2854- [
2855- ["number_column" , "number_column" , "number_column" ],
2856- ["A" , "B" , "C" ]
2857- ],
2858- names = (None , "string_column" )
2857+ multi_index = MultiIndex .from_arrays (
2858+ [["number_column" , "number_column" , "number_column" ], ["A" , "B" , "C" ]],
2859+ names = (None , "string_column" ),
28592860 )
2860- df_expected = pd .DataFrame (
2861- [
2862- [1.0 , np .nan , np .nan ],
2863- [np .nan , 2.0 , np .nan ],
2864- [np .nan , np .nan , 3.0 ]
2865- ],
2866- columns = multi_index
2861+ df_expected = DataFrame (
2862+ [[1.0 , np .nan , np .nan ], [np .nan , 2.0 , np .nan ], [np .nan , np .nan , 3.0 ]],
2863+ columns = multi_index ,
2864+ )
2865+ tm .assert_frame_equal (
2866+ df , df_expected , check_dtype = False , check_column_type = False
28672867 )
2868- tm .assert_frame_equal (df , df_expected , check_dtype = False , check_column_type = False )
2869-
0 commit comments