Fix #2720 (#2729)

dask · Oct 1, 2017 · 0f29e7b · 0f29e7b
1 parent 68d8c95
commit 0f29e7b
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/dask/dataframe/io/parquet.py b/dask/dataframe/io/parquet.py
@@ -91,8 +91,9 @@ def _read_fastparquet(fs, paths, myopen, columns=None, filters=None,
                              dtypes)
 
     for cat in categories:
-        meta[cat] = pd.Series(pd.Categorical([],
-                              categories=[UNKNOWN_CATEGORIES]))
+        if cat in meta:
+            meta[cat] = pd.Series(pd.Categorical([],
+                                  categories=[UNKNOWN_CATEGORIES]))
 
     if index_col:
         meta = meta.set_index(index_col)

diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
@@ -536,3 +536,19 @@ def test_drill_scheme(fn):
     out = df.compute()
     assert 'dir0' in out
     assert (np.unique(out.dir0) == ['test_data1', 'test_data2']).all()
+
+
+def test_parquet_select_cats(fn):
+    df = pd.DataFrame({
+        'categories': pd.Series(
+            np.random.choice(['a', 'b', 'c', 'd', 'e', 'f'], size=100),
+            dtype='category'),
+        'ints': pd.Series(list(range(0, 100)), dtype='int'),
+        'floats': pd.Series(list(range(0, 100)), dtype='float')})
+
+    ddf = dd.from_pandas(df, 1)
+    ddf.to_parquet(fn)
+    rddf = dd.read_parquet(fn, columns=['ints'])
+    assert list(rddf.columns) == ['ints']
+    rddf = dd.read_parquet(fn)
+    assert list(rddf.columns) == list(df)