在函数中使用 return 创建新的数据框列

Question

在下面的DataFrame中，我需要检查shop列，并根据shop添加两列

datetime	dtnum	time	tnum	shop
02-03-2022 05:25	20220302052500	05:25:00	52500	PRESS1
02-03-2022 05:26	20220302052600	05:26:00	52600	BODY1
02-03-2022 05:27	20220302052700	05:27:00	52700	BODY2

我尝试使用以下代码。即使用相同的 return 制作两列

# Apply shop_groupcode
def shop_chk(column):
    if column['shop'] == 'PRESS1' or 'PRESS1_A' or 'E176' or 'E177' or 'E184' or 'E185' or 'E186' or 'E187':
        return 1, 1  # shop code , plant code
    if column['shop'] == 'BODY1' or 'BODY1_A' or 'E179' or 'E180' or 'E181' or 'E208' or \
            'E216' or 'E217' or 'E218' or 'E232':
        return 2, 1 # shop code , plant code
    if column['shop'] == 'BODY2' or 'BODY2_A' or 'E196' or 'E197' or 'E198' or 'E199' or 'E200':
        return 8, 2 # shop code , plant code

df['shop_code', 'plant_code'] = df.apply(shop_chk, axis=1, result_type="expand")

代码运行没有错误，创建了两列 - 但列值都是 1

要求：我想知道代码是否有错误，或者有什么有效的方法，因为我还有一些商店条件需要检查

datetime	dtnum	time	tnum	shop	shop_code	plant_code
02-03-2022 05:25	20220302052500	05:25:00	52500	PRESS1	1	1
02-03-2022 05:26	20220302052600	05:26:00	52600	BODY1	2	1
02-03-2022 05:27	20220302052700	05:27:00	52700	BODY2	8	2

Answer 1

您需要结合使用 isin 和提供值列表：

    if column['shop'].isin(['PRESS1', 'PRESS1_A', 'E176', 'E177', 'E184', 'E185', 'E186', 'E187']):

否则它将始终停留在第一个 if-statement，因为 or 之后的部分始终为 True，因此结果始终为 1。

Answer 2

您可以使用 isin 而不是多个 OR 并将这些条件存储在列表中并使用 numpy.select:

import numpy as np
conditions = [df['shop'].isin(['PRESS1','PRESS1_A','E176','E177','E184','E185','E186','E187']), 
              df['shop'].isin(['BODY1','BODY1_A','E179','E180','E181','E208','E216','E217','E218','E232']),
              df['shop'].isin(['BODY2','BODY2_A','E196','E197','E198','E199','E200'])
             ]
df['shop_code'] = np.select(conditions, [1, 2, 8])
df['plant_code'] = np.select(conditions, [1, 1, 2])

输出：

           datetime           dtnum      time   tnum    shop  shop_code  plant_code  
0  02-03-2022 05:25  20220302052500  05:25:00  52500  PRESS1          1           1  
1  02-03-2022 05:26  20220302052600  05:26:00  52600   BODY1          2           1
2  02-03-2022 05:27  20220302052700  05:27:00  52700   BODY2          8           2

仅供参考，正确的语法是：

(column['shop'] == 'PRESS1') or (column['shop'] == 'PRESS1_A') or ...

因为

column['shop'] == 'PRESS1' or 'PRESS1_A' or ...

returns True 或 'PRESS1_A'（不是 truth-value）。

在函数中使用 return 创建新的数据框列

Create new data frame columns with return in function

python

pandas

dataframe

function

boolean-expression