用pickle格式保存金融tick数据 减少内存占用

前言

降低数据集文件的内存占用是Quant需要解决的一个问题，特别是当你面对着几千只股票、上千个交易日的tick数据的时候。相比于csv文件，使用pickle文件进行存储也许是一个不错的解决方案，本文将探讨这一思路。当然，肯定还有更好的方案，本文就权当抛砖引玉了。

使用csv文件进行保存

作为例子，我这里直接生成了一个10万行的数据，如何获取tick数据的文章后期再补。

import pandas as pd
import os

ls = [[1, 2, 3] for _ in range(100000)]
data = pd.DataFrame(ls, columns=['a', 'b', 'c'])
data.to_csv("test.csv", index=False)
print(f"file size of csv: {os.path.getsize('test.csv')}")

运行结果是

file size of csv: 700007

os.path.getsize()得到的结果是字节，这个csv文件的大小是684KB。

使用pickle文件进行保存

import pandas as pd
import os

ls = [[1, 2, 3] for _ in range(100000)]
data = pd.DataFrame(ls, columns=['a', 'b', 'c'])
data.to_pickle("test_before.pkl")
print(f"file size of pickle_before: {os.path.getsize('test_before.pkl')}")

运行结果是

file size of pickle_before: 2400709

保存的pickle文件的大小是2345KB
对，没看错，就是2345KB，可以试试，就是这么多，此车已翻~
使用pickle文件格式进行存储，不仅没有减少内存占用，反而多占了2、3倍的空间~

优化内存占用

刚刚翻车的原因是csv中文件的存储格式，可以使用下面的代码进行查看

print(data.dtypes)

运行结果是

a    int64
b    int64
c    int64

全是int64，岂能不大！！！
一个int64占用8个字节，所以刚刚保存pickle文件中的2400709字节其实很好理解，约等于$8\ast 3\ast 100000$
问题找到了，解决方案也就有了。这里，实际上使用int8足以，因此

import pandas as pd
import os

ls = [[1, 2, 3] for _ in range(100000)]
data = pd.DataFrame(ls, columns=['a', 'b', 'c'])
data = data.astype('int8')
data.to_pickle("test_after.pkl")
print(f"file size of pickle_after: {os.path.getsize('test_after.pkl')}")

运行结果是

file size of pickle_after: 300707

300707字节约等于$1\ast 3\ast 100000$

总结

本文对比了使用csv文件格式、pickle文件格式和优化数据类型后的pickle文件格式分别占用的内存大小，结果汇总如下

file size of csv: 700007
file size of pickle_before: 2400709
file size of pickle_after: 300707

从结果中可以看出，文件大小从之前的700007下降到了300707，下降了一倍有余，而结果并不会受影响。

import pickle
with open("test_after.pkl", "rb") as f:
    data = pickle.load(f)
print(data)

运行结果是

       a  b  c
0      1  2  3
1      1  2  3
2      1  2  3
3      1  2  3
4      1  2  3
...   .. .. ..
99995  1  2  3
99996  1  2  3
99997  1  2  3
99998  1  2  3
99999  1  2  3

[100000 rows x 3 columns]