@@ -846,58 +846,74 @@ def test_consistency_for_boxed(self, box):
846
846
assert_frame_equal (result , expected )
847
847
848
848
849
- def zip_frames (* frames ):
849
+ def zip_frames (* frames , axis = 1 ):
850
850
"""
851
- take a list of frames, zip the columns together for each
852
- assume that these all have the first frame columns
851
+ take a list of frames, zip them together under the
852
+ assumption that these all have the first frames' index/ columns.
853
853
854
- return a new frame
854
+ Returns
855
+ -------
856
+ new_frame : DataFrame
855
857
"""
856
- columns = frames [0 ].columns
857
- zipped = [f [c ] for c in columns for f in frames ]
858
- return pd .concat (zipped , axis = 1 )
858
+ if axis == 1 :
859
+ columns = frames [0 ].columns
860
+ zipped = [f .loc [:, c ] for c in columns for f in frames ]
861
+ return pd .concat (zipped , axis = 1 )
862
+ else :
863
+ index = frames [0 ].index
864
+ zipped = [f .loc [i , :] for i in index for f in frames ]
865
+ return pd .DataFrame (zipped )
859
866
860
867
861
868
class TestDataFrameAggregate (TestData ):
862
869
863
- def test_agg_transform (self ):
870
+ def test_agg_transform (self , axis ):
871
+ other_axis = abs (axis - 1 )
864
872
865
873
with np .errstate (all = 'ignore' ):
866
874
867
- f_sqrt = np .sqrt (self .frame )
868
875
f_abs = np .abs (self .frame )
876
+ f_sqrt = np .sqrt (self .frame )
869
877
870
878
# ufunc
871
- result = self .frame .transform (np .sqrt )
879
+ result = self .frame .transform (np .sqrt , axis = axis )
872
880
expected = f_sqrt .copy ()
873
881
assert_frame_equal (result , expected )
874
882
875
- result = self .frame .apply (np .sqrt )
883
+ result = self .frame .apply (np .sqrt , axis = axis )
876
884
assert_frame_equal (result , expected )
877
885
878
- result = self .frame .transform (np .sqrt )
886
+ result = self .frame .transform (np .sqrt , axis = axis )
879
887
assert_frame_equal (result , expected )
880
888
881
889
# list-like
882
- result = self .frame .apply ([np .sqrt ])
890
+ result = self .frame .apply ([np .sqrt ], axis = axis )
883
891
expected = f_sqrt .copy ()
884
- expected .columns = pd .MultiIndex .from_product (
885
- [self .frame .columns , ['sqrt' ]])
892
+ if axis == 0 :
893
+ expected .columns = pd .MultiIndex .from_product (
894
+ [self .frame .columns , ['sqrt' ]])
895
+ else :
896
+ expected .index = pd .MultiIndex .from_product (
897
+ [self .frame .index , ['sqrt' ]])
886
898
assert_frame_equal (result , expected )
887
899
888
- result = self .frame .transform ([np .sqrt ])
900
+ result = self .frame .transform ([np .sqrt ], axis = axis )
889
901
assert_frame_equal (result , expected )
890
902
891
903
# multiple items in list
892
904
# these are in the order as if we are applying both
893
905
# functions per series and then concatting
894
- expected = zip_frames (f_sqrt , f_abs )
895
- expected .columns = pd .MultiIndex .from_product (
896
- [self .frame .columns , ['sqrt' , 'absolute' ]])
897
- result = self .frame .apply ([np .sqrt , np .abs ])
906
+ result = self .frame .apply ([np .abs , np .sqrt ], axis = axis )
907
+ expected = zip_frames (f_abs , f_sqrt , axis = other_axis )
908
+ if axis == 0 :
909
+ expected .columns = pd .MultiIndex .from_product (
910
+ [self .frame .columns , ['absolute' , 'sqrt' ]])
911
+ else :
912
+ expected .index = pd .MultiIndex .from_product (
913
+ [self .frame .index , ['absolute' , 'sqrt' ]])
898
914
assert_frame_equal (result , expected )
899
915
900
- result = self .frame .transform (['sqrt' , np . abs ] )
916
+ result = self .frame .transform ([np . abs , 'sqrt' ], axis = axis )
901
917
assert_frame_equal (result , expected )
902
918
903
919
def test_transform_and_agg_err (self , axis ):
@@ -985,13 +1001,16 @@ def test_agg_dict_nested_renaming_depr(self):
985
1001
986
1002
def test_agg_reduce (self , axis ):
987
1003
other_axis = abs (axis - 1 )
988
- name1 , name2 = self .frame .axes [other_axis ].unique ()[:2 ]
1004
+ name1 , name2 = self .frame .axes [other_axis ].unique ()[:2 ]. sort_values ()
989
1005
990
1006
# all reducers
991
- expected = zip_frames (self .frame .mean (axis = axis ).to_frame (),
992
- self .frame .max (axis = axis ).to_frame (),
993
- self .frame .sum (axis = axis ).to_frame ()).T
994
- expected .index = ['mean' , 'max' , 'sum' ]
1007
+ expected = pd .concat ([self .frame .mean (axis = axis ),
1008
+ self .frame .max (axis = axis ),
1009
+ self .frame .sum (axis = axis ),
1010
+ ], axis = 1 )
1011
+ expected .columns = ['mean' , 'max' , 'sum' ]
1012
+ expected = expected .T if axis == 0 else expected
1013
+
995
1014
result = self .frame .agg (['mean' , 'max' , 'sum' ], axis = axis )
996
1015
assert_frame_equal (result , expected )
997
1016
@@ -1001,7 +1020,7 @@ def test_agg_reduce(self, axis):
1001
1020
expected = Series ([self .frame .loc (other_axis )[name1 ].mean (),
1002
1021
self .frame .loc (other_axis )[name2 ].sum ()],
1003
1022
index = [name1 , name2 ])
1004
- assert_series_equal (result . reindex_like ( expected ) , expected )
1023
+ assert_series_equal (result , expected )
1005
1024
1006
1025
# dict input with lists
1007
1026
func = {name1 : ['mean' ], name2 : ['sum' ]}
@@ -1011,7 +1030,8 @@ def test_agg_reduce(self, axis):
1011
1030
index = ['mean' ]),
1012
1031
name2 : Series ([self .frame .loc (other_axis )[name2 ].sum ()],
1013
1032
index = ['sum' ])})
1014
- assert_frame_equal (result .reindex_like (expected ), expected )
1033
+ expected = expected .T if axis == 1 else expected
1034
+ assert_frame_equal (result , expected )
1015
1035
1016
1036
# dict input with lists with multiple
1017
1037
func = {name1 : ['mean' , 'sum' ],
@@ -1024,7 +1044,8 @@ def test_agg_reduce(self, axis):
1024
1044
name2 : Series ([self .frame .loc (other_axis )[name2 ].sum (),
1025
1045
self .frame .loc (other_axis )[name2 ].max ()],
1026
1046
index = ['sum' , 'max' ])})
1027
- assert_frame_equal (result .reindex_like (expected ), expected )
1047
+ expected = expected .T if axis == 1 else expected
1048
+ assert_frame_equal (result , expected )
1028
1049
1029
1050
def test_nuiscance_columns (self ):
1030
1051
0 commit comments