Fix issues pandas-dev#219 and pandas-dev#220 (pandas-dev#221)

Author: bmoscon

arctic/store/_pickle_store.py

@@ -3,12 +3,13 @@
 from bson.errors import InvalidDocument
 from six.moves import cPickle, xrange
 import io
-import lz4
+from .._compression import compress, decompress, compress_array
 import pymongo
 
 from arctic.store._version_store_utils import checksum, pickle_compat_load
 
 _MAGIC_CHUNKED = '__chunked__'
+_MAGIC_CHUNKEDV2 = '__chunked__V2'
 _CHUNK_SIZE = 15 * 1024 * 1024  # 15MB
 _MAX_BSON_ENCODE = 256 * 1024  # 256K - don't fill up the version document with encoded bson
 
@@ -28,15 +29,19 @@ def get_info(self, version):
     def read(self, mongoose_lib, version, symbol, **kwargs):
         blob = version.get("blob")
         if blob is not None:
-            if blob == _MAGIC_CHUNKED:
+            if blob == _MAGIC_CHUNKEDV2:
+                collection = mongoose_lib.get_top_level_collection()
+                data = b''.join(decompress(x['data']) for x in collection.find({'symbol': symbol,
+                                                                                'parent': version['_id']},
+                                                                               sort=[('segment', pymongo.ASCENDING)]))
+            elif blob == _MAGIC_CHUNKED:
                 collection = mongoose_lib.get_top_level_collection()
                 data = b''.join(x['data'] for x in collection.find({'symbol': symbol,
-                                                                   'parent': version['_id']},
+                                                                    'parent': version['_id']},
                                                                    sort=[('segment', pymongo.ASCENDING)]))
+                data = decompress(data)
             else:
-                data = blob
-            # Backwards compatibility
-            data = lz4.decompress(data)
+                data = decompress(blob)
             return pickle_compat_load(io.BytesIO(data))
         return version['data']
 
@@ -53,13 +58,16 @@ def write(self, arctic_lib, version, symbol, item, previous_version):
         # Pickle, chunk and store the data
         collection = arctic_lib.get_top_level_collection()
         # Try to pickle it. This is best effort
-        version['blob'] = _MAGIC_CHUNKED
-        pickled = lz4.compressHC(cPickle.dumps(item, protocol=cPickle.HIGHEST_PROTOCOL))
+        version['blob'] = _MAGIC_CHUNKEDV2
+        pickled = cPickle.dumps(item, protocol=cPickle.HIGHEST_PROTOCOL)
+
+        data = compress_array([pickled[i * _CHUNK_SIZE: (i + 1) * _CHUNK_SIZE] for i in xrange(int(len(pickled) / _CHUNK_SIZE + 1))])
 
-        for i in xrange(int(len(pickled) / _CHUNK_SIZE + 1)):
-            segment = {'data': Binary(pickled[i * _CHUNK_SIZE : (i + 1) * _CHUNK_SIZE])}
+        for seg, d in enumerate(data):
+            segment = {'data': Binary(d)}
+            segment['segment'] = seg
+            seg += 1
             sha = checksum(symbol, segment)
-            segment['segment'] = i
-            collection.update_one({'symbol': symbol, 'sha': sha}, {'$set': segment,
-                                                               '$addToSet': {'parent': version['_id']}},
-                                       upsert=True)
+            collection.update_one({'symbol': symbol, 'sha': sha},
+                                  {'$set': segment, '$addToSet': {'parent': version['_id']}},
+                                  upsert=True)

tests/integration/store/test_pickle_store.py

@@ -2,7 +2,6 @@
 from datetime import datetime as dt, timedelta
 from mock import patch
 import numpy as np
-import re
 
 from arctic.arctic import Arctic
 
@@ -13,13 +12,25 @@ def test_save_read_bson(library):
     saved_blob = library.read('BLOB').data
     assert blob == saved_blob
 
+'''
+Run test at your own discretion. Takes > 60 secs
+def test_save_read_MASSIVE(library):
+    import pandas as pd
+    df = pd.DataFrame(data={'data': [1] * 150000000})
+    data = (df, df)
+    library.write('BLOB', data)
+    saved_blob = library.read('BLOB').data
+    assert(saved_blob[0].equals(df))
+    assert(saved_blob[1].equals(df))
+'''
+
 
 def test_save_read_big_encodable(library):
     blob = {'foo': 'a' * 1024 * 1024 * 20}
     library.write('BLOB', blob)
     saved_blob = library.read('BLOB').data
     assert blob == saved_blob
-    
+
 
 def test_save_read_bson_object(library):
     blob = {'foo': dt(2015, 1, 1), 'object': Arctic}

tests/unit/store/test_pickle_store.py

@@ -1,8 +1,5 @@
 import pandas as pd
-try:
-    import cPickle as pickle
-except ImportError:
-    import pickle
+from six.moves import cPickle
 import lz4
 import pytest
 import sys
