为何PyArrow数组求和速度优于NumPy？适用场景解析

测试代码与结果 #

初始化代码 #

import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
import timeit

# 生成NumPy数组
np_days = np.random.randint(0, 100, 100000000, dtype=np.int8)
np_months = np.random.randint(0, 100, 100000000, dtype=np.int8)
np_years = np.random.randint(0, 100, 100000000, dtype=np.int8)

# 转换为PyArrow数组
days = pa.array(np_days, type=pa.int8())
months = pa.array(np_months, type=pa.int8())
years = pa.array(np_years, type=pa.int8())

循环求和测试 #

NumPy 计算耗时

%%time
for i in range(100):
    np.sum(np_days)
    np.sum(np_months)
    np.sum(np_years)

CPU times: user 12.1 s, sys: 0 ns, total: 12.1 s
Wall time: 12.1 s

PyArrow 计算耗时

%%time
for i in range(100):
    pc.sum(days)
    pc.sum(months)
    pc.sum(years)

CPU times: user 4.51 s, sys: 0 ns, total: 4.51 s
Wall time: 4.51 s

单数组重复求和测试 #

print('pyarrow:', timeit.timeit(lambda: pc.sum(days), number=1000))

pyarrow: 0.02044958801707253

print('NumPy:', timeit.timeit(lambda: np.sum(np_days), number=1000))

NumPy: 0.04707089299336076

问题解答 #

1. 求和速度差异的原因 #

• 内存布局与向量化优化：PyArrow数组采用列存布局，连续的同类型数据能更高效利用CPU缓存和SIMD指令。而NumPy的np.sum处理int8这类小整数类型时，会自动将数据提升为int64再求和，额外的类型转换增加了开销；PyArrow的求和操作可直接针对int8类型做优化，省去了这一步骤。
• 计算路径轻量化：PyArrow compute模块专为列式计算设计，避免了NumPy通用计算逻辑中的额外兼容开销，在单一类型的聚合场景下，执行路径更简洁直接。
• 底层优化策略：PyArrow的计算实现会根据CPU特性自动启用更高效的并行或向量化策略，在小整数聚合这类场景下，比NumPy的优化更具针对性。

2. 优先选择PyArrow的计算场景 #

• 大规模列式数据处理：当数据以列存格式（如Parquet、Feather）存储，或本身就是PyArrow数组时，直接使用PyArrow compute可避免数据转换开销，同时利用列存的缓存友好性提升性能。
• 多列批量聚合：对多列同类型数据执行求和、均值等聚合操作时，PyArrow的列式处理能批量高效执行，整体开销低于逐列用NumPy处理。
• 跨语言/跨系统数据交互：需要在Python、C++、R等语言间传递数据，或与Spark、Flink等大数据系统交互时，PyArrow作为通用内存格式，无需转换为NumPy数组即可计算，端到端效率更高。
• 小宽度数值类型的聚合：针对int8、int16这类小整数类型的聚合计算，PyArrow无需频繁类型提升，性能优势显著。
• 零拷贝数据访问场景：PyArrow支持零拷贝读取多种数据源，计算时直接基于原始内存操作，避免了NumPy数组转换带来的内存拷贝开销。

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