pandas 使用方法链接重命名列

Question

我有一个数据框并进行了一些特征工程，现在想更改列名。 如果我做一个新的任务，我知道如何改变它们，但我想用方法链来做。我尝试了以下（rename 行）但它不起作用。我怎么能写出来让它工作？

df = pd.DataFrame({'ID':[1,2,2,3,3,3], 'date': ['2021-10-12','2021-10-16','2021-10-15','2021-10-10','2021-10-19','2021-10-01'], 
                   'location':['up','up','down','up','up','down'], 
                   'code':[False, False, False, True, False, False]})

df = (df
     .assign(date = lambda x: pd.to_datetime(x.date))
     .assign(entries_per_ID = lambda x: x.groupby('ID').ID.transform('size'))
     .pivot_table(values=['entries_per_ID'], index=['ID','date','code'],
                   columns=['location'], aggfunc=np.max)
     .reset_index()
     #.rename(columns=lambda x: dict(zip(x.columns, ['_'.join(col).strip() if col[1]!='' else col[0] for col in x.columns.values])))
     )

这里可行，但我不想这样写。

df.columns = ['_'.join(col).strip() if col[1]!='' else col[0] for col in df.columns.values ]

Answer 1

重命名链中的列

在axis=1 旁边使用set_axis：

df.set_axis(['foo', 'bar', 'baz'], axis=1)

使用groupby、pivot、melt等

如果新列依赖于链中的某个较早步骤，请将set_axis 与pipe 结合使用。例如，要将链中的旋转列大写：

• 我们不能直接链接set_axis：

  # does NOT work since df.columns are the original columns, not pivoted columns
  df.pivot(...).set_axis(df.columns.str.upper(), axis=1))
  

• 但我们可以将 pipe 转为 set_axis 的结果：

  # does work since we've piped the pivoted df
  df.pivot(...).pipe(lambda piv: piv.set_axis(piv.columns.str.upper(), axis=1)))
  #                         ^    ^            ^
  

OP 的例子

由于 OP 创建了一个 pivot_table 并希望有条件地折叠那些旋转的 MultiIndex，我们将 pipe 和 pivot_table 放入列表理解中：

(df.assign(date=pd.to_datetime(df.date))
   .assign(entries_per_ID=df.groupby('ID').ID.transform('size'))
   .pivot_table(index=['ID', 'date', 'code'],
                columns='location',
                values='entries_per_ID',
                aggfunc='max')
   .reset_index()
   .pipe(lambda piv: piv.set_axis(['_'.join(col).strip() if col[1] else col[0] for col in piv.columns],
                                  axis=1)))

#    ID        date   code  entries_per_ID_down  entries_per_ID_up
# 0   1  2021-10-12  False                  NaN                1.0
# 1   2  2021-10-15  False                  2.0                NaN
# 2   2  2021-10-16  False                  NaN                2.0
# 3   3  2021-10-01  False                  3.0                NaN
# 4   3  2021-10-10   True                  NaN                3.0
# 5   3  2021-10-19  False                  NaN                3.0