一种比较省内存的稀疏矩阵Python存储方案

原文：http://www.pythontab.com/html/2014/pythonhexinbiancheng_0522/788.html

推荐系统中经常需要处理类似user_id, item_id, rating这样的数据，其实就是数学里面的稀疏矩阵，scipy中提供了sparse模块来解决这个问题，但scipy.sparse有很多问题不太合用：1、不能很好的同时支持data[i, ...]、data[..., j]、data[i, j]快速切片；2、由于数据保存在内存中，不能很好的支持海量数据处理。

要支持data[i, ...]、data[..., j]的快速切片，需要i或者j的数据集中存储；同时，为了保存海量的数据，也需要把数据的一部分放在硬盘上，用内存做buffer。这里的解决方案比较简单，用一个类Dict的东西来存储数据，对于某个i（比如9527），它的数据保存在dict['i9527']里面，同样的，对于某个j（比如3306），它的全部数据保存在dict['j3306']里面，需要取出data[9527, ...]的时候，只要取出dict['i9527']即可，dict['i9527']原本是一个dict对象，储存某个j对应的值，为了节省内存空间，我们把这个dict以二进制字符串形式存储，直接上代码：

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112

''' Sparse Matrix ''' import struct import numpy as np import bsddb from cStringIO import StringIO    class DictMatrix():     def __init__(self, container = {}, dft = 0.0):         self._data  = container         self._dft   = dft         self._nums  = 0        def __setitem__(self, index, value):         try:             i, j = index         except:             raise IndexError('invalid index')            ik = ('i%d' % i)         # 为了节省内存，我们把j, value打包成字二进制字符串         ib = struct.pack('if', j, value)         jk = ('j%d' % j)         jb = struct.pack('if', i, value)            try:             self._data[ik] += ib         except:             self._data[ik] = ib         try:             self._data[jk] += jb         except:             self._data[jk] = jb         self._nums += 1        def __getitem__(self, index):         try:             i, j = index         except:             raise IndexError('invalid index')            if (isinstance(i, int)):             ik = ('i%d' % i)             if not self._data.has_key(ik): return self._dft             ret = dict(np.fromstring(self._data[ik], dtype = 'i4,f4'))             if (isinstance(j, int)): return ret.get(j, self._dft)            if (isinstance(j, int)):             jk = ('j%d' % j)             if not self._data.has_key(jk): return self._dft             ret = dict(np.fromstring(self._data[jk], dtype = 'i4,f4'))            return ret        def __len__(self):         return self._nums        def __iter__(self):         pass        '''     从文件中生成matrix     考虑到dbm读写的性能不如内存，我们做了一些缓存，每1000W次批量写入一次     考虑到字符串拼接性能不太好，我们直接用StringIO来做拼接     '''     def from_file(self, fp, sep = 't'):         cnt = 0         cache = {}         for l in fp:             if 10000000 == cnt:                 self._flush(cache)                 cnt = 0                 cache = {}             i, j, v = [float(i) for i in l.split(sep)]                ik = ('i%d' % i)             ib = struct.pack('if', j, v)             jk = ('j%d' % j)             jb = struct.pack('if', i, v)                try:                 cache[ik].write(ib)             except:                 cache[ik] = StringIO()                 cache[ik].write(ib)                try:                 cache[jk].write(jb)             except:                 cache[jk] = StringIO()                 cache[jk].write(jb)                cnt += 1             self._nums += 1            self._flush(cache)         return self._nums        def _flush(self, cache):         for k,v in cache.items():             v.seek(0)             s = v.read()             try:                 self._data[k] += s             except:                 self._data[k] = s    if __name__ == '__main__':     db = bsddb.btopen(None, cachesize = 268435456)     data = DictMatrix(db)     data.from_file(open('/path/to/log.txt', 'r'), ',')

测试4500W条rating数据（整形,整型,浮点格式），922MB文本文件导入，采用内存dict储存的话，12分钟构建完毕，消耗内存1.2G，采用示例代码中的bdb存储，20分钟构建完毕，占用内存300～400MB左右，比cachesize大不了多少，数据读取测试：

1 2	import timeit timeit.Timer('foo = __main__.data[9527, ...]', 'import __main__').timeit(number = 1000)

消耗1.4788秒，大概读取一条数据1.5ms。

采用类Dict来存储数据的另一个好处是你可以随便用内存Dict或者其他任何形式的DBM，甚至传说中的Tokyo Cabinet….

好的，码完收工。