Merge pull request pandas-dev#148 from manahl/chunkstore_get_info

Add get info to chunkstore

Author: bmoscon

CHANGES.md

@@ -3,6 +3,7 @@
 ### 1.26
 
   * Bugfix: Faster TickStore querying for multiple symbols simultaneously
+  * Bugfix: #147 Add get_info method to ChunkStore
 
 ### 1.25 (2016-05-23)
 

arctic/chunkstore/chunkstore.py

@@ -2,6 +2,7 @@
 import pymongo
 import numpy as np
 import bson
+import ast
 
 from bson.binary import Binary
 from pandas import Series, DataFrame
@@ -131,7 +132,7 @@ def read(self, symbol, chunk_range=None):
 
         data = b''.join(segments)
 
-        dtype = PandasSerializer()._dtype(sym['dtype'], sym.get('dtype_metadata', {}))
+        dtype = PandasSerializer._dtype(sym['dtype'], sym.get('dtype_metadata', {}))
         records = np.fromstring(data, dtype=dtype).reshape(sym.get('shape', (-1)))
 
         data = deserialize(records, sym['type'])
@@ -377,3 +378,17 @@ def update(self, symbol, item):
                 sym['append_size'] += seg_len
                 sym['append_count'] += seg_count
             self._symbols.replace_one({'symbol': symbol}, sym)
+
+    def get_info(self, symbol):
+        sym = self._get_symbol_info(symbol)
+        ret = {}
+        dtype = PandasSerializer._dtype(sym['dtype'], sym['dtype_metadata'])
+        length = sym['len']
+        ret['size'] = dtype.itemsize * length
+        ret['chunk_count'] = sym['chunk_count']
+        ret['dtype'] = sym['dtype']
+        ret['type'] = sym['type']
+        ret['rows'] = length
+        ret['col_names'] = sym['dtype_metadata']
+        ret['dtype'] = ast.literal_eval(sym['dtype'])
+        return ret

arctic/serialization/pandas_serializer.py

@@ -42,7 +42,8 @@ def _to_primitive(arr, string_max_len=None):
 
 class PandasSerializer(object):
 
-    def _dtype(self, string, metadata=None):
+    @staticmethod
+    def _dtype(string, metadata=None):
         if metadata is None:
             metadata = {}
         if string.startswith('['):

tests/integration/chunkstore/test_chunkstore.py

@@ -595,3 +595,24 @@ def test_delete(chunkstore_lib):
     assert ('test_df' in chunkstore_lib.list_symbols())
     chunkstore_lib.delete('test_df')
     assert (chunkstore_lib.list_symbols() == [])
+
+
+def test_get_info(chunkstore_lib):
+    df = DataFrame(data={'data': [1, 2, 3]},
+                   index=MultiIndex.from_tuples([(dt(2016, 1, 1), 1),
+                                                 (dt(2016, 1, 2), 1),
+                                                 (dt(2016, 1, 3), 1)],
+                                                names=['date', 'id'])
+                   )
+    chunkstore_lib.write('test_df', df, 'D')
+    info = {'rows': 3,
+            'dtype': [('date', '<M8[ns]'), ('id', '<i8'), ('data', '<i8')],
+            'chunk_count': 3,
+            'col_names': {
+                            u'index': [u'date', u'id'],
+                            u'index_tz': [None, None],
+                            u'columns': [u'data']
+                         },
+            'type': u'df',
+            'size': 72}
+    assert(chunkstore_lib.get_info('test_df') == info)