使用groupby和agg(sum)时出现KeyError: 'sum'错误

问题 #

我想把DataFrame按sum值降序排序，用了这段代码：

new_df = df.groupby(
  ['Category'])[['Total']].agg(
  ['sum', 'mean', 'count']).sort_values(
  'sum', ascending=False).copy()

运行后抛出KeyError：

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/tmp/ipykernel_557/4086857446.py in <module>
      4 
      5 sector.head()
----> 6 sector = data.sort_values('sum', ascending=False)

/usr/local/lib/python3.7/site-packages/pandas/util/_decorators.py in wrapper(*args, **kwargs)
    309                     stacklevel=stacklevel,
    310                 )
--> 311             return func(*args, **kwargs)
    312 
    313         return wrapper

/usr/local/lib/python3.7/site-packages/pandas/core/frame.py in sort_values(self, by, axis, ascending, inplace, kind, na_position, ignore_index, key)
   6257 
   6258             by = by[0]
--> 6259             k = self._get_label_or_level_values(by, axis=axis)
   6260 
   6261             # need to rewrap column in Series to apply key function

/usr/local/lib/python3.7/site-packages/pandas/core/generic.py in _get_label_or_level_values(self, key, axis)
   1777             values = self.axes[axis].get_level_values(key)._values
   1778         else:
--> 1779             raise KeyError(key)
   1780 
   1781         # Check for duplicates

KeyError: 'sum'

错误原因 #

你用agg(['sum', 'mean', 'count'])之后，生成的DataFrame是多级列结构，列名是元组形式：('Total', 'sum')、('Total', 'mean')、('Total', 'count')。直接用'sum'作为排序的列名，pandas找不到这个单独的列，所以抛出KeyError。

解决方法 #

有两种常用的修正方式：

方式一：指定完整的多级列名 #

直接把排序的列名写成元组形式，对应多级结构：

new_df = df.groupby(
  ['Category'])[['Total']].agg(
  ['sum', 'mean', 'count']).sort_values(
  ('Total', 'sum'), ascending=False).copy()

方式二：扁平化多级列名 #

把多级列合并成单个名称，后续操作更方便：

# 先聚合，再扁平化列名
temp_df = df.groupby(['Category'])[['Total']].agg(['sum', 'mean', 'count'])
temp_df.columns = ['_'.join(col) for col in temp_df.columns]
# 现在可以直接用'Total_sum'排序
new_df = temp_df.sort_values('Total_sum', ascending=False).copy()

或者一步到位：

new_df = (df.groupby(['Category'])[['Total']]
          .agg(['sum', 'mean', 'count'])
          .pipe(lambda x: x.set_axis(['Total_sum', 'Total_mean', 'Total_count'], axis=1))
          .sort_values('Total_sum', ascending=False)
          .copy())

内容的提问来源于stack exchange，提问作者anakula