Get all columns after GroupBy Operation Dask/Pandas

Viewed 110

I have a simple dataframe like this:

    id sim  col1    col2    year    extra
0   1   1    a       b      2021    H
1   2   2    a       b      2024    A

I want to groupby this df by col1 & col2 and count the sim that comes in similar grouped columns. so i do this and i get this too..

ddf.groupby(["col1", "col2"]).agg({"sim":"list"}).compute()

                  
col1    col2     sim
a        b       [1, 2]

I also want to keep the row with highest year. So I did this too:

ddf.groupby(["col1", "col2"]).agg({"sim":"list", "year":"max"}).compute()

and got this df:

col1    col2    sim     year    
a         b    [1, 2]   2024

Note: See the max year was 2024 between the grouped cols, so i kept that.

But i lose other columns like id and extra. I also want to get those column based on maximum year value. So , wherever the year is maximum keep that row.

But i loose other columns like id, year and extra. So how can i get this.

Expected Result Df:

 id   col1  col2    sim     year    extra
  2    a      b    [1, 2]   2024      A
1 Answers

Sort the values on year then groupby and aggregate using agg dict: d where sim is aggregated using list while the rest of the columns are aggregated using last

d = {**dict.fromkeys(ddf.columns, 'last'), 'sim': 'list'}
ddf.sort_values('year').groupby(['col1', 'col2']).agg(d).reset_index(drop=True)

Result

   id     sim col1 col2  year extra
0   2  [2, 1]    a    b  2024     A
Related