@@ -487,53 +487,30 @@ def test_nth_multi_index_as_expected():
487
487
tm .assert_frame_equal (result , expected )
488
488
489
489
490
- def test_groupby_head_tail ():
490
+ @pytest .mark .parametrize (
491
+ "op, n, expected_rows" ,
492
+ [
493
+ ("head" , - 1 , []),
494
+ ("head" , 0 , []),
495
+ ("head" , 1 , [0 , 2 ]),
496
+ ("head" , 7 , [0 , 1 , 2 ]),
497
+ ("tail" , - 1 , []),
498
+ ("tail" , 0 , []),
499
+ ("tail" , 1 , [1 , 2 ]),
500
+ ("tail" , 7 , [0 , 1 , 2 ]),
501
+ ],
502
+ )
503
+ @pytest .mark .parametrize ("columns" , [None , [], ["A" ], ["B" ], ["A" , "B" ]])
504
+ @pytest .mark .parametrize ("as_index" , [True , False ])
505
+ def test_groupby_head_tail (op , n , expected_rows , columns , as_index ):
491
506
df = DataFrame ([[1 , 2 ], [1 , 4 ], [5 , 6 ]], columns = ["A" , "B" ])
492
- g_as = df .groupby ("A" , as_index = True )
493
- g_not_as = df .groupby ("A" , as_index = False )
494
-
495
- # as_index= False, much easier
496
- tm .assert_frame_equal (df .loc [[0 , 2 ]], g_not_as .head (1 ))
497
- tm .assert_frame_equal (df .loc [[1 , 2 ]], g_not_as .tail (1 ))
498
-
499
- empty_not_as = DataFrame (columns = df .columns , index = Index ([], dtype = df .index .dtype ))
500
- empty_not_as ["A" ] = empty_not_as ["A" ].astype (df .A .dtype )
501
- empty_not_as ["B" ] = empty_not_as ["B" ].astype (df .B .dtype )
502
- tm .assert_frame_equal (empty_not_as , g_not_as .head (0 ))
503
- tm .assert_frame_equal (empty_not_as , g_not_as .tail (0 ))
504
- tm .assert_frame_equal (empty_not_as , g_not_as .head (- 1 ))
505
- tm .assert_frame_equal (empty_not_as , g_not_as .tail (- 1 ))
506
-
507
- tm .assert_frame_equal (df , g_not_as .head (7 )) # contains all
508
- tm .assert_frame_equal (df , g_not_as .tail (7 ))
509
-
510
- # as_index=True, (used to be different)
511
- df_as = df
512
-
513
- tm .assert_frame_equal (df_as .loc [[0 , 2 ]], g_as .head (1 ))
514
- tm .assert_frame_equal (df_as .loc [[1 , 2 ]], g_as .tail (1 ))
515
-
516
- empty_as = DataFrame (index = df_as .index [:0 ], columns = df .columns )
517
- empty_as ["A" ] = empty_not_as ["A" ].astype (df .A .dtype )
518
- empty_as ["B" ] = empty_not_as ["B" ].astype (df .B .dtype )
519
- tm .assert_frame_equal (empty_as , g_as .head (0 ))
520
- tm .assert_frame_equal (empty_as , g_as .tail (0 ))
521
- tm .assert_frame_equal (empty_as , g_as .head (- 1 ))
522
- tm .assert_frame_equal (empty_as , g_as .tail (- 1 ))
523
-
524
- tm .assert_frame_equal (df_as , g_as .head (7 )) # contains all
525
- tm .assert_frame_equal (df_as , g_as .tail (7 ))
526
-
527
- # test with selection
528
- tm .assert_frame_equal (g_as [[]].head (1 ), df_as .loc [[0 , 2 ], []])
529
- tm .assert_frame_equal (g_as [["A" ]].head (1 ), df_as .loc [[0 , 2 ], ["A" ]])
530
- tm .assert_frame_equal (g_as [["B" ]].head (1 ), df_as .loc [[0 , 2 ], ["B" ]])
531
- tm .assert_frame_equal (g_as [["A" , "B" ]].head (1 ), df_as .loc [[0 , 2 ]])
532
-
533
- tm .assert_frame_equal (g_not_as [[]].head (1 ), df_as .loc [[0 , 2 ], []])
534
- tm .assert_frame_equal (g_not_as [["A" ]].head (1 ), df_as .loc [[0 , 2 ], ["A" ]])
535
- tm .assert_frame_equal (g_not_as [["B" ]].head (1 ), df_as .loc [[0 , 2 ], ["B" ]])
536
- tm .assert_frame_equal (g_not_as [["A" , "B" ]].head (1 ), df_as .loc [[0 , 2 ]])
507
+ g = df .groupby ("A" , as_index = as_index )
508
+ expected = df .iloc [expected_rows ]
509
+ if columns is not None :
510
+ g = g [columns ]
511
+ expected = expected [columns ]
512
+ result = getattr (g , op )(n )
513
+ tm .assert_frame_equal (result , expected )
537
514
538
515
539
516
def test_group_selection_cache ():
0 commit comments