Categoricals

partd/pandas.py

@@ -9,15 +9,18 @@
 from .utils import extend
 
 
+dumps = partial(pickle.dumps, protocol=pickle.HIGHEST_PROTOCOL)
+
+
 class PandasColumns(Interface):
     def __init__(self, partd=None):
         self.partd = Numpy(partd)
         Interface.__init__(self)
 
     def append(self, data, **kwargs):
         for k, df in data.items():
-            self.iset(extend(k, '.columns'), pickle.dumps(list(df.columns)))
-            self.iset(extend(k, '.index-name'), pickle.dumps(df.index.name))
+            self.iset(extend(k, '.columns'), dumps(list(df.columns)))
+            self.iset(extend(k, '.index-name'), dumps(df.index.name))
 
         # TODO: don't use values, it does some work.  Look at _blocks instead
         #       pframe/cframe do this well
@@ -104,16 +107,23 @@ def serialize(df):
     Uses Pandas blocks, snappy, and blosc to deconstruct an array into bytes
     """
     blocks, index, index_name, columns, placement = to_blocks(df)
+    categories = [(b.ordered, b.categories)
+                  if isinstance(b, pd.Categorical)
+                  else None
+                  for b in blocks]
+    blocks = [b.codes if isinstance(b, pd.Categorical) else b
+              for b in blocks]
     b_blocks = [pnp.compress(pnp.serialize(block), block.dtype)
                 for block in blocks]  # this can be slightly faster if we merge both operations
     b_index = pnp.compress(pnp.serialize(index), index.dtype)
-    frames = [pickle.dumps(index_name),
-              pickle.dumps(columns),
-              pickle.dumps(placement),
-              pickle.dumps(index.dtype),
+    frames = [dumps(index_name),
+              dumps(columns),
+              dumps(placement),
+              dumps(index.dtype),
               b_index,
-              pickle.dumps([block.dtype for block in blocks]),
-              pickle.dumps([block.shape for block in blocks])] + b_blocks
+              dumps([block.dtype for block in blocks]),
+              dumps([block.shape for block in blocks]),
+              dumps(categories)] + b_blocks
 
     return b''.join(map(frame, frames))
 
@@ -128,9 +138,14 @@ def deserialize(bytes):
     index = pnp.deserialize(pnp.decompress(frames[4], dt), dt, copy=True)
     dtypes = pickle.loads(frames[5])
     shapes = pickle.loads(frames[6])
-    b_blocks = frames[7:]
+    categories = pickle.loads(frames[7])
+    b_blocks = frames[8:]
     blocks = [pnp.deserialize(pnp.decompress(block, dt), dt, copy=True).reshape(shape)
                 for block, dt, shape in zip(b_blocks, dtypes, shapes)]
+    blocks = [pd.Categorical.from_codes(b, cat[1], ordered=cat[0])
+              if cat is not None
+              else b
+              for cat, b in zip(categories, blocks)]
 
     return from_blocks(blocks, index, index_name, columns, placement)
 

partd/tests/test_pandas.py

@@ -8,7 +8,7 @@
 import os
 import shutil
 
-from partd.pandas import PandasColumns, PandasBlocks
+from partd.pandas import PandasColumns, PandasBlocks, serialize, deserialize
 
 
 df1 = pd.DataFrame({'a': [1, 2, 3],
@@ -69,3 +69,14 @@ def test_PandasBlocks():
             result = p.get(['x'], lock=False)
 
     assert not os.path.exists(p.partd.path)
+
+
+@pytest.mark.parametrize('ordered', [False, True])
+def test_serialize_categoricals(ordered):
+    df = pd.DataFrame({'x': [1, 2, 3, 4],
+                       'y': pd.Categorical(['c', 'a', 'b', 'a'],
+                                           ordered=ordered)})
+
+    df2 = deserialize(serialize(df))
+    tm.assert_frame_equal(df, df2)
+    assert df.y.cat.ordered == df2.y.cat.ordered