因子平台的机器只有 908MB 内存。我在 app/services/expr.py 里留过一段话:站内唯一的 N² 算子(广义秩)整行一次算,要同时存活 diff、absd、sign·absd 三份 (N,N) 的 float64,全市场 N≈5200 时约 650MB,被内核 OOM 掉过。后来按列分块,峰值降到 32MB。修是修好了,但宽表这条路什么时候会撞墙、撞在哪一步,我一直没量过。
所以我把同一批截面因子计算,用三种写法各跑一遍,看耗时和内存分别是什么量级。
数据和题目
本机 A 股日线库导成一份 Parquet,0.26 GiB。A 股 5225 只、1668 万行、8778 个交易日;文件里还留着三个加密货币标的(btc_usdt、eth_usdt、pepe_usdt),一共 5228 个代码。四个工作负载是我平台上真实会做的四类操作:
- mom20:20 日动量,close 除以自己的滞后 20 期
- 截面 z-score:日收益率按当天所有股票标准化
- roll_std20:日收益率的 20 期滚动标准差,样本标准差
- composite:上面三个百分位排名相加,每天取分数最高的 50 只,再求这 50 只的均值
三种写法:numpy 走宽表(透视成 date×code 矩阵,滚动窗口用 cumsum 做);polars 走长表表达式加 over 分组;duckdb 走 SQL 窗口函数。数值口径我尽量对齐了(min_samples=20、ddof、NULL 处理),float 精度我用 float32,因为 float64 在这台机器上装不下。
先踩的坑:日历被撑长了
第一版跑出来,numpy 的 roll_std20 只有 1869 个有效值,另外两个引擎是 205 万。我以为自己 cumsum 写错了,查了半天,问题是行情库里除了 A 股还存了三个加密货币标的。它们的日历是 7×24,把 2025 年以来的日期轴从 426 天撑到 645 天。A 股那些列在周末位置上全是 NaN,20 格滑动窗口就永远凑不满 20 个有效样本。
三个币圈标的一共一万行出头,占全库不到千分之一,却把宽表的时间轴改了样子。长表路线(polars、duckdb)不受影响,因为它们的窗口是按"这只股票自己的前 20 根"算的,跟日历无关。
这是我第一次真切感受到宽表对日历多敏感。加了一句只留 A 股代码前缀的过滤,三个引擎的日期数都是 426 了。
843 万行(2019 年起,1882 个交易日)
| 引擎 | 载入 | 透视 | mom20 | 截面z | 滚动std20 | 复合+Top50 | 计算合计 | 峰值内存 |
|---|---|---|---|---|---|---|---|---|
| numpy 宽表 | 8.05s | 0.24s | 0.07s | 0.27s | 1.07s | 2.87s | 4.28s | 2012 MB |
| polars | 0.17s | — | 0.39s | 1.25s | 0.84s | 6.14s | 8.62s | 3041 MB |
| duckdb | 0.19s | — | 1.16s | 2.58s | 3.41s | 11.22s | 18.37s | 929 MB |
算得最快的是 numpy,四个负载 4.28 秒,比 polars 快一倍、比 duckdb 快四倍。但端到端加进去读数据,顺序就反过来了:numpy 12.6 秒、polars 8.8 秒、duckdb 18.6 秒。因为 pandas 读那份 Parquet 花了 8 秒,polars 只要 0.17 秒。
也就是说宽表路线真正的成本不在矩阵运算上,在把长表摊成矩阵之前那次 DataFrame 落地。我在 duckwide.py 里为了压这段已经写过一次教训:列名没显式投影时,DuckDB 不把列裁剪下推到读 parquet,全区间装载 38.9 秒,显式投影 1.19 秒。同一个地方,第二次撞上同一类问题。
内存那一列反过来:duckdb 只用 929MB,是 numpy 的一半、polars 的三分之一。
三边结果对不对得上
polars 和 duckdb 是两套独立写法,三个输出的非空数和求和逐位相同(mom20 都是 8,320,620 个、和 79,913.8)。numpy 差一点:
| 输出 | numpy 非空数 | 另两家 | 差 |
|---|---|---|---|
| mom20 | 8,308,821 | 8,320,620 | 0.14% |
| 截面 z | 8,416,746 | 8,419,793 | 0.04% |
| 滚动 std20 | 8,260,933 | 8,320,620 | 0.72% |
| Top50 日均值合计 | 4,664.8(1862 天) | 4,652.5 / 4,668.4 | 0.34% |
差异来源是 lag 的定义。宽表里"往前 20 格"是 20 个日历位置,这只股票要是在中间停过牌,那一格是 NaN,因子就没了;长表里 PARTITION BY code ORDER BY date 的"往前 20 行"是这只股票自己的前 20 根 K 线,停牌的日子直接被跳过。窗口越长、停牌越多,差得越明显,所以滚动 std20 的偏差最大。
这条我没去改成完全一致,因为它不是 bug,是宽表的固有性质。要一致就得在矩阵里再存一张"该股第几根 bar"的下标,那又是另一份内存。我现在把它当成已知偏差记着:凡是宽表算出来的因子,停牌股的样本会少 0.1% 到 0.7%,而且这 0.7% 不是随机丢的,专挑停过牌的股票。
全量 1668 万行
另两个引擎的峰值内存按 843 万行外推要 4 到 6GB,机器当时只剩 1.3GB 可用,我没硬跑。DuckDB 能设内存上限,我给它 1GB 跑完了全量:
mom20 3.97s,截面 z 6.44s,滚动 std20 12.09s,复合加 Top50 31.72s,合计 54.22 秒,约 30.8 万行每秒。超出内存的部分它自己落盘了,我没有额外配置。
这条对生产机最实用:908MB 的机器上,只有这条路能给出一个"内存封顶、慢就慢一点"的保证。numpy 和 polars 都是要么给够内存、要么被系统杀掉。
我以后怎么写
截面排名、分组聚合这类,用 polars 或者 duckdb,别为了矩阵好看先摊成宽表。滚动窗口密集的链路保留宽表,它确实快一倍以上,但要先固定交易日历,并且接受停牌股的样本偏差。生产机上跑不动的大批量,走 duckdb 加内存上限。
还没做的两件事:一是没测 4000 多条因子批量重算的总耗时(那才是平台的真实负载,单因子外推不一定准);二是没测并发,多用户同时点回测时三个引擎的争用行为肯定不一样。
复现
cd D:\Axm\股票助手\cross-bench
python export_parquet.py
python bench.py --from 2019
"D:\Axm\股票助手\gpzs\.venv\Scripts\python.exe" duckdb_worker.py --mem-limit 1GB
四个脚本:numpy_worker.py、polars_worker.py、duckdb_worker.py、bench.py(父进程用 psutil 采样峰值内存)。三个引擎各自起独立子进程跑,表里的秒数是单进程单次的实测值,没取平均。脚本是我和 AI 协作写的,数字都是这台机器实跑出来的输出。