Repository: joblib/joblib

Issue:
joblib.dump bug with too large object-type arrays
issue from
https://github.com/scikit-learn/scikit-learn/issues/4889
The code
```python
import joblib
import os
import numpy as np
from gzip import GzipFile
from io import BytesIO
from urllib2 import urlopen
from os.path import join
from sklearn.datasets import get_data_home

URL10 = ('http://archive.ics.uci.edu/ml/'
         'machine-learning-databases/kddcup99-mld/kddcup.data_10_percent.gz')

data_home = get_data_home()
kddcup_dir = join(data_home, "test")
samples_path = join(kddcup_dir, "samples")
os.makedirs(kddcup_dir)
f = BytesIO(urlopen(URL10).read())
file = GzipFile(fileobj=f, mode='r')

X = []
for line in file.readlines():
    X.append(line.replace('\n', '').split(','))

file.close()
X = np.asarray(X, dtype=object)
joblib.dump(X, samples_path, compress=9)
X = joblib.load(samples_path)
```
More precisely, it works if X has less than 300000 lines or if X is not of dtype `object`:
```python
Y = X[:300000,:]  ### works
joblib.dump(Y, samples_path, compress=9)
Y = joblib.load(samples_path)
###
Y = X[:400000,:].astype(str)  ### works
joblib.dump(Y, samples_path, compress=9)
Y = joblib.load(samples_path)
###
Y = X[:400000,:]  ### doesnt work
joblib.dump(Y, samples_path, compress=9)
Y = joblib.load(samples_path)
```

The error raised by joblib.load is 
```
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-32-25fe17420f08> in <module>()
----> 1 Y = joblib.load(samples_path)

/usr/lib/python2.7/dist-packages/joblib/numpy_pickle.pyc in load(filename, mmap_mode)
    422 
    423         try:
--> 424             obj = unpickler.load()
    425         finally:
    426             if hasattr(unpickler, 'file_handle'):

/usr/lib/python2.7/pickle.pyc in load(self)
    856             while 1:
    857                 key = read(1)
--> 858                 dispatch[key](self)
    859         except _Stop, stopinst:
    860             return stopinst.value

/usr/lib/python2.7/dist-packages/joblib/numpy_pickle.pyc in load_build(self)
    288                         "but numpy didn't import correctly")
    289             nd_array_wrapper = self.stack.pop()
--> 290             array = nd_array_wrapper.read(self)
    291             self.stack.append(array)
    292 

/usr/lib/python2.7/dist-packages/joblib/numpy_pickle.pyc in read(self, unpickler)
    158         array = unpickler.np.core.multiarray._reconstruct(*self.init_args)
    159         with open(filename, 'rb') as f:
--> 160             data = read_zfile(f)
    161         state = self.state + (data,)
    162         array.__setstate__(state)

/usr/lib/python2.7/dist-packages/joblib/numpy_pickle.pyc in read_zfile(file_handle)
     69     assert len(data) == length, (
     70         "Incorrect data length while decompressing %s."
---> 71         "The file could be corrupted." % file_handle)
     72     return data
     73 

AssertionError: Incorrect data length while decompressing <open file '/home/nicolas/scikit_learn_data/test/samples_01.npy.z', mode 'rb' at 0x7efcea714db0>.The file could be corrupted.
```
@lesteve what do you think?

Discussion:
lesteve: > Maybe I'm wrong, but why not adding here a condition "not obj.dtype.hasobject" and keep the previous behavior (using the except) ?

We could do that I am just slightly worried that this try except is not a very clean way of hiding failures. A few weird behaviours we recently saw:
* for Python 2.6 and @aabadie's PR with `with gzip.GZipFile as f`, you get a failure when you exit the context manager, meaning that you actually write a valid .z auxiliary file but then still pickle the array inside the master file because you encountered an error. This auxiliary file is never deleted.
* for Python 3, when trying to dump an array with dtype object, you encountered an error quite early on but you still write an invalid .z auxiliary file with just the header. This auxiliary file is never deleted.
* for Python 2.7, when trying to dump an array with dtype object, you don't even encounter an error and you end up writing str(array). You only see the problem at load time.

@aabadie as an aside, just a gentle reminder: always better to comment in the "Files changed" tab rather than in individual commits.


