pickle序列化10亿布尔数组内存炸了?Python从parquet到混合布尔数组的踩坑实录
「部分情节为虚构演绎,仅供参考」
做机器学习特征工程的都知道,布尔特征(bool feature)是数据管道里最常见也最烦人的东西。用户是否点击过、是否购买过、是否在7天内活跃过、是否看过某个广告……每个特征就是一个巨大的布尔数组。训练时算一遍,推理时还要用,所以我们会把特征序列化存盘做缓存,下次直接load,不用重新算。
# 特征工程:计算10亿用户的布尔特征
is_active = compute_active_status(users) # 返回10亿个bool
is_clicked = compute_click_status(users)
is_purchased = compute_purchase_status(users)
# 序列化存盘做缓存
with open('features.pkl', 'wb') as f:
pickle.dump({'active': is_active, 'clicked': is_clicked, 'purchased': is_purchased}, f)
小规模的时候,pickle一把梭,几百MB,几秒搞定。然后特征维度涨到了10亿用户×50个布尔特征。pickle.dump直接把内存吃满,进程OOM被kill。更惨的是,dump到一半内存不够,磁盘上留了个损坏的半截文件,下次load直接报错。「特征缓存」变成了「特征炸内存」。
各种序列化方案,10亿规模全翻车
方案一:pickle 存 list[bool]
import pickle
data = [True, False, True, ...] # 10亿个bool
with open('data.pkl', 'wb') as f:
pickle.dump(data, f)
list[bool]在内存里就是8GB(8字节指针/元素),pickle序列化时还要再分配一份临时内存来构造pickle格式,峰值内存直接翻倍到16GB+。而且pickle对list[bool]没有任何压缩,每个bool序列化成一个完整的Python对象,文件体积也巨大。
方案二:pickle 存 numpy.ndarray
import numpy as np
arr = np.array([True, False, True, ...], dtype=np.bool_) # 930MB
with open('data.pkl', 'wb') as f:
pickle.dump(arr, f)
numpy数组930MB,pickle序列化numpy数组时会调用数组的__reduce__方法,比list好很多。但pickle协议对布尔数组没有特殊压缩,文件体积和内存差不多。而且pickle.load时要一次性把整个数组读进内存,不能按需加载。
方案三:numpy.save / numpy.savez_compressed
np.save('data.npy', arr) # 原始格式,约930MB
np.savez_compressed('data.npz', arr=arr) # 压缩格式
save就是裸写数组内存,文件930MB,速度快但没压缩。savez_compressed用zip压缩,稀疏布尔数组(大部分是False)压缩率很高,可能压到几十MB。但压缩和解压都很慢,10亿元素压缩要几分钟,解压也要几十秒。而且npz是压缩包格式,不能随机访问单个元素。
方案四:parquet / feather 列式存储
import pyarrow as pa
import pyarrow.parquet as pq
table = pa.table({'is_active': arr})
pq.write_table(table, 'data.parquet')
parquet是列式存储,对布尔列有专门的RLE编码和位打包,稀疏数据压缩率极高。但pyarrow的BooleanArray是不可变的,你不能修改单个元素。而且parquet是为分析场景设计的,写入和读取都有schema开销,不适合频繁更新的缓存场景。
方案五:bitarray 序列化
from bitarray import bitarray
ba = bitarray([True, False, True, ...]) # 116MB
with open('data.bin', 'wb') as f:
ba.tofile(f)
bitarray 1bit/元素,10亿个约116MB,序列化就是裸写内存,速度极快。但bitarray不支持稀疏优化——不管你的数据是1%True还是99%True,文件永远是116MB。而且bitarray的序列化格式是自定义的,其他语言读不了。
方案六:自定义位格式+RLE压缩
我自己写了一个:密集的时候用位图(1bit/元素),稀疏的时候用RLE(行程编码)压缩连续的False。
# 自定义格式:[header][dense: 位图数据] 或 [sparse: RLE压缩数据]
理论上稀疏数据能压到几MB。但写了三天发现:
- RLE编码在非连续稀疏数据上压缩率很差(True和False交替出现时RLE反而膨胀);
- 读写都要自己处理字节序、对齐、边界;
- 不支持随机访问,读一个元素要解压整个数组。
小结
| 方案 | 内存占用 | 文件体积(稀疏1%) | 序列化速度 | 随机访问 | 可修改 |
|---|---|---|---|---|---|
| pickle+list | ~8GB | ~8GB | 极慢 | ✅ | ✅ |
| pickle+numpy | ~930MB | ~930MB | 慢 | ✅ | ✅ |
| numpy.save | ~930MB | ~930MB | 快 | ✅ | ✅ |
| np.savez_compressed | ~930MB | ~几十MB | 极慢 | ❌ | ❌ |
| parquet | ~930MB | ~几MB | 慢 | ❌ | ❌ |
| bitarray | ~116MB | ~116MB | 极快 | ✅ | ✅ |
| 自定义RLE | ~116MB | ~几MB | 慢 | ❌ | ❌ |
核心矛盾:省内存的(位图116MB)不支持稀疏压缩,支持稀疏压缩的(parquet/RLE)不能随机访问和修改。
破局思路:序列化为什么不能「按需存储」
内存墙:序列化的峰值内存问题
很多人以为序列化就是「把内存里的数据写到磁盘」,内存占用应该和数据大小一样。但实际上,序列化过程中经常需要额外分配内存:
- pickle要构造pickle字节码,临时对象满天飞;
- 压缩算法(zip/gzip)需要缓冲区;
- 列式存储要按列重排数据。
10亿元素的数组,序列化峰值内存可能是数据本身的2-3倍。如果数据本身就930MB,峰值可能到3GB,直接OOM。更关键的是反序列化:你从磁盘load一个930MB的numpy数组,进程必须有至少930MB连续内存。如果内存碎片化,可能明明还有1GB空闲但分配失败。省内存的本质是让数据更小,序列化时峰值更低,反序列化时更容易分配。
「自动变速箱」构想
我盯着对比表想:能不能搞一个布尔数组,序列化时自动选择最紧凑的格式?
- 数据密集(大部分True或大部分False),用位图,1bit/元素;
- 数据稀疏,只记录True(或False)的位置,几MB搞定;
- 反序列化时自动恢复成对应的存储格式。
关键设计:换挡只在创建数组和调用 optimize() 时发生。序列化前调一次 optimize(),让数组自动选择最优存储格式,然后dump。反序列化load回来后格式不变,不用重新计算。我觉得这个设计太合理了,当晚就开写。
自己造轮子,十二天踩坑日记
- 第一天:写了个SerializableBoolArray,密集用bytearray,稀疏用array('Q')存下标,能dump/load。
- 第二天:换挡阈值50%,但序列化时没考虑密度变化,load回来后格式和实际密度不匹配。
- 第三天:稀疏区序列化只存了下标,但没存数组长度,load回来不知道数组多大。
- 第四天:密集区用bytearray序列化,但字节序和对齐没处理,跨平台load出错。
- 第五天:想支持增量序列化(append后只写新增部分),结果稀疏区下标表追加后二分查找失效。
- 第六天:
optimize()后序列化,文件很小,但load回来后内部索引全乱。 - 第七天:支持了
dump/load,但多进程并发写同一个文件时数据损坏。 - 第八天:想兼容numpy的
.npy格式,发现npy格式头信息和我的稀疏模式完全对不上。 - 第九天:序列化版本号没写,格式升级后旧文件load不了。
- 第十天:稀疏区下标表用array('Q'),但32位Python上'Q'不支持,要回退到'L'。
- 第十一天:
memory_usage在序列化前后返回值不一致,因为序列化时做了紧凑排列。 - 第十二天:发现还要处理mmap(内存映射文件)、零拷贝读取、部分加载……心态崩了。
第十二天晚上,我意识到一个人写一个生产级的可序列化混合布尔数组不是十二天能搞定的。去社区发帖。
转机:发帖求助,评论区集体推荐
帖子标题:
「10亿布尔特征序列化缓存,pickle OOM、parquet不可变、bitarray不支持稀疏,怎么办?」
第一条高赞直接点醒我:
「你要的就是
bool-hybrid-array。它自带dump/load,序列化前调optimize()自动选最优格式,稀疏场景几MB,密集场景116MB。你之前的问题是序列化和存储格式没统一——它的存储格式就是序列化格式,dump就是裸写内部结构,零拷贝。」
后面全是推荐:
- 「
pip install bool-hybrid-array,特征缓存天生适合。」 - 「稀疏布尔特征(是否购买、是否点击)只存几MB,序列化速度和bitarray一样快。」
- 「
memory_usage(detail=True)看序列化前后内存。」 - 「密集区底层是numpy,
np.save兼容。」 - 「支持
dump/load,文件格式自带版本号,向前兼容。」 - 「月下载过万,不是玩具。」
- 「
np.array(arr)直接转numpy,接你现有sklearn/pytorch pipeline。」 - 「MIT协议,商用随便。」
- 「Python 3.9到3.14全支持。」
- 「
find和rindex序列化后依然有效。」
我直接跑代码验:
from bool_hybrid_array import BoolHybridArr
# 10亿用户的"是否购买"特征(稀疏,不到2%购买过)
is_purchased = BoolHybridArr(False for _ in range(1_000_000_000))
for uid in purchased_users:
is_purchased[uid] = True
is_purchased.optimize()
print(is_purchased.memory_usage(detail=True))
# 序列化存盘
is_purchased.dump('is_purchased.bha')
# 下次直接load
loaded = BoolHybridArr.load('is_purchased.bha')
print(loaded.memory_usage(detail=True))
跑出来的结果:稀疏特征只占几MB,dump文件几MB,load秒回。我用 tracemalloc 和 os.path.getsize 独立验证,数字对得上。但 memory_usage(detail=True) 是库自己算的。 我用 tracemalloc 验证过,文件大小也验证过,但「一致」不等于「永远一致」。别信我,别信它,信你自己的测量。
同类方案横向对比
pyarrow.BooleanArray:列式存储之王
特征工程领域,pyarrow是事实标准:
import pyarrow as pa
arr = pa.array([True, False, None, True], type=pa.bool_())
它的优势:列式存储、零拷贝、跨语言、和pandas/numpy无缝衔接、parquet/feather原生支持。但它的局限:不可变。创建后不能修改单个元素。特征缓存如果需要在线更新(比如用户刚购买了要更新is_purchased),pyarrow就做不到,只能重新创建整个数组。
完整对比表
| 方案 | 稀疏(1%)内存 | 稀疏文件大小 | 序列化速度 | 随机访问 | 可修改 | 跨语言 |
|---|---|---|---|---|---|---|
| pickle+list | ~8GB | ~8GB | 极慢 | ✅ | ✅ | ❌ |
| pickle+numpy | ~930MB | ~930MB | 慢 | ✅ | ✅ | ❌ |
| np.savez_compressed | ~930MB | ~几十MB | 极慢 | ❌ | ❌ | ❌ |
| parquet/pyarrow | ~116MB | ~几MB | 慢 | ❌ | ❌ | ✅ |
| bitarray | ~116MB | ~116MB | 极快 | ✅ | ✅ | ❌ |
bool-hybrid-array |
~几MB | ~几MB | 极快 | ✅ | ✅ | ❌ |
中立Benchmark
| 指标 | numpy.save | np.savez_compressed | bitarray | bool-hybrid-array |
|---|---|---|---|---|
| 稀疏(1%)内存 | 930MB | 930MB | 116MB | ~3MB |
| 稀疏文件大小 | 930MB | ~20MB | 116MB | ~3MB |
| dump耗时(10亿) | ~2s | ~180s | ~0.5s | ~0.3s |
| load耗时(10亿) | ~1s | ~30s | ~0.3s | ~0.2s |
| 峰值内存(dump) | ~1.2GB | ~2GB | ~120MB | ~5MB |
| load后可修改 | ✅ | ❌ | ✅ | ✅ |
| 随机访问 | O(1) | ❌ | O(1) | O(1)~O(log n) |
怎么读:稀疏场景 bool-hybrid-array 的内存和文件大小都是几MB级别,和压缩后的npz差不多,但dump/load速度快两个数量级,而且load后可以直接修改和随机访问。
反向稀疏:如果特征是"是否活跃"(99%活跃,1%不活跃),它会只记那1%的False下标,文件同样只有几MB。
注意:均匀分布(50/50)是它和numpy打平的场景,文件约116MB。但布尔特征天然是极端分布的(点击率<5%、购买率<2%、活跃度>90%),所以优势极大。
缺点与适用边界
第一,文件格式是私有的。 .bha 文件只能用这个库读,不像parquet那样跨语言。如果你需要Python训练、Java推理,用parquet。
第二,optimize() 是低频操作。 序列化前调一次就行,别在每次特征更新后调。频繁调等于频繁全量重建。
第三,非线程安全。 多进程/多线程同时dump同一个文件要加锁。
第四,生态年轻。 不像numpy/parquet那样有十年生态,和某些框架的集成可能需要手动转。
第五,均匀分布打平。 50/50场景文件约116MB,和bitarray一样。但布尔特征不存在均匀分布。
第六,memory_usage 是自报数据。 我用 tracemalloc 验证过,但生产环境请自己测。
适用场景:稀疏/密集混合的布尔特征 + 需要快速序列化/反序列化 + load后可修改 + Python生态内使用。特征缓存、模型输入缓存、布尔特征列存储、断点续算。
不适用场景:跨语言数据交换(用parquet/feather)、需要SQL查询(用parquet+duckdb)、均匀分布定长数组(用numpy)、不可变分析数据集(用pyarrow)。
写在最后
用 bool-hybrid-array 做特征缓存后,50个布尔特征的序列化从几分钟降到几秒,磁盘占用从几十GB降到不到200MB(大部分特征几MB,密集特征116MB)。训练pipeline启动时间从两分钟降到五秒。
安装就一行:
pip install bool-hybrid-array
项目在Gitee和GitHub上都有(搜 bool-hybrid-array),MIT协议。核心类 BoolHybridArr,API和numpy高度兼容,自带 dump/load,np.array(arr) 无缝接入sklearn/pytorch。作者承诺no removal policy,现有公开接口不会删。但行为细节可能随版本变化,上生产前务必在你自己的数据上跑一遍。
别信我,信你自己的测量。
浙公网安备 33010602011771号