Using groupby group names in function

Question

I have data something like:

df = pd.DataFrame({'user': np.random.choice(['a', 'b','c'], size=100, replace=True),
                   'value1': np.random.randint(10, size=100),
                   'value2': np.random.randint(20, size=100)})

I'm using it to produce some results, e.g.,

grouped = df.groupby('user')
results = pd.DataFrame()
results['value2_sum'] = grouped['value2'].sum()

For one of he columns of this result dataframe, I'd like to pass the user names to a different function, which uses data outside of the dataframe.

I tried something like:

results['user_result'] = grouped.apply(lambda x: my_func(x.index))

But couldn't figure out a syntax that worked.

Answer 1

results['user_result'] = results.index.values

To pass the index value to your function, you can use a list comprehension.

def my_func(val):
    return val + "_" + val

results['my_func'] = [my_func(idx) for idx in results.index]

>>> results
      value2_sum user_result my_func
user                                
a            417           a     a_a
b            306           b     b_b
c            331           c     c_c

Answer 2

You want the .name attribute to access a groups index value:

In [6]:
grouped = df.groupby('user')
results = pd.DataFrame()
results['value2_sum'] = grouped['value2'].sum()
results['user_result'] = grouped.apply(lambda x: x.name)
results

Out[6]:
      value2_sum user_result
user                        
a            342           a
b            333           b
c            308           c