Pandas Groupby Aggregate Return Customized Output (not one row)

Viewed 138

I have a dataframe in the following format

| col1 | col2  | 
|  1   |  day1 |
|  1   |  day2 |
|  1   |  day3 |
|  2   |  day1 | 
|  2   |  day3 | 

I have a pretty customized logic/function that works on a single dataframe for one value of col1 that I will like to apply to all values. It is sort of the groupby. However, I seem not to get the aggregate part right, as to me it looks like only return 1 value, (like max/min/count/..). What if my aggregate function return a dataframe and have more than 1 rows?

For example, the expected output could be

| col1 | col2       | 
|  1   |  day1-day2 |
|  1   |  day2-day3 |
|  2   |  day1-day3 |

As you can tell, there are two rows generated out of group1 and one row out of group2. And the aggregation logic is every two consecutive rows concatenation, or even more complex. It is sort of the map reduce idea in Spark/Hadoop but couldn't get it working in group.aggregate...

Update:

people usually do groupby().agg(sum), it works as sum return 1 row or 1 number. However, I have a function that return a dataframe which could be 0,1 or many rows, it reduce the number of rows for sure but not quite as collapsing to 1 row yet, like groupby().agg(func), is it possible to do this in groupby.agg?

def func(xdf):
    res = []
    for i in range(len(xdf)-1):
        res.append(xdf.iloc[i] + '-' + xdf.iloc[i+1])
    return pd.DataFrame(res) # return a dataframe, not a number, not a row.
2 Answers

Instead of using groupby.agg, you can use groupby.apply, like this with your data and function func, you get

print (df.groupby('col1').apply(func))
             col2
col1             
1    0  day1-day2
     1  day2-day3
2    0  day1-day3

and for getting the expected output format, you can use reset_index as well

print (df.groupby('col1').apply(func)
         .reset_index(level=0)
         .reset_index(drop=True))
   col1       col2
0     1  day1-day2
1     1  day2-day3
2     2  day1-day3

​but in your real case, you might not need it

I think this gets you to where you want.

Data:

df = pd.DataFrame({'col1': [1, 1, 1, 2, 2],
'col2': ['  day1 ', '  day2 ', '  day3 ', '  day1 ', '  day3 '],
'col3': ['  day1 ', '  day2 ', '  day3 ', '  day1 ', '  day3 ']})

Groupby each column

gb = df.groupby(['col1', 'col2']).nth(0)

Then groupby again and dropna's

gb.groupby(['col1']).shift(-1).dropna()

Which yields

col1    col2    col3
1     day1    day2 
1     day2    day3 
2     day1    day3 
Related