Pass on file_scheme to fastparquet (dask#2714)

* Pass on file_scheme to fastparquet

* test drill cat in dask dataframe

* directly columns

Author: martindurant

dask/dataframe/io/parquet.py

@@ -103,7 +103,8 @@ def _read_fastparquet(fs, paths, myopen, columns=None, filters=None,
 
     dsk = {(name, i): (_read_parquet_row_group, myopen, pf.row_group_filename(rg),
                        index_col, all_columns, rg, out_type == Series,
-                       categories, pf.schema, pf.cats, pf.dtypes)
+                       categories, pf.schema, pf.cats, pf.dtypes,
+                       pf.file_scheme)
            for i, rg in enumerate(rgs)}
 
     if not dsk:
@@ -131,15 +132,15 @@ def _read_fastparquet(fs, paths, myopen, columns=None, filters=None,
 
 
 def _read_parquet_row_group(open, fn, index, columns, rg, series, categories,
-                            schema, cs, dt, *args):
+                            schema, cs, dt, scheme, *args):
     if not isinstance(columns, (tuple, list)):
         columns = (columns,)
         series = True
     if index and index not in columns:
         columns = columns + type(columns)([index])
     df, views = _pre_allocate(rg.num_rows, columns, categories, index, cs, dt)
     read_row_group_file(fn, rg, columns, categories, schema, cs,
-                        open=open, assign=views)
+                        open=open, assign=views, scheme=scheme)
 
     if series:
         return df[df.columns[0]]

dask/dataframe/io/tests/test_parquet.py

@@ -513,3 +513,26 @@ def test_timestamp96(fn):
     assert pf._schema[1].type == fastparquet.parquet_thrift.Type.INT96
     out = dd.read_parquet(fn).compute()
     assert_eq(out, df)
+
+
+def test_drill_scheme(fn):
+    N = 5
+    df1 = pd.DataFrame({c: np.random.random(N)
+                        for i, c in enumerate(['a', 'b', 'c'])})
+    df2 = pd.DataFrame({c: np.random.random(N)
+                        for i, c in enumerate(['a', 'b', 'c'])})
+    files = []
+    for d in ['test_data1', 'test_data2']:
+        dn = os.path.join(fn, d)
+        if not os.path.exists(dn):
+            os.mkdir(dn)
+        files.append(os.path.join(dn, 'data1.parq'))
+
+    fastparquet.write(files[0], df1)
+    fastparquet.write(files[1], df2)
+
+    df = dd.read_parquet(files)
+    assert 'dir0' in df.columns
+    out = df.compute()
+    assert 'dir0' in out
+    assert (np.unique(out.dir0) == ['test_data1', 'test_data2']).all()