Python: имя нового столбца в функции с assign ()

Question

В следующем коде я создал функцию для возврата DataFrame с новым столбцом new_col со сдвинутыми на один период значениями ori_col.Но новый столбец вывода имеет имя new_col, а не B.Что мне здесь не хватает?

import pandas as pd

def add_col_diff(dframe, ori_col, new_col):
    shift_1 = dframe[ori_col].shift(1)
    return dframe.assign(new_col=shift_1)

data = {'A': [10, 11, 12]}
df = pd.DataFrame.from_dict(data)
df = add_col_diff(df, 'A', 'B')
df

Это дает

     A  new_col
0   10  NaN
1   11  10.0
2   12  11.0

Answer 1

assign установит ключевые аргументы в качестве имен столбцов.Просто наличие assign(new_col=shift_l) сделает столбец с именем new_col. Аргументы ключевых слов не оцениваются как переменные.

Если вы хотите использовать переменные для определения имен столбцов, вы можете распаковать их из словаря .Попробуйте:

dframe.assign(**{new_col: shift_1})

---

Альтернативой является использование:

dframe[new_col] = shift_1