to_castra protects list of categories

mrocklin · Jul 15, 2015 · 3c3ac7a · 3c3ac7a
1 parent 5ffaa8a
commit 3c3ac7a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
@@ -652,6 +652,8 @@ def to_castra(self, fn=None, categories=None):
         """
         from castra import Castra
         name = 'to-castra' + next(tokens)
+        if isinstance(categories, list):
+            categories = (list, categories)
         dsk = {name: (Castra, fn, (self._name, 0), categories)}
         dsk.update(dict(((name, i), (Castra.extend, name, (self._name, i)))
                         for i in range(self.npartitions)))

diff --git a/dask/dataframe/tests/test_io.py b/dask/dataframe/tests/test_io.py
@@ -398,12 +398,21 @@ def test_from_dask_array_raises():
 
 def test_to_castra():
     pytest.importorskip('castra')
-    df = pd.DataFrame({'x': ['a', 'b', 'c', 'D'],
+    df = pd.DataFrame({'x': ['a', 'b', 'c', 'd'],
                        'y': [1, 2, 3, 4]})
     a = dd.from_pandas(df, 2)
 
     c = a.to_castra()
-    assert eq(a, c[:])
+    try:
+        assert eq(a, c[:])
+    finally:
+        c.drop()
+
+    c = a.to_castra(categories=['x'])
+    try:
+        assert c[:].dtypes['x'] == 'category'
+    finally:
+        c.drop()
 
 
 def test_to_hdf():