Merge pull request pandas-dev#58 from bmoscon/master

Fix issue with new versions of NumPy (v1.10.1) that cause dtype metadata to be lost

Author: jamesblackburn

arctic/store/_ndarray_store.py

@@ -210,35 +210,38 @@ def _do_read(self, collection, version, symbol, index_range=None):
         rtn = np.fromstring(data, dtype=dtype).reshape(version.get('shape', (-1)))
         return rtn
 
-    def _promote_types(self, item, dtype_str):
-        if dtype_str == str(item.dtype):
-            return item.dtype
+    def _promote_types(self, dtype, dtype_str):
+        if dtype_str == str(dtype):
+            return dtype
         prev_dtype = self._dtype(dtype_str)
-        if item.dtype.names is None:
-            rtn = np.promote_types(item.dtype, prev_dtype)
+        if dtype.names is None:
+            rtn = np.promote_types(dtype, prev_dtype)
         else:
-            rtn = _promote_struct_dtypes(item.dtype, prev_dtype)
-        rtn = np.dtype(rtn, metadata=dict(item.dtype.metadata or {}))
+            rtn = _promote_struct_dtypes(dtype, prev_dtype)
+        rtn = np.dtype(rtn, metadata=dict(dtype.metadata or {}))
         return rtn
 
-    def append(self, arctic_lib, version, symbol, item, previous_version):
+    def append(self, arctic_lib, version, symbol, item, previous_version, dtype=None):
         collection = arctic_lib.get_top_level_collection()
         if previous_version.get('shape', [-1]) != [-1, ] + list(item.shape)[1:]:
             raise UnhandledDtypeException()
 
-        if previous_version['up_to'] == 0:
+        if not dtype:
             dtype = item.dtype
+        
+        if previous_version['up_to'] == 0:
+            dtype = dtype
         elif len(item) == 0:
             dtype = self._dtype(previous_version['dtype'])
         else:
-            dtype = self._promote_types(item, previous_version['dtype'])
+            dtype = self._promote_types(dtype, previous_version['dtype'])
         item = item.astype(dtype)
         if str(dtype) != previous_version['dtype']:
             logger.debug('Converting %s from %s to %s' % (symbol, previous_version['dtype'], str(dtype)))
             if item.dtype.hasobject:
                 raise UnhandledDtypeException()
-            version['dtype'] = str(item.dtype)
-            version['dtype_metadata'] = dict(item.dtype.metadata or {})
+            version['dtype'] = str(dtype)
+            version['dtype_metadata'] = dict(dtype.metadata or {})
             version['type'] = self.TYPE
 
             old_arr = self._do_read(collection, previous_version, symbol).astype(dtype)
@@ -368,20 +371,22 @@ def checksum(self, item):
         sha.update(item.tostring())
         return Binary(sha.digest())
 
-    def write(self, arctic_lib, version, symbol, item, previous_version):
+    def write(self, arctic_lib, version, symbol, item, previous_version, dtype=None):
         collection = arctic_lib.get_top_level_collection()
         if item.dtype.hasobject:
             raise UnhandledDtypeException()
 
-        version['dtype'] = str(item.dtype)
+        if not dtype:
+            dtype = item.dtype
+        version['dtype'] = str(dtype)
         version['shape'] = (-1,) + item.shape[1:]
-        version['dtype_metadata'] = dict(item.dtype.metadata or {})
+        version['dtype_metadata'] = dict(dtype.metadata or {})
         version['type'] = self.TYPE
         version['up_to'] = len(item)
         version['sha'] = self.checksum(item)
-
+        
         if previous_version:
-            if version['dtype'] == str(item.dtype) \
+            if version['dtype'] == str(dtype) \
                     and 'sha' in previous_version \
                     and self.checksum(item[:previous_version['up_to']]) == previous_version['sha']:
                 #The first n rows are identical to the previous version, so just append.

arctic/store/_pandas_ndarray_store.py

@@ -87,15 +87,18 @@ def to_records(self, df):
         arrays = map(_to_primitive, arrays)
         dtype = np.dtype([(str(x), v.dtype) if len(v.shape) == 1 else (str(x), v.dtype, v.shape[1]) for x, v in zip(names, arrays)],
                          metadata=metadata)
+
         rtn = np.rec.fromarrays(arrays, dtype=dtype, names=names)
-        #For some reason the dtype metadata is lost in the line above.
-        rtn.dtype = dtype
-        return rtn
+        # For some reason the dtype metadata is lost in the line above
+        # and setting rtn.dtype to dtype does not preserve the metadata
+        # see https://github.com/numpy/numpy/issues/6771
+
+        return (rtn, dtype)
 
     def can_convert_to_records_without_objects(self, df, symbol):
         # We can't easily distinguish string columns from objects
         try:
