python pandas, DF.groupby().agg(), column reference in agg()
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.
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'}
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