@@ -12,6 +9,7 @@
 from distutils.version import LooseVersion
 from mock import create_autospec, sentinel, Mock, call
 
+from arctic._compression import compress, decompress
 from arctic.store._pickle_store import PickleStore
 from arctic.store._version_store_utils import checksum
 
@@ -32,13 +30,12 @@ def test_write_object():
     PickleStore.write(self, arctic_lib, version, 'sentinel.symbol', sentinel.item, sentinel.previous_version)
     assert 'data' not in version
 
-    assert version['blob'] == '__chunked__'
+    assert version['blob'] == '__chunked__V2'
     coll = arctic_lib.get_top_level_collection.return_value
-    assert coll.update_one.call_args_list == [call({'sha': checksum('sentinel.symbol',
-                                                                    {'data': Binary(lz4.compressHC(pickle.dumps(sentinel.item, pickle.HIGHEST_PROTOCOL)))}), 'symbol': 'sentinel.symbol'},
-                                               {'$set': {'segment': 0,
-                                                         'data': Binary(lz4.compressHC(pickle.dumps(sentinel.item, pickle.HIGHEST_PROTOCOL)), 0)},
-                                                         '$addToSet': {'parent': version['_id']}}, upsert=True)]
+    assert coll.update_one.call_args_list == [call({'sha': checksum('sentinel.symbol', {'segment':0, 'data': Binary(compress(cPickle.dumps(sentinel.item, cPickle.HIGHEST_PROTOCOL)))}),
+                                                    'symbol': 'sentinel.symbol'},
+                                                   {'$set': {'segment': 0, 'data': Binary(compress(cPickle.dumps(sentinel.item, cPickle.HIGHEST_PROTOCOL)), 0)},
+                                                    '$addToSet': {'parent': version['_id']}}, upsert=True)]
 
 
 def test_read():
@@ -49,7 +46,7 @@ def test_read():
 
 def test_read_object_backwards_compat():
     self = create_autospec(PickleStore)
-    version = {'blob': Binary(lz4.compressHC(pickle.dumps(object)))}
+    version = {'blob': Binary(lz4.compressHC(cPickle.dumps(object)))}
     assert PickleStore.read(self, sentinel.arctic_lib, version, sentinel.symbol) == object
 
 
@@ -59,7 +56,7 @@ def test_read_object_2():
                'blob': '__chunked__'}
     coll = Mock()
     arctic_lib = Mock()
-    coll.find.return_value = [{'data': Binary(lz4.compressHC(pickle.dumps(object))),
+    coll.find.return_value = [{'data': Binary(lz4.compressHC(cPickle.dumps(object))),
                                'symbol': 'sentinel.symbol'}
                               ]
     arctic_lib.get_top_level_collection.return_value = coll
@@ -79,10 +76,10 @@ def test_read_backward_compatibility():
     if PANDAS_VERSION >= LooseVersion("0.16.1"):
         if sys.version_info[0] >= 3:
             with pytest.raises(UnicodeDecodeError), open(fname) as fh:
-                pickle.load(fh)
+                cPickle.load(fh)
         else:    
             with pytest.raises(TypeError), open(fname) as fh:
-                pickle.load(fh)
+                cPickle.load(fh)
 
     # Verify that PickleStore() uses a backwards compatible unpickler.
     store = PickleStore()
@@ -101,7 +98,7 @@ def test_unpickle_highest_protocol():
     container has been pickled with HIGHEST_PROTOCOL.
     """
     version = {
-        'blob': lz4.compressHC(pickle.dumps(pd.Series(), protocol=pickle.HIGHEST_PROTOCOL)),
+        'blob': lz4.compressHC(cPickle.dumps(pd.Series(), protocol=cPickle.HIGHEST_PROTOCOL)),
     }
 
     store = PickleStore()
@@ -111,3 +108,23 @@ def test_unpickle_highest_protocol():
     assert (ps == expected).all()
 
 
+def test_pickle_chunk_V1_read():
+    data = {'foo': b'abcdefghijklmnopqrstuvwxyz'}
+    version = {'_id': sentinel._id,
+               'blob': '__chunked__'}
+    coll = Mock()
+    arctic_lib = Mock()
+    datap = lz4.compressHC(cPickle.dumps(data, protocol=cPickle.HIGHEST_PROTOCOL))
+    data_1 = datap[0:5]
+    data_2 = datap[5:]
+    coll.find.return_value = [{'data': Binary(data_1),
+                               'symbol': 'sentinel.symbol',
+                               'segment': 0},
+                              {'data': Binary(data_2),
+                               'symbol': 'sentinel.symbol',
+                               'segment': 1},
+                              ]
+    arctic_lib.get_top_level_collection.return_value = coll
+
+    ps = PickleStore()
+    assert(data == ps.read(arctic_lib, version, sentinel.symbol))