@@ -2832,9 +2832,12 @@ 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+ # Create dataframe with categorical column
28362836 df = (
2837- pd .DataFrame ([("A" , 1 ), ("B" , 2 ), ("C" , 3 )], columns = ["string_column" , "number_column" ])
2837+ DataFrame (
2838+ [("A" , 1 ), ("B" , 2 ), ("C" , 3 )],
2839+ columns = ["string_column" , "number_column" ],
2840+ )
28382841 .astype ({"string_column" : "string" , "number_column" : "float32" })
28392842 .astype ({"string_column" : "category" , "number_column" : "float32" })
28402843 )
@@ -2845,25 +2848,18 @@ def test_pivot_with_pyarrow_categorical(self):
28452848 buffer .seek (0 ) # Reset buffer position
28462849 df = pd .read_parquet (buffer , dtype_backend = "pyarrow" )
28472850
2848-
28492851 # Check that pivot works
28502852 df = df .pivot (columns = ["string_column" ], values = ["number_column" ])
28512853
28522854 # 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" )
2855+ multi_index = MultiIndex .from_arrays (
2856+ [["number_column" , "number_column" , "number_column" ], ["A" , "B" , "C" ]],
2857+ names = (None , "string_column" ),
28592858 )
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
2859+ df_expected = DataFrame (
2860+ [[1.0 , np .nan , np .nan ], [np .nan , 2.0 , np .nan ], [np .nan , np .nan , 3.0 ]],
2861+ columns = multi_index ,
2862+ )
2863+ tm .assert_frame_equal (
2864+ df , df_expected , check_dtype = False , check_column_type = False
28672865 )
2868- tm .assert_frame_equal (df , df_expected , check_dtype = False , check_column_type = False )
2869-
0 commit comments