-            arr = self.to_records(df)
+            arr,_ = self.to_records(df)
         except Exception as e:
             # This exception will also occur when we try to write the object so we fall-back to saving using Pickle
             log.info('Pandas dataframe %s caused exception "%s" when attempting to convert to records. Saving as Blob.'
@@ -249,12 +252,12 @@ def can_write(self, version, symbol, data):
         return False
 
     def write(self, arctic_lib, version, symbol, item, previous_version):
-        item = self.to_records(item)
-        super(PandasSeriesStore, self).write(arctic_lib, version, symbol, item, previous_version)
+        item, md = self.to_records(item)
+        super(PandasSeriesStore, self).write(arctic_lib, version, symbol, item, previous_version, dtype=md)
 
     def append(self, arctic_lib, version, symbol, item, previous_version):
-        item = self.to_records(item)
-        super(PandasSeriesStore, self).append(arctic_lib, version, symbol, item, previous_version)
+        item, md = self.to_records(item)
+        super(PandasSeriesStore, self).append(arctic_lib, version, symbol, item, previous_version, dtype=md)
 
     def read(self, arctic_lib, version, symbol, **kwargs):
         item = super(PandasSeriesStore, self).read(arctic_lib, version, symbol, **kwargs)
@@ -287,12 +290,12 @@ def can_write(self, version, symbol, data):
         return False
 
     def write(self, arctic_lib, version, symbol, item, previous_version):
-        item = self.to_records(item)
-        super(PandasDataFrameStore, self).write(arctic_lib, version, symbol, item, previous_version)
+        item, md = self.to_records(item)
+        super(PandasDataFrameStore, self).write(arctic_lib, version, symbol, item, previous_version, dtype=md)
 
     def append(self, arctic_lib, version, symbol, item, previous_version):
-        item = self.to_records(item)
-        super(PandasDataFrameStore, self).append(arctic_lib, version, symbol, item, previous_version)
+        item, md = self.to_records(item)
+        super(PandasDataFrameStore, self).append(arctic_lib, version, symbol, item, previous_version, dtype=md)
 
     def read(self, arctic_lib, version, symbol, **kwargs):
         item = super(PandasDataFrameStore, self).read(arctic_lib, version, symbol, **kwargs)

tests/unit/store/test_pandas_ndarray_store.py

@@ -21,7 +21,7 @@ def test_can_convert_to_records_without_objects_returns_false_on_exception_in_to
 
 def test_can_convert_to_records_without_objects_returns_false_when_records_have_object_dtype():
     store = PandasStore()
-    store.to_records = Mock(return_value=np.array(['a', 'b', None, 'd']))
+    store.to_records = Mock(return_value=(np.array(['a', 'b', None, 'd']), None))
 
     with patch('arctic.store._pandas_ndarray_store.log') as mock_log:
         assert store.can_convert_to_records_without_objects(sentinel.df, 'my_symbol') is False
@@ -32,8 +32,8 @@ def test_can_convert_to_records_without_objects_returns_false_when_records_have_
 
 def test_can_convert_to_records_without_objects_returns_false_when_records_have_arrays_in_them():
     store = PandasStore()
-    store.to_records = Mock(return_value=np.rec.array([(1356998400000000000L, ['A', 'BC'])],
-                                                      dtype=[('index', '<M8[ns]'), ('values', 'S2', (2,))]))
+    store.to_records = Mock(return_value=(np.rec.array([(1356998400000000000L, ['A', 'BC'])],
+                                                       dtype=[('index', '<M8[ns]'), ('values', 'S2', (2,))]), None))
 
     with patch('arctic.store._pandas_ndarray_store.log') as mock_log:
         assert store.can_convert_to_records_without_objects(sentinel.df, 'my_symbol') is False
@@ -44,8 +44,8 @@ def test_can_convert_to_records_without_objects_returns_false_when_records_have_
 
 def test_can_convert_to_records_without_objects_returns_true_otherwise():
     store = PandasStore()
-    store.to_records = Mock(return_value=np.rec.array([(1356998400000000000L, 'a')],
-                                                      dtype=[('index', '<M8[ns]'), ('values', 'S2')]))
+    store.to_records = Mock(return_value=(np.rec.array([(1356998400000000000L, 'a')],
+                                                       dtype=[('index', '<M8[ns]'), ('values', 'S2')]), None))
 
     with patch('arctic.store._pandas_ndarray_store.log') as mock_log:
         assert store.can_convert_to_records_without_objects(sentinel.df, 'my_symbol') is True
@@ -54,15 +54,6 @@ def test_can_convert_to_records_without_objects_returns_true_otherwise():
     store.to_records.assert_called_once_with(sentinel.df)
 
 
-def test_to_records_raises_when_object_dtypes_present():
-    store = PandasDataFrameStore()
-    df = pd.DataFrame(data=dict(A=['a', 'b', None, 'c'], B=[1., 2., 3., 4.]), index=range(4))
-    with raises(TypeError) as e:
-        store.to_records(df)
-
-    assert "Cannot change data-type for object array." in str(e)
-
-
 def test_panel_converted_to_dataframe_and_stacked_to_write():
     store = PandasPanelStore()
     panel = Mock(shape=(1, 2, 3), axes=[Mock(names=['n%d' % i]) for i in range(3)])