如何在海量数据帧上提高 lambda 函数的性能

Question

我有一个 df 有数亿行。

     latitude                    longitude                   time                    VAL     
0   -39.20000076293945312500    140.80000305175781250000    1972-01-19 13:00:00     1.20000004768371582031  
1   -39.20000076293945312500    140.80000305175781250000    1972-01-20 13:00:00     0.89999997615814208984 
2   -39.20000076293945312500    140.80000305175781250000    1972-01-21 13:00:00     1.50000000000000000000 
3   -39.20000076293945312500    140.80000305175781250000    1972-01-22 13:00:00     1.60000002384185791016 
4   -39.20000076293945312500    140.80000305175781250000    1972-01-23 13:00:00     1.20000004768371582031
... ...

它包含一个 time 列，其类型为 UTC 中的 datetime64。以下代码是创建一个新列isInDST 来指示time 是否处于本地时区的夏令时。

df['isInDST'] = pd.DatetimeIndex(df['time']).tz_localize('UTC').tz_convert('Australia/Victoria').map(lambda x : x.dst().total_seconds()!=0)

处理 15,223,160 行大约需要 400 秒。

有没有更好的方法以更好的性能实现这一目标？ vectorize 是不是更好？

Answer 1

所有结果均基于 1M 数据点计算。

赛通 + np.vectorize

比原代码快 7.2 倍

%%cython
from cpython.datetime cimport datetime
cpdef bint c_is_in_dst(datetime dt):
    return dt.dst().total_seconds() != 0 

%%timeit
df['isInDST'] = np.vectorize(c_is_in_dst)(df['time'].dt.tz_localize('UTC').dt.tz_convert('Australia/Victoria').dt.to_pydatetime())

1.08 s ± 10.2 ms per loop per loop

np.vectorize

比原始代码快 6.5 倍

def is_in_dst(dt):
    return dt.dst().total_seconds() != 0 

%%timeit
df['isInDST'] = np.vectorize(is_in_dst)(df['time'].dt.tz_localize('UTC').dt.tz_convert('Australia/Victoria').dt.to_pydatetime())

1.2 s ± 29.3 ms 每循环每循环

基于documentation (The implementation is essentially a for loop)，我预计结果与列表理解的结果相同，但它始终比列表理解好一点。

列表理解

比原始代码快 5.9 倍

%%timeit
df['isInDST'] = [x.dst().total_seconds()!=0 for x in pd.DatetimeIndex(df['time']).tz_localize('UTC').tz_convert('Australia/Victoria')]

每个循环 1.33 秒 ± 48.4 毫秒

这个结果表明 pandas map/apply 非常慢，它增加了额外的开销，这可以通过使用 python for 循环来消除。

原始方法（map on pandas DatetimeIndex）

%%timeit
df['isInDST'] = pd.DatetimeIndex(df['time']).tz_localize('UTC').tz_convert('Australia/Victoria').map(lambda x : x.dst().total_seconds()!=0)

每个循环 7.82 秒 ± 84.3 毫秒

在 100 万行虚拟数据上测试

N = 1_000_000
df = pd.DataFrame({"time": [datetime.datetime.now().replace(hour=random.randint(0,23),minute=random.randint(0,59)) for _ in range(N)]})

另外，在 100K 和 10M 行上运行代码 - 结果线性依赖于行数