Pandas: добавление столбца к фрейму данных на основе существующих данных в качестве связанной операции

Question

Я создаю фрейм данных:

import pandas as pd
df = pd.DataFrame({'FOO': [0,1,2], 'BAR': ['a','b','c']})

    FOO BAR
0   0   a
1   1   b
2   2   c

Теперь я отфильтрую строки и добавлю столбец на основе значений существующего столбца:

df = df[lambda x: x['FOO']>0]
df['BAZ'] = df['BAR'].map(lambda x: x+'z')

    FOO BAR BAZ
1   1   b   bz
2   2   c   cz

Можно ли выполнять фильтрацию и добавление столбца в одну строку способом цепочки операций? Я могу сделать это:

df = df[lambda x: x['FOO']>0].join(df[lambda x: x['FOO']>0]['BAR'].map(lambda x: x+'z').rename('BAZ'))

Но это нежелательно, так как я должен повторить фильтрацию.

Answer 1

Вы можете assign после фильтра с lambda:

df.loc[df.FOO.gt(0)].assign(BAZ = lambda x: x.BAR+'z')

#   FOO BAR BAZ
#1    1   b  bz
#2    2   c  cz

Также возможно назначить сначала, а затем фильтр, но это не будет столь же эффективно, если вы отфильтруете много строк:

df.assign(BAZ = df.BAR+'z').loc[df.FOO.gt(0)]