如何根据DataFrame的rule_id列动态调用对应规则函数生成新列？

现有简化版DataFrame：

|    | amount  | other_amt |  rule_id |
|---:|:--------|:----------|---------:|
|  0 | 2       | 0         |      101 |
|  1 | 20      | 0.5       |      102 |
|  2 | 300     | 0         |        0 |
|  3 | 50      | 1         |      101 |

已定义规则函数：

def rule_101(df):
    return df['amount'] / 2

def rule_102(df):
    return df['other_amt']

需求：生成新列new_col，根据每行rule_id的值调用对应的rule_xxx(df)函数。

尝试代码（存在错误）：

df['new_col'] = np.where(df['rule_id'] == '0',
                         df['amount']),
                         locals()[f'rule_{df.rule_id}'](df)) 

错误信息：

KeyError: 'rule_0      0\n1      0\n2      0\n3      0\n4      0\n      ..\n495    0\n496    0\n497    0\n498    0\n499    0\nName: rule_id, Length: 500, dtype: object'

错误原因：f'rule_{df.rule_id}'会将整个rule_id Series转为字符串，无法匹配到对应的函数名。

可行解决方案 #

方案1：使用apply逐行处理（灵活适配现有函数） #

先将规则函数映射到字典中，避免直接使用locals()带来的风险，再通过apply逐行调用：

import pandas as pd

# 定义规则映射，包含rule_id=0的处理逻辑
rule_map = {
    101: rule_101,
    102: rule_102,
    0: lambda df: df['amount']
}

def process_row(row):
    # 获取当前行对应的规则函数，默认返回amount
    func = rule_map.get(row['rule_id'], lambda df: df['amount'])
    # 将单行转为DataFrame，适配现有规则函数的参数要求
    return func(pd.DataFrame([row]))

df['new_col'] = df.apply(process_row, axis=1)

如果允许修改规则函数，改为接收Series参数，效率会更高：

# 修改规则函数为接收单行Series
def rule_101(row):
    return row['amount'] / 2

def rule_102(row):
    return row['other_amt']

rule_map = {
    101: rule_101,
    102: rule_102,
    0: lambda row: row['amount']
}

df['new_col'] = df.apply(lambda row: rule_map[row['rule_id']](row), axis=1)

方案2：按rule_id分组处理（大数据量推荐） #

分组处理比逐行apply效率更高，适合数据量较大的场景：

rule_map = {
    101: rule_101,
    102: rule_102,
    0: lambda df: df['amount']
}

# 初始化新列
df['new_col'] = 0

# 遍历每个rule_id对应的分组，应用规则
for rule_id, group in df.groupby('rule_id'):
    func = rule_map.get(rule_id, lambda df: df['amount'])
    df.loc[group.index, 'new_col'] = func(group)

方案3：使用np.select实现向量运算（规则数量少的场景） #

如果规则数量不多，直接用np.select做向量运算，效率最高：

import numpy as np

# 定义匹配条件和对应结果
conditions = [
    df['rule_id'] == 101,
    df['rule_id'] == 102,
    df['rule_id'] == 0
]
choices = [
    rule_101(df),
    rule_102(df),
    df['amount']
]

# 生成新列，默认返回amount
df['new_col'] = np.select(conditions, choices, default=df['amount'])

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