Patch:
diff --git a/joblib/numpy_pickle.py b/joblib/numpy_pickle.py
index 6d727d6..f2d43eb 100644
--- a/joblib/numpy_pickle.py
+++ b/joblib/numpy_pickle.py
@@ -215,7 +215,7 @@ class NumpyPickler(Pickler):
         else:
             self.file = BytesIO()
         # Count the number of npy files that we have created:
-        self._npy_counter = 0
+        self._npy_counter = 1
         # By default we want a pickle protocol that only changes with
         # the major python version and not the minor one
         if protocol is None:
@@ -257,8 +257,8 @@ class NumpyPickler(Pickler):
             files, rather than pickling them. Of course, this is a
             total abuse of the Pickler class.
         """
-        if self.np is not None and type(obj) in (self.np.ndarray,
-                                            self.np.matrix, self.np.memmap):
+        if (self.np is not None and type(obj) in
+                (self.np.ndarray, self.np.matrix, self.np.memmap)):
             size = obj.size * obj.itemsize
             if self.compress and size < self.cache_size * _MEGA:
                 # When compressing, as we are not writing directly to the
@@ -267,19 +267,21 @@ class NumpyPickler(Pickler):
                     # Pickling doesn't work with memmaped arrays
                     obj = self.np.asarray(obj)
                 return Pickler.save(self, obj)
-            self._npy_counter += 1
-            try:
-                filename = '%s_%02i.npy' % (self._filename,
-                                            self._npy_counter)
-                # This converts the array in a container
-                obj, filename = self._write_array(obj, filename)
-                self._filenames.append(filename)
-            except:
-                self._npy_counter -= 1
-                # XXX: We should have a logging mechanism
-                print('Failed to save %s to .npy file:\n%s' % (
+
+            if not obj.dtype.hasobject:
+                try:
+                    filename = '%s_%02i.npy' % (self._filename,
+                                                self._npy_counter)
+                    # This converts the array in a container
+                    obj, filename = self._write_array(obj, filename)
+                    self._filenames.append(filename)
+                    self._npy_counter += 1
+                except Exception:
+                    # XXX: We should have a logging mechanism
+                    print('Failed to save %s to .npy file:\n%s' % (
                         type(obj),
                         traceback.format_exc()))
+
         return Pickler.save(self, obj)
 
     def close(self):


Test patch:
diff --git a/joblib/test/test_numpy_pickle.py b/joblib/test/test_numpy_pickle.py
index ab80354..c6b8960 100644
--- a/joblib/test/test_numpy_pickle.py
+++ b/joblib/test/test_numpy_pickle.py
@@ -260,6 +260,7 @@ def test_z_file():
 def test_compressed_pickle_dump_and_load():
     expected_list = [np.arange(5, dtype=np.int64),
                      np.arange(5, dtype=np.float64),
+                     np.array([1, 'abc', {'a': 1, 'b': 2}]),
                      # .tostring actually returns bytes and is a
                      # compatibility alias for .tobytes which was
                      # added in 1.9.0
@@ -269,17 +270,23 @@ def test_compressed_pickle_dump_and_load():
     with tempfile.NamedTemporaryFile(suffix='.gz', dir=env['dir']) as f:
         fname = f.name
 
-    try:
-        numpy_pickle.dump(expected_list, fname, compress=1)
-        result_list = numpy_pickle.load(fname)
-        for result, expected in zip(result_list, expected_list):
-            if isinstance(expected, np.ndarray):
-                nose.tools.assert_equal(result.dtype, expected.dtype)
-                np.testing.assert_equal(result, expected)
-            else:
-                nose.tools.assert_equal(result, expected)
-    finally:
-        os.remove(fname)
+    # Need to test both code branches (whether array size is greater
+    # or smaller than cache_size)
+    for cache_size in [0, 1e9]:
+        try:
+            dumped_filenames = numpy_pickle.dump(
+                expected_list, fname, compress=1,
+                cache_size=cache_size)
+            result_list = numpy_pickle.load(fname)
+            for result, expected in zip(result_list, expected_list):
+                if isinstance(expected, np.ndarray):
+                    nose.tools.assert_equal(result.dtype, expected.dtype)
+                    np.testing.assert_equal(result, expected)
+                else:
+                    nose.tools.assert_equal(result, expected)
+        finally:
+            for fn in dumped_filenames:
+                os.remove(fn)
 
 
 def _check_pickle(filename, expected_list):