Applying a dataframe function to a pandas groupby object

Question

I am trying to apply a function to each group in a pandas dataframe where the function requires access to the entire group (as opposed to just one row). For this I am iterating over each group in the groupby object. Is this the best way to achieve this?

import pandas as pd
df = pd.DataFrame({'id': [1,1,1,1,2,2,2], 
                   'value': [70,10,20,100,50,5,33], 
                   'other_value': [2.3, 3.3, 7.4, 1.1, 5, 10.3, 12]})
def clean_df(df, v_col, other_col):
    '''This function is just a made up example and might 
       get more complex in real life. ;)
    '''
    prev_points = df[v_col].shift(1)
    next_points = df[v_col].shift(-1)
    return df[(prev_points > 50) | (next_points < 20)]  
grouped = df.groupby('id')
pd.concat([clean_df(group, 'value', 'other_value') for _, group in grouped])

The original dataframe is

    id  other_value value
0   1   2.3         70
1   1   3.3         10
2   1   7.4         20
3   1   1.1         100
4   2   5.0         50
5   2   10.3        5
6   2   12.0        33

The code will reduce it to

    id  other_value value
0   1   2.3         70
1   1   3.3         10
4   2   5.0         50

Answer 1

You can directly use apply on the grouped dataframe and it will be passed the whole group:

def clean_df(df, v_col='value', other_col='other_value'):
    '''This function is just a made up example and might 
       get more complex in real life. ;)
    '''
    prev_points = df[v_col].shift(1)
    next_points = df[v_col].shift(-1)
    return df[(prev_points > 50) | (next_points < 20)]  

df.groupby('id').apply(clean_df).reset_index(level=0, drop=True)
#    id  other_value  value
# 0   1          2.3     70
# 1   1          3.3     10
# 4   2          5.0     50

Note that I had to give the other arguments default values, since the function that is applied needs to have only one argument. Another way around this is to make a function that returns the function:

def clean_df(v_col, other_col):
    '''This function is just a made up example and might 
       get more complex in real life. ;)
    '''
    def wrapper(df):
        prev_points = df[v_col].shift(1)
        next_points = df[v_col].shift(-1)
        return df[(prev_points > 50) | (next_points < 20)]  
    return wrapper

Which you can use like this:

df.groupby('id').apply(clean_df('value', 'other_value')).reset_index(level=0, drop=True)

Or you can use functools.partial with your clean_df:

from functools import partial

df.groupby('id') \
  .apply(partial(clean_df, v_col='value', other_col='other_value')) \
  .reset_index(level=0, drop=True)