Python中调用自定义函数生成DataFrame新列的问题排查

示例数据 #

首先定义测试用DataFrame：

df = pd.DataFrame({'A': ['foo', 'bar', 'baz', 'foo'], 
                   'B': ['qux', 'quux', 'quuz', 'xyz']})

正确实现的基础方法 #

以下代码可正常生成符合预期的新列：

# 直接向量化运算生成col2
df['col2'] = (df['A'] == 'foo') & (df['B'] == 'qux' )
# 使用assign结合lambda生成col3
df = df.assign(col3 = lambda x: (df['A'] == 'foo') & (df['B'] == 'qux' ))

错误UDF的问题分析 #

为处理更复杂逻辑编写的test函数，调用后生成的col1、col4、col5全部为True，完全不符合预期：

def test(df):
    if ((df['A'] == 'foo') & (df['B'] == 'qux' )).any():
        x = True
    else:
        x = False
    return x 

# 三种错误调用方式
df['col1'] = test(df)               
df = df.assign(col4 = lambda x: test(df))
df = df.assign(col5 = df.apply(lambda x: test(df), axis=1))

错误原因 #

1. 函数逻辑错误：test函数接收整个DataFrame作为参数，(df['A'] == 'foo') & (df['B'] == 'qux')生成的是布尔Series，.any()会检查整个Series中是否存在至少一个True，只要有一行满足条件就返回True，最终函数只会返回单个布尔值，赋值给列时会把这个值广播到所有行。
2. 调用方式错误：
   • 直接调用test(df)返回单个值，赋值给列时所有行都被填充这个值。
   • assign中的lambda参数x是当前DataFrame，但调用test(df)还是传入整个原始DataFrame，结果同样是单个值广播。
   • df.apply(lambda x: test(df), axis=1)中，lambda的参数x是每行数据，但调用test(df)仍然传入整个DataFrame，所以每次循环都返回同一个全局的True，导致所有行结果一致。

正确的逐行处理UDF #

后续编写的test3函数可以正常工作，因为它接收单行数据作为参数，针对每行单独计算：

def test3(row):
    x = True
    if ((row['A'] == 'foo') & (row['B'] == 'qux')):
        x = False
    elif ((row['A'] == 'foo') & (row['B'] == 'xyz')):
        x = 'NA'
    return x  

df['col'] = df.apply(lambda x: test3(x), axis=1)

大型DataFrame的高效实现方案 #

df.apply(axis=1)本质是逐行循环，在DataFrame规模较大时效率极低，推荐使用以下高效替代方案：

1. 向量化运算 #

优先使用Pandas的向量化操作，避免逐行循环：

# 使用numpy.where实现多条件判断
import numpy as np

df['col_vectorized'] = np.where(
    (df['A'] == 'foo') & (df['B'] == 'qux'),
    False,
    np.where(
        (df['A'] == 'foo') & (df['B'] == 'xyz'),
        'NA',
        True
    )
)

2. 使用np.select处理多分支条件 #

当条件分支较多时，np.select结构更清晰：

conditions = [
    (df['A'] == 'foo') & (df['B'] == 'qux'),
    (df['A'] == 'foo') & (df['B'] == 'xyz')
]
choices = [False, 'NA']

df['col_select'] = np.select(conditions, choices, default=True)

3. 使用swifter加速apply #

如果必须使用自定义函数（比如逻辑极度复杂无法向量化），可以用swifter自动选择最优执行方式（向量化或并行处理）：

import swifter

df['col_swifter'] = df.swifter.apply(test3, axis=1)

4. 自定义向量化函数 #

直接编写接收Series的函数，利用numpy进行批量运算：

def test_vectorized(A, B):
    res = np.full(len(A), True)
    mask1 = (A == 'foo') & (B == 'qux')
    res[mask1] = False
    mask2 = (A == 'foo') & (B == 'xyz')
    res[mask2] = 'NA'
    return res

df['col_udf_vectorized'] = test_vectorized(df['A'], df['B'])

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