在pandas apply循环中途保存DataFrame的技术方案问询

问题原因 #

pandas的apply方法会先完整计算出整个Series，再将其赋值给DataFrame的新列。一旦计算中途抛出异常，整个Series都不会被创建，自然也不会添加到原DataFrame中，所以保存的文件里只有原始的x列。

解决方案：逐元素处理并实时赋值 #

放弃一次性用apply生成整列，改为逐个处理每个元素，每处理完一个就立即赋值到DataFrame的y列中。这样中途出现异常时，已经计算完成的元素值已经存在于DataFrame里，保存时就能保留这些已计算的数据。

改进后的代码示例 #

import pandas as pd
from math import log

data = pd.DataFrame()
data['x'] = [1,2,-1,4,5]
# 先初始化y列，避免后续赋值时出现KeyError
data['y'] = None

try:
    for idx, val in data['x'].items():
        data.loc[idx, 'y'] = log(val)
except Exception as e:
    data.to_pickle('data.pkl')
    raise e

验证效果 #

执行上述代码后，读取保存的data.pkl：

data_bis = pd.read_pickle('data.pkl')
print(data_bis)

输出会显示前两行的y列已经计算完成，未处理到的行保留初始的None，x列完整保留：

x         y
0  1  0.000000
1  2  0.693147
2 -1       NaN
3  4       NaN
4  5       NaN

优化建议（针对大任务） #

如果你的数据量极大（对应14小时的处理任务），逐行遍历可能效率稍低，可以考虑分块处理：将DataFrame分成若干块，每处理完一块就将结果赋值回原DataFrame，这样既能保证中途异常时保存已处理块的数据，又能兼顾效率。示例代码如下：

import pandas as pd
from math import log

data = pd.DataFrame()
data['x'] = [1,2,-1,4,5]
data['y'] = None

# 自定义分块大小，可根据数据量调整
chunk_size = 2
try:
    for start in range(0, len(data), chunk_size):
        end = start + chunk_size
        chunk = data.loc[start:end-1, 'x']
        # 对块内数据用apply批量计算，块级实时赋值
        data.loc[start:end-1, 'y'] = chunk.apply(log)
except Exception as e:
    data.to_pickle('data.pkl')
    raise e

这种方式在块内用apply提升计算速度，同时保证块级的实时赋值，平衡了处理效率和容错需求。

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