python pandas, DF.groupby().agg(), column reference in agg()

Question

On a concrete problem, say I have a DataFrame DF

     word  tag count 0    a     S    30 1    the   S    20 2    a     T    60 3    an    T    5 4    the   T    10  

I want to find, for every "word", the "tag" that has the most "count". So the return would be something like

     word  tag count 1    the   S    20 2    a     T    60 3    an    T    5 

I don't care about the count column or if the order/Index is original or messed up. Returning a dictionary {'the' : 'S', ...} is just fine.

I hope I can do

DF.groupby(['word']).agg(lambda x: x['tag'][ x['count'].argmax() ] ) 

but it doesn't work. I can't access column information.

More abstractly, what does the function in agg(function) see as its argument?

btw, is .agg() the same as .aggregate() ?

Many thanks.

Answer 1

agg is the same as aggregate. It's callable is passed the columns (Series objects) of the DataFrame, one at a time.

---

You could use idxmax to collect the index labels of the rows with the maximum count:

idx = df.groupby('word')['count'].idxmax() print(idx) 

yields

word a       2 an      3 the     1 Name: count 

and then use loc to select those rows in the word and tag columns:

print(df.loc[idx, ['word', 'tag']]) 

yields

  word tag 2    a   T 3   an   T 1  the   S 

Note that idxmax returns index labels. df.loc can be used to select rows by label. But if the index is not unique -- that is, if there are rows with duplicate index labels -- then df.loc will select all rows with the labels listed in idx. So be careful that df.index.is_unique is True if you use idxmax with df.loc

---

Alternative, you could use apply. apply's callable is passed a sub-DataFrame which gives you access to all the columns:

import pandas as pd df = pd.DataFrame({'word':'a the a an the'.split(),                    'tag': list('SSTTT'),                    'count': [30, 20, 60, 5, 10]})  print(df.groupby('word').apply(lambda subf: subf['tag'][subf['count'].idxmax()])) 

yields

word a       T an      T the     S 

---

Using idxmax and loc is typically faster than apply, especially for large DataFrames. Using IPython's %timeit:

N = 10000 df = pd.DataFrame({'word':'a the a an the'.split()*N,                    'tag': list('SSTTT')*N,                    'count': [30, 20, 60, 5, 10]*N}) def using_apply(df):     return (df.groupby('word').apply(lambda subf: subf['tag'][subf['count'].idxmax()]))  def using_idxmax_loc(df):     idx = df.groupby('word')['count'].idxmax()     return df.loc[idx, ['word', 'tag']]  In [22]: %timeit using_apply(df) 100 loops, best of 3: 7.68 ms per loop  In [23]: %timeit using_idxmax_loc(df) 100 loops, best of 3: 5.43 ms per loop 

---

If you want a dictionary mapping words to tags, then you could use set_index and to_dict like this:

In [36]: df2 = df.loc[idx, ['word', 'tag']].set_index('word')  In [37]: df2 Out[37]:       tag word     a      T an     T the    S  In [38]: df2.to_dict()['tag'] Out[38]: {'a': 'T', 'an': 'T', 'the': 'S'} 

Answer 2

Here's a simple way to figure out what is being passed (the unutbu) solution then 'applies'!

In [33]: def f(x): ....:     print type(x) ....:     print x ....:       In [34]: df.groupby('word').apply(f) <class 'pandas.core.frame.DataFrame'>   word tag  count 0    a   S     30 2    a   T     60 <class 'pandas.core.frame.DataFrame'>   word tag  count 0    a   S     30 2    a   T     60 <class 'pandas.core.frame.DataFrame'>   word tag  count 3   an   T      5 <class 'pandas.core.frame.DataFrame'>   word tag  count 1  the   S     20 4  the   T     10 

your function just operates (in this case) on a sub-section of the frame with the grouped variable all having the same value (in this cas 'word'), if you are passing a function, then you have to deal with the aggregation of potentially non-string columns; standard functions, like 'sum' do this for you

Automatically does NOT aggregate on the string columns

In [41]: df.groupby('word').sum() Out[41]:        count word        a        90 an        5 the      30 

You ARE aggregating on all columns

In [42]: df.groupby('word').apply(lambda x: x.sum()) Out[42]:          word tag count word                   a         aa  ST    90 an        an   T     5 the   thethe  ST    30 

You can do pretty much anything within the function

In [43]: df.groupby('word').apply(lambda x: x['count'].sum()) Out[43]:  word a       90 an       5 the     30