求教一个dataframe 分组问题

Question

这样的一个表，想用key1分组，取data2大的那个行，做一个表怎么做？
试过df.groupby('key1',as_index=False).max()取的是所有列的最大，这样不行，请教怎么做？
这是原数据：
df = pd.DataFrame({'key1':['a','a','b','b','a'],'key2':['one','two','one','two','one'],'data1':np.random.randn(5),'data2':np.random.randn(5)})
df

Answer 1

def f(df, col=1):
    return df[df['data2'] == max(df['data2'])]

df1 = df.groupby(['key1']).apply(f)

Answer 2

l = df.groupby('key1')['data2'].max().tolist()
df = df.query('data2 in @l')