Pandas apply(raw=True)结合返回Optional[int]的函数时的报错问题及行为合理性咨询

你观察到的这个现象完全是预期行为，根源就在于raw=True时Pandas底层依赖的numpy apply_along_axis的工作机制，咱们一步步拆解来看：

为什么会触发TypeError？ #

当你设置raw=True时，Pandas不会把每行数据包装成pd.Series，而是直接传递numpy数组给你的自定义函数。而numpy的apply_along_axis有个关键特性：它会根据函数返回的第一个结果来自动推断最终输出数组的数据类型。

比如你第一个例子里，第一行输入是1，函数返回0（int类型），numpy就会创建一个int64类型的数组来存储后续结果。但None是Python对象，无法被存入numpy的数值类型数组，所以当后续函数返回None时，就会抛出TypeError: int() argument must be a string, a bytes-like object or a real number, not 'NoneType'。

为什么第一个返回值是None就正常？ #

如果函数第一个返回的是None，numpy会推断输出数组的dtype为object——这种类型的数组可以容纳任意Python对象，包括None，所以后续的int值和None都能正常存储，不会报错。

那能不能用raw=True结合返回Optional[int]的函数？ #

当然可以，但需要主动控制数据类型，避免numpy自动推断带来的问题，这里给你两个实用方案：

方案1：用numpy.nan替代None，配合Pandas的Nullable整数类型

把函数里的return None改成return np.nan，然后将结果转换为Pandas的Int64类型（支持空值的整数类型），示例代码：

import pandas as pd
import numpy as np
from typing import Optional

def func(a: int) -> Optional[int]:
    if a % 3 == 0:
        return 1
    if a % 3 == 1:
        return 0
    else:
        return np.nan  # 用nan替代None

df = pd.DataFrame([[1], [2], [3], [4], [5], [6]])
result = df.apply(lambda row: func(row[0]), axis=1, raw=True).astype('Int64')
print(result)

这样既保留了raw=True的性能优势，又能正确处理空值。

方案2：手动指定输出为object类型

如果你坚持要返回None，可以在apply后显式将结果转换为object类型：

result = df.apply(lambda row: func(row[0]), axis=1, raw=True).astype(object)

不过这种方式的性能会略低于方案1，因为object类型数组的操作效率不如数值类型。

针对PySpark pandas_udf的额外建议 #

你提到是在PySpark的@pandas_udf场景下使用，这里额外提一句：如果你的逻辑能用PySpark内置的when/otherwise表达式实现，性能会比自定义pandas_udf更好；如果必须用自定义函数，上面的方案1完全适用，而且Int64类型和PySpark的IntegerType（支持空值）也能很好兼容。

内容来源于stack exchange