ENH: sparse data structures in HDFStore. #85

Author: Chang She

pandas/io/pytables.py

@@ -12,6 +12,8 @@
 from pandas import (
     Series, TimeSeries, DataFrame, Panel, Index, MultiIndex, Int64Index
 )
+from pandas.sparse.api import SparseSeries, SparseDataFrame, SparsePanel
+from pandas.sparse.array import BlockIndex, IntIndex
 from pandas.tseries.api import PeriodIndex, DatetimeIndex
 from pandas.core.common import adjoin
 from pandas.core.algorithms import match, unique
@@ -28,17 +30,23 @@
 # reading and writing the full object in one go
 _TYPE_MAP = {
     Series     : 'series',
+    SparseSeries : 'sparse_series',
     TimeSeries : 'series',
     DataFrame  : 'frame',
-    Panel  : 'wide'
+    SparseDataFrame : 'sparse_frame',
+    Panel  : 'wide',
+    SparsePanel : 'sparse_panel'
 }
 
 _NAME_MAP = {
     'series' : 'Series',
     'time_series' : 'TimeSeries',
+    'sparse_series' : 'SparseSeries',
     'frame' : 'DataFrame',
+    'sparse_frame' : 'SparseDataFrame',
     'frame_table' : 'DataFrame (Table)',
     'wide' : 'Panel',
+    'sparse_panel' : 'SparsePanel',
     'wide_table' : 'Panel (Table)',
     'long' : 'LongPanel',
     # legacy h5 files
@@ -406,6 +414,78 @@ def _write_series(self, group, series):
         self._write_array(group, 'values', series.values)
         group._v_attrs.name = series.name
 
+    def _write_sparse_series(self, group, series):
+        self._write_index(group, 'index', series.index)
+        self._write_index(group, 'sp_index', series.sp_index)
+        self._write_array(group, 'sp_values', series.sp_values)
+        group._v_attrs.name = series.name
+        group._v_attrs.fill_value = series.fill_value
+        group._v_attrs.kind = series.kind
+
+    def _read_sparse_series(self, group, where=None):
+        index = self._read_index(group, 'index')
+        sp_values = _read_array(group, 'sp_values')
+        sp_index = self._read_index(group, 'sp_index')
+        name = getattr(group._v_attrs, 'name', None)
+        fill_value = getattr(group._v_attrs, 'fill_value', None)
+        kind = getattr(group._v_attrs, 'kind', 'block')
+        return SparseSeries(sp_values, index=index, sparse_index=sp_index,
+                            kind=kind, fill_value=fill_value,
+                            name=name)
+
+    def _write_sparse_frame(self, group, sdf):
+        for name, ss in sdf.iteritems():
+            key = 'sparse_series_%s' % name
+            if key not in group._v_children:
+                node = self.handle.createGroup(group, key)
+            else:
+                node = getattr(group, key)
+            self._write_sparse_series(node, ss)
+        setattr(group._v_attrs, 'default_fill_value',
+                sdf.default_fill_value)
+        setattr(group._v_attrs, 'default_kind',
+                sdf.default_kind)
+        self._write_index(group, 'columns', sdf.columns)
+
+    def _read_sparse_frame(self, group, where=None):
+        columns = self._read_index(group, 'columns')
+        sdict = {}
+        for c in columns:
+            key = 'sparse_series_%s' % c
+            node = getattr(group, key)
+            sdict[c] = self._read_sparse_series(node)
+        default_kind = getattr(group._v_attrs, 'default_kind')
+        default_fill_value = getattr(group._v_attrs, 'default_fill_value')
+        return SparseDataFrame(sdict, columns=columns,
+                               default_kind=default_kind,
+                               default_fill_value=default_fill_value)
+
+    def _write_sparse_panel(self, group, swide):
+        setattr(group._v_attrs, 'default_fill_value', swide.default_fill_value)
+        setattr(group._v_attrs, 'default_kind', swide.default_kind)
+        self._write_index(group, 'items', swide.items)
+
+        for name, sdf in swide.iteritems():
+            key = 'sparse_frame_%s' % name
+            if key not in group._v_children:
+                node = self.handle.createGroup(group, key)
+            else:
+                node = getattr(group, key)
+            self._write_sparse_frame(node, sdf)
+
+    def _read_sparse_panel(self, group, where=None):
+        default_fill_value = getattr(group._v_attrs, 'default_fill_value')
+        default_kind = getattr(group._v_attrs, 'default_kind')
+        items = self._read_index(group, 'items')
+
+        sdict = {}
+        for name in items:
+            key = 'sparse_frame_%s' % name
+            node = getattr(group, key)
+            sdict[name] = self._read_sparse_frame(node)
+        return SparsePanel(sdict, items=items, default_kind=default_kind,
+                           default_fill_value=default_fill_value)
+
     def _write_frame(self, group, df):
         self._write_block_manager(group, df._data)
 
@@ -474,21 +554,32 @@ def _read_wide_table(self, group, where=None):
         return self._read_panel_table(group, where)
 
     def _write_index(self, group, key, index):
-        if len(index) == 0:
-            raise ValueError('Can not write empty structure, axis length was 0')
-
         if isinstance(index, MultiIndex):
+            if len(index) == 0:
+                raise ValueError('Can not write empty structure, '
+                                 'axis length was 0')
+
             setattr(group._v_attrs, '%s_variety' % key, 'multi')
             self._write_multi_index(group, key, index)
+        elif isinstance(index, BlockIndex):
+            setattr(group._v_attrs, '%s_variety' % key, 'block')
+            self._write_block_index(group, key, index)
+        elif isinstance(index, IntIndex):
+            setattr(group._v_attrs, '%s_variety' % key, 'sparseint')
+            self._write_sparse_intindex(group, key, index)
         else:
+            if len(index) == 0:
+                raise ValueError('Can not write empty structure, '
+                                 'axis length was 0')
+
             setattr(group._v_attrs, '%s_variety' % key, 'regular')
             converted, kind, _ = _convert_index(index)
             self._write_array(group, key, converted)
             node = getattr(group, key)
             node._v_attrs.kind = kind
             node._v_attrs.name = index.name
 
-            if isinstance(index, (DatetimeIndex, PeriodIndex)):
+            if isinstance(index, (DatetimeIndex, PeriodIndex, IntIndex)):
                 node._v_attrs.index_class = type(index)
 
             if hasattr(index, 'freq'):
@@ -499,12 +590,36 @@ def _read_index(self, group, key):
 
         if variety == 'multi':
             return self._read_multi_index(group, key)
+        elif variety == 'block':
+            return self._read_block_index(group, key)
+        elif variety == 'sparseint':
+            return self._read_sparse_intindex(group, key)
         elif variety == 'regular':
             _, index = self._read_index_node(getattr(group, key))
             return index
         else:  # pragma: no cover
             raise Exception('unrecognized index variety: %s' % variety)
 
+    def _write_block_index(self, group, key, index):
+        self._write_array(group, '%s_blocs' % key, index.blocs)
+        self._write_array(group, '%s_blengths' % key, index.blengths)
+        setattr(group._v_attrs, '%s_length' % key, index.length)
+
+    def _read_block_index(self, group, key):
+        length = getattr(group._v_attrs, '%s_length' % key)
+        blocs = _read_array(group, '%s_blocs' % key)
+        blengths = _read_array(group, '%s_blengths' % key)
+        return BlockIndex(length, blocs, blengths)
+
+    def _write_sparse_intindex(self, group, key, index):
+        self._write_array(group, '%s_indices' % key, index.indices)
+        setattr(group._v_attrs, '%s_length' % key, index.length)
+
+    def _read_sparse_intindex(self, group, key):
+        length = getattr(group._v_attrs, '%s_length' % key)
+        indices = _read_array(group, '%s_indices' % key)
+        return IntIndex(length, indices)
+
     def _write_multi_index(self, group, key, index):
         setattr(group._v_attrs, '%s_nlevels' % key, index.nlevels)
 

pandas/io/tests/test_pytables.py

@@ -24,7 +24,7 @@
 _default_compressor = LooseVersion(tables.__version__) >= '2.2' \
                       and 'blosc' or 'zlib'
 
-class TesttHDFStore(unittest.TestCase):
+class TestHDFStore(unittest.TestCase):
     path = '__test__.h5'
     scratchpath = '__scratch__.h5'
 
@@ -201,6 +201,53 @@ def test_series(self):
         ts = tm.makeTimeSeries()
         self._check_roundtrip(ts, tm.assert_series_equal)
 
+    def test_sparse_series(self):
+        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({i : tm.makeDataFrame() 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 test_float_index(self):
         # GH #454
         index = np.random.randn(10)
@@ -486,7 +533,7 @@ def test_select_filter_corner(self):
         result = self.store.select('frame', [crit])
         tm.assert_frame_equal(result, df.ix[:, df.columns[:75]])
 
-    def _check_roundtrip(self, obj, comparator, compression=False):
+    def _check_roundtrip(self, obj, comparator, compression=False, **kwargs):
         options = {}
         if compression:
             options['complib'] = _default_compressor
@@ -495,7 +542,7 @@ def _check_roundtrip(self, obj, comparator, compression=False):
         try:
             store['obj'] = obj
             retrieved = store['obj']
-            comparator(retrieved, obj)
+            comparator(retrieved, obj, **kwargs)
         finally:
             store.close()
             os.remove(self.scratchpath)

pandas/util/testing.py

@@ -119,7 +119,10 @@ def assert_dict_equal(a, b, compare_keys=True):
 
 def assert_series_equal(left, right, check_dtype=True,
                         check_index_type=False,
-                        check_index_freq=False):
+                        check_index_freq=False,
+                        check_series_type=False):
+    if check_series_type:
+        assert(type(left) == type(right))
     assert_almost_equal(left.values, right.values)
     if check_dtype:
         assert(left.dtype == right.dtype)
@@ -133,7 +136,10 @@ def assert_series_equal(left, right, check_dtype=True,
                getattr(right, 'freqstr', None))
 
 def assert_frame_equal(left, right, check_index_type=False,
-                       check_column_type=False):
+                       check_column_type=False,
+                       check_frame_type=False):
+    if check_frame_type:
+        assert(type(left) == type(right))
     assert(isinstance(left, DataFrame))
     assert(isinstance(right, DataFrame))
     for col, series in left.iterkv():
@@ -152,7 +158,10 @@ def assert_frame_equal(left, right, check_index_type=False,
         assert(left.columns.dtype == right.columns.dtype)
         assert(left.columns.inferred_type == right.columns.inferred_type)
 
-def assert_panel_equal(left, right):
+def assert_panel_equal(left, right, check_panel_type=False):
+    if check_panel_type:
+        assert(type(left) == type(right))
+
     assert(left.items.equals(right.items))
     assert(left.major_axis.equals(right.major_axis))
     assert(left.minor_axis.equals(right.minor_axis))