在np.where()中传入自定义函数实现条件输出的问题（Pandas、Python、Numpy）

问题现象 #

运行现有代码后，结果DataFrame保留了多余的SAS列，且SAS Program Language列中LOB=BNK的行显示None，无法得到预期的拼接SAS语句；LOB=XYZ的行可正常显示"Pending XYZ Logic"。

问题原因 #

1. logic_for_bnk()函数直接修改全局DataFrame，但无返回值，因此np.where调用它时会得到None，导致LOB=BNK的行被填充为None。
2. 函数内部新增了SAS列，后续删除操作未包含该列，最终结果保留了此列。

解决方案 #

步骤1：重构logic_for_bnk()函数 #

让函数接收数据子集（仅LOB=BNK的行），计算后返回对应的SAS语句序列，而非直接修改全局DataFrame：

import pandas as pd
import numpy as np

def logic_for_bnk(sub_df):
    # 国家和州标识逻辑
    country_state_flag_conditions = [
        (sub_df['US_ALL'] == True)  & (sub_df['US_50_States'] == True),
        (sub_df['US_ALL'] == False) & (sub_df['US_50_States'] == False),
        (sub_df['US_ALL'] == True)  & (sub_df['US_50_States'] == False),
        (sub_df['US_ALL'] == False) & (sub_df['US_50_States'] == True),
    ]
    country_state_flag_values = [
        """%keep(criteria="country = 'US' and states_50 = 1", desc="Keep only US and in 50 states customers");""",
        "",
        """%keep(criteria="country = 'US'",desc="Keep customers in the US");""",
        """%keep(criteria="states_50 = 1", desc="Keep customers in 50 states");"""
    ]
    country_state_logic = np.select(country_state_flag_conditions, country_state_flag_values, "")
    
    # 主副账户逻辑
    primary_secondary_flag_conditions = [
        (sub_df['Primary'] == True)  & (sub_df['Secondary'] == True),
        (sub_df['Primary'] == False) & (sub_df['Secondary'] == False),
        (sub_df['Primary'] == True)  & (sub_df['Secondary'] == False),
        (sub_df['Primary'] == False) & (sub_df['Secondary'] == True)
    ]
    primary_secondary_flag_values = [
        """%keep(criteria="acct_ownership = '1' or acct_ownership = '2'",desc="Keep primary and secondary ownership");""",
        """%keep(criteria="acct_ownership = '1' or acct_ownership = '2'",desc="Keep primary and secondary ownership");""",
        """%keep(criteria="acct_ownership = '1'",desc="Keep primary ownership");""",
        """%keep(criteria="acct_ownership = '2'",desc="Keep secondary ownership");"""
    ]
    primary_secondary_logic = np.select(primary_secondary_flag_conditions, primary_secondary_flag_values, "")
    
    # 拼接并返回结果
    return country_state_logic + primary_secondary_logic

步骤2：生成目标列并清理数据 #

先初始化SAS Program Language列，再分别为不同LOB的行赋值，最后删除冗余列：

# 创建初始DataFrame
data = {'ID':            ['14','15','16','18','19','20','21'],
        'LOB':           ["BNK", "BNK", "BNK", "BNK", "XYZ", "XYZ", "XYZ"],
        'US_ALL':        [False, False, True, True, True, False, True],
        'US_50_States':  [True, False, True, False, True, False, False],
        'Primary':       [False, True, True, False, True, False, True],
        'Secondary':     [True, False, True, False, False, True, True]}
df = pd.DataFrame(data)

# 初始化目标列
df['SAS Program Language'] = ""

# 为LOB=BNK的行赋值
bnk_mask = df['LOB'] == "BNK"
df.loc[bnk_mask, 'SAS Program Language'] = logic_for_bnk(df[bnk_mask])

# 为LOB=XYZ的行赋值
xyz_mask = df['LOB'] == "XYZ"
df.loc[xyz_mask, 'SAS Program Language'] = "Pending XYZ Logic"

# 删除不需要的列
df.drop(['US_ALL','US_50_States','Primary', 'Secondary'], axis=1, inplace=True)

# 查看结果
print(df)

最终效果 #

输出的DataFrame仅保留ID、LOB、SAS Program Language三列：

• LOB=BNK的行显示拼接后的SAS语句
• LOB=XYZ的行显示"Pending XYZ Logic"

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