ENH: added support for Panel,SparseSeries,SparseDataFrame,SparsePanel,IntIndex,BlockIndex

Author: jreback

doc/source/io.rst

@@ -1008,6 +1008,21 @@ You can pass a list of objects and you will receive them back on deserialization
    pd.to_msgpack('foo.msg', df, 'foo', np.array([1,2,3]), s)
    pd.read_msgpack('foo.msg')
 
+You can pass ``iterator=True`` to iterator over the unpacked results
+
+.. ipython:: python
+
+   for o in pd.read_msgpack('foo.msg',iterator=True):
+       print o
+
+
+You can pass ``append=True`` to the writer to append to an existing pack
+
+.. ipython:: python
+
+   df.to_msgpack('foo.msg',append=True)
+   pd.read_msgpack('foo.msg')
+
 .. ipython:: python
    :suppress:
    :okexcept:

doc/source/v0.12.0.txt

@@ -24,6 +24,13 @@ Enhancements
           pd.to_msgpack('foo.msg', df, s)
           pd.read_msgpack('foo.msg')
 
+    You can pass ``iterator=True`` to iterator over the unpacked results
+
+    .. ipython:: python
+
+          for o in pd.read_msgpack('foo.msg',iterator=True):
+              print o
+
     .. ipython:: python
           :suppress:
           :okexcept:

pandas/io/packers.py

@@ -57,6 +57,7 @@
 from pandas.tseries.api import PeriodIndex, DatetimeIndex
 from pandas.core.index import Int64Index, _ensure_index
 import pandas.core.common as com
+from pandas.core.generic import NDFrame
 from pandas.core.common import needs_i8_conversion
 from pandas.core.internals import BlockManager, make_block
 import pandas.core.internals as internals
@@ -162,6 +163,7 @@ def encode(obj):
     Data encoder
     """
 
+    tobj = type(obj)
     if isinstance(obj, Index):
         if isinstance(obj, PeriodIndex):
             return {'typ' : 'period_index',
@@ -191,25 +193,47 @@ def encode(obj):
                     'data': obj.tolist() }
     elif isinstance(obj, Series):
         if isinstance(obj, SparseSeries):
-            import pdb; pdb.set_trace()
+            d = {'typ' : 'sparse_series',
+                 'klass' : obj.__class__.__name__,
+                 'dtype': obj.dtype.name,
+                 'index' : obj.index,
+                 'sp_index' : obj.sp_index,
+                 'sp_values' : convert(obj.sp_values)}
+            for f in ['name','fill_value','kind']:
+                d[f] = getattr(obj,f,None)
+            return d
         else:
             return {'typ' : 'series',
                     'klass' : obj.__class__.__name__,
                     'name' : getattr(obj,'name',None),
                     'index' : obj.index,
                     'dtype': obj.dtype.name,
                     'data': convert(obj.values) }
-    elif isinstance(obj, DataFrame):
+    elif issubclass(tobj, NDFrame):
         if isinstance(obj, SparseDataFrame):
-            import pdb; pdb.set_trace()
+            d = {'typ' : 'sparse_dataframe',
+                 'klass' : obj.__class__.__name__,
+                 'columns' : obj.columns }
+            for f in ['default_fill_value','default_kind']:
+                d[f] = getattr(obj,f,None)
+            d['data'] = dict([ (name,ss) for name,ss in obj.iteritems() ])
+            return d
+        elif isinstance(obj, SparsePanel):
+            d = {'typ' : 'sparse_panel',
+                 'klass' : obj.__class__.__name__,
+                 'items' : obj.items }
+            for f in ['default_fill_value','default_kind']:
+                d[f] = getattr(obj,f,None)
+            d['data'] = dict([ (name,df) for name,df in obj.iteritems() ])
+            return d
         else:
 
             data = obj._data
             if not data.is_consolidated():
                 data = data.consolidate()
 
            # the block manager
-            return {'typ' : 'dataframe',
+            return {'typ' : 'block_manager',
                     'klass'  : obj.__class__.__name__,
                     'axes'   : data.axes,
                     'blocks' : [ { 'items'  : b.items, 
@@ -237,6 +261,17 @@ def encode(obj):
         return {'typ' : 'period',
                 'ordinal' : obj.ordinal,
                 'freq' : obj.freq }
+    elif isinstance(obj, BlockIndex):
+        return { 'typ' : 'block_index',
+                 'klass' : obj.__class__.__name__,
+                 'blocs' : obj.blocs,
+                 'blengths' : obj.blengths,
+                 'length' : obj.length }
+    elif isinstance(obj, IntIndex):
+        return { 'typ' : 'int_index',
+                 'klass' : obj.__class__.__name__,
+                 'indices' : obj.indices,
+                 'length' : obj.length }
     elif isinstance(obj, np.ndarray):
         return {'typ' : 'ndarray',
                 'shape': obj.shape,
@@ -288,7 +323,7 @@ def decode(obj):
         dtype = dtype_for(obj['dtype'])
         index = obj['index']
         return globals()[obj['klass']](obj['data'],index=index,dtype=dtype,name=obj['name'])
-    elif typ == 'dataframe':
+    elif typ == 'block_manager':
         axes = obj['axes']
 
         def create_block(b):
@@ -300,6 +335,20 @@ def create_block(b):
     elif typ == 'datetime':
         import pdb; pdb.set_trace()
         return datetime.fromtimestamp(obj['data'])
+    elif typ == 'sparse_series':
+        dtype = dtype_for(obj['dtype'])
+        return globals()[obj['klass']](np.array(obj['sp_values'],dtype=dtype),sparse_index=obj['sp_index'],
+                                       index=obj['index'],fill_value=obj['fill_value'],kind=obj['kind'],name=obj['name'])
+    elif typ == 'sparse_dataframe':
+        return globals()[obj['klass']](obj['data'],
+                                       columns=obj['columns'],default_fill_value=obj['default_fill_value'],default_kind=obj['default_kind'])
+    elif typ == 'sparse_panel':
+        return globals()[obj['klass']](obj['data'],
+                                       items=obj['items'],default_fill_value=obj['default_fill_value'],default_kind=obj['default_kind'])
+    elif typ == 'block_index':
+        return globals()[obj['klass']](obj['length'],obj['blocs'],obj['blengths'])
+    elif typ == 'int_index':
+        return globals()[obj['klass']](obj['length'],obj['indices'])
     elif typ == 'ndarray':
         return np.array(obj['data'],
                         dtype=np.typeDict[obj['dtype']],

pandas/io/tests/test_packers.py

@@ -14,6 +14,10 @@
 from pandas.util.testing import ensure_clean
 from pandas.tests.test_series import assert_series_equal
 from pandas.tests.test_frame import assert_frame_equal
+from pandas.tests.test_panel import assert_panel_equal
+
+import pandas
+from pandas.sparse.tests.test_sparse import assert_sp_series_equal, assert_sp_frame_equal
 from pandas import concat, Timestamp, tslib
 
 from numpy.testing.decorators import slow
@@ -32,6 +36,8 @@ def check_arbitrary(a, b):
         assert(len(a) == len(b))
         for a_, b_ in zip(a,b):
             check_arbitrary(a_,b_)
+    elif isinstance(a,Panel):
+        assert_panel_equal(a,b)
     elif isinstance(a,DataFrame):
         assert_frame_equal(a,b)
     elif isinstance(a,Series):
@@ -225,10 +231,10 @@ def test_basic(self):
             i_rec = self.encode_decode(i)
             assert_series_equal(i,i_rec)
 
-class TestFrame(Test):
+class TestNDFrame(Test):
 
     def setUp(self):
-        super(TestFrame, self).setUp()
+        super(TestNDFrame, self).setUp()
 
         data = {
             'A': [0., 1., 2., 3., np.nan],
@@ -238,98 +244,107 @@ def setUp(self):
             'E' : [0., 1, Timestamp('20100101'),'foo',2.],
             }
 
-        self.d = { 'float' : DataFrame(dict(A = data['A'], B = Series(data['A']) + 1)),
-                   'int'   : DataFrame(dict(A = data['B'], B = Series(data['B']) + 1)),
-                   'mixed' : DataFrame(dict([ (k,data[k]) for k in ['A','B','C','D']])) }
+        self.frame = { 'float' : DataFrame(dict(A = data['A'], B = Series(data['A']) + 1)),
+                       'int'   : DataFrame(dict(A = data['B'], B = Series(data['B']) + 1)),
+                       'mixed' : DataFrame(dict([ (k,data[k]) for k in ['A','B','C','D']])) }
+        
+        self.panel = { 'float' : Panel(dict(ItemA = self.frame['float'], ItemB = self.frame['float']+1)) }
 
-    def test_basic(self):
+    def test_basic_frame(self):
 
-        for s, i in self.d.items():
+        for s, i in self.frame.items():
             i_rec = self.encode_decode(i)
             assert_frame_equal(i,i_rec)
 
+    def test_basic_panel(self):
+
+        for s, i in self.panel.items():
+            i_rec = self.encode_decode(i)
+            assert_panel_equal(i,i_rec)
+
     def test_multi(self):
 
-        i_rec = self.encode_decode(self.d)
-        for k in self.d.keys():
-            assert_frame_equal(self.d[k],i_rec[k])
+        i_rec = self.encode_decode(self.frame)
+        for k in self.frame.keys():
+            assert_frame_equal(self.frame[k],i_rec[k])
 
-        l = tuple([ self.d['float'], self.d['float'].A, self.d['float'].B, None ])
+        l = tuple([ self.frame['float'], self.frame['float'].A, self.frame['float'].B, None ])
         l_rec = self.encode_decode(l)
         check_arbitrary(l,l_rec)
 
         # this is an oddity in that packed lists will be returned as tuples
-        l = [ self.d['float'], self.d['float'].A, self.d['float'].B, None ]
+        l = [ self.frame['float'], self.frame['float'].A, self.frame['float'].B, None ]
         l_rec = self.encode_decode(l)
         self.assert_(isinstance(l_rec,tuple))
         check_arbitrary(l,l_rec)
 
     def test_iterator(self):
 
-        l = [ self.d['float'], self.d['float'].A, self.d['float'].B, None ]
+        l = [ self.frame['float'], self.frame['float'].A, self.frame['float'].B, None ]
 
         with ensure_clean(self.path) as path:
             to_msgpack(path,*l)
             for i, packed in enumerate(read_msgpack(path, iterator=True)):
                 check_arbitrary(packed,l[i])
 
-def _create_sp_series():
+class TestSparse(Test):
 
-    # nan-based
-    arr = np.arange(15, dtype=float)
-    index = np.arange(15)
-    arr[7:12] = nan
-    arr[-1:] = nan
+    def _check_roundtrip(self, obj, comparator, **kwargs):
 
-    date_index = bdate_range('1/1/2011', periods=len(index))
-    bseries = SparseSeries(arr, index=index, kind='block')
-    bseries.name = 'bseries'
-    return bseries
+        i_rec = self.encode_decode(obj)
+        comparator(obj,i_rec,**kwargs)
 
-def _create_sp_frame():
+    def test_sparse_series(self):
 
-    data = {'A': [nan, nan, nan, 0, 1, 2, 3, 4, 5, 6],
-            'B': [0, 1, 2, nan, nan, nan, 3, 4, 5, 6],
-            'C': np.arange(10),
-            'D': [0, 1, 2, 3, 4, 5, nan, nan, nan, nan]}
-    
-    dates = bdate_range('1/1/2011', periods=10)
-    return SparseDataFrame(data, index=dates)
+        s = tm.makeStringSeries()
+        s[3:5] = np.nan
+        ss = s.to_sparse()
+        self._check_roundtrip(ss, tm.assert_series_equal,
+                              check_series_type=True)
+        
+        ss2 = s.to_sparse(kind='integer')
+        self._check_roundtrip(ss2, tm.assert_series_equal,
+                              check_series_type=True)
+        
+        ss3 = s.to_sparse(fill_value=0)
+        self._check_roundtrip(ss3, tm.assert_series_equal,
+                              check_series_type=True)
+
+    def test_sparse_frame(self):
+
+        s = tm.makeDataFrame()
+        s.ix[3:5, 1:3] = np.nan
+        s.ix[8:10, -2] = np.nan
+        ss = s.to_sparse()
+
+        self._check_roundtrip(ss, tm.assert_frame_equal,
+                              check_frame_type=True)
+        
+        ss2 = s.to_sparse(kind='integer')
+        self._check_roundtrip(ss2, tm.assert_frame_equal,
+                              check_frame_type=True)
+        
+        ss3 = s.to_sparse(fill_value=0)
+        self._check_roundtrip(ss3, tm.assert_frame_equal,
+                              check_frame_type=True)
+            
+    def test_sparse_panel(self):
+
+        items = ['x', 'y', 'z']
+        p = Panel(dict((i, tm.makeDataFrame().ix[:2, :2]) for i in items))
+        sp = p.to_sparse()
+        
+        self._check_roundtrip(sp, tm.assert_panel_equal,
+                              check_panel_type=True)
+        
+        sp2 = p.to_sparse(kind='integer')
+        self._check_roundtrip(sp2, tm.assert_panel_equal,
+                              check_panel_type=True)
+        
+        sp3 = p.to_sparse(fill_value=0)
+        self._check_roundtrip(sp3, tm.assert_panel_equal,
+                              check_panel_type=True)
 
-def create_data():
-    """ create the pickle data """
-    
-    data = {
-        'A': [0., 1., 2., 3., np.nan],
-        'B': [0, 1, 0, 1, 0],
-        'C': ['foo1', 'foo2', 'foo3', 'foo4', 'foo5'],
-        'D': date_range('1/1/2009', periods=5),
-        'E' : [0., 1, Timestamp('20100101'),'foo',2.],
-        }
-    
-    index  = dict(int   = Index(np.arange(10)),
-                  date  = date_range('20130101',periods=10))
-    mi     = dict(reg   = MultiIndex.from_tuples(zip([['bar', 'bar', 'baz', 'baz', 'foo', 'foo', 'qux', 'qux'],
-                                                      ['one', 'two', 'one', 'two', 'one', 'two', 'one', 'two']]),
-                                                 names=['first', 'second']))
-    series = dict(float = Series(data['A']),
-                  int   = Series(data['B']),
-                  mixed = Series(data['E']))
-    frame  = dict(float = DataFrame(dict(A = series['float'], B = series['float'] + 1)),
-                  int   = DataFrame(dict(A = series['int']  , B = series['int']   + 1)),
-                  mixed = DataFrame(dict([ (k,data[k]) for k in ['A','B','C','D']])))
-    panel  = dict(float = Panel(dict(ItemA = frame['float'], ItemB = frame['float']+1)))
-
- 
-
-    return dict( series = series, 
-                 frame  = frame, 
-                 panel  = panel,
-                 index  = index,
-                 mi     = mi,
-                 sp_series = dict(float = _create_sp_series()),
-                 sp_frame  = dict(float = _create_sp_frame())
-                 )
 
 if __name__ == '__main__':
     import nose