Repository: dask/dask

Issue:
`astype` broken for partitioned parquet files
Computing `dask.dataframe.astype` on partitioned parquet files either fails with `ValueError` or (worse) messes up the column order. Self-contained example below:

```python
import numpy as np
import pandas as pd
import dask, dask.dataframe as dd

print('dask:', dask.__version__)

data = {
    'p': np.repeat(np.arange(3), 2).astype(np.int8),
    'b': np.repeat(-1, 6).astype(np.int16),
    'c': np.repeat(-2, 6).astype(np.float32),
    'd': np.repeat(-3, 6).astype(np.float64),
}
pdf = pd.DataFrame(data)
print('Original df:\n{}\n'.format(pdf))

path = 'test_df.pq'
ddf = dd.from_pandas(pdf, npartitions=2)
ddf.to_parquet(path, engine='pyarrow', partition_on='p')

ddfs = [
    dd.read_parquet(path),
    dd.read_parquet(path, columns=['p', 'b']),
    dd.read_parquet(path, columns=['p', 'b', 'd']),
]

for i, df in enumerate(ddfs):
    try:
        converted = df.astype({'b': np.float32}).compute()
    except ValueError:
        print('ddfs[{}]: Failed to convert\n'.format(i))
    else:
        print('ddfs[{}] head:\n{}\n'.format(i, converted.head()))
```
Output:
```
dask: 0.17.5
Original df:
   b    c    d  p
0 -1 -2.0 -3.0  0
1 -1 -2.0 -3.0  0
2 -1 -2.0 -3.0  1
3 -1 -2.0 -3.0  1
4 -1 -2.0 -3.0  2
5 -1 -2.0 -3.0  2

ddfs[0]: Failed to convert

ddfs[1] head:
     p  b
0 -1.0  0
1 -1.0  0
0 -1.0  0
1 -1.0  0
0 -1.0  0

ddfs[2] head:
     p    b  d
0 -1.0 -3.0  0
1 -1.0 -3.0  0
0 -1.0 -3.0  0
1 -1.0 -3.0  0
0 -1.0 -3.0  0
```

Possibly related to #3361.


Discussion:
Thank you for the excellent bug report @gsakkis . My apologies for the delay.

@martindurant, this seems to be in your realm.  If you have time can you take a look at this and give your impression about what might be amiss?
I'll take a look.
I would probably fix it on the dask side by, within `_read_pyarrow`, passing `meta` to `_read_pyarrow_parquet_piece` and enforcing columns/types on the output of that function; although it may well be that the fault is in the detail of `_read_pyarrow_parquet_piece`.
@gsakkis , any progress here?
@martindurant my current workaround is to call `dd.read_parquet()` with a `columns` list parameter that includes the partition column as the last column (along with a XXX comment of why this is needed) but I'm not familiar enough with dask/pyarrow internals to fix it myself.

Resolving patch:
diff --git a/dask/dataframe/io/parquet.py b/dask/dataframe/io/parquet.py
--- a/dask/dataframe/io/parquet.py
+++ b/dask/dataframe/io/parquet.py
@@ -505,6 +505,11 @@ def _read_pyarrow(fs, fs_token, paths, columns=None, filters=None,
         columns = list(columns)
 
     dataset = pq.ParquetDataset(paths, filesystem=get_pyarrow_filesystem(fs))
+    if dataset.partitions is not None:
+        partitions = [n for n in dataset.partitions.partition_names
+                      if n is not None]
+    else:
+        partitions = []
     schema = dataset.schema.to_arrow_schema()
     has_pandas_metadata = schema.metadata is not None and b'pandas' in schema.metadata
 
@@ -528,6 +533,7 @@ def _read_pyarrow(fs, fs_token, paths, columns=None, filters=None,
             index_names = [name_storage_mapping.get(name, name)
                            for name in index_names]
 
+    column_names += [p for p in partitions if p not in column_names]
     column_names, index_names, out_type = _normalize_index_columns(
         columns, column_names, index, index_names)
 
@@ -764,7 +770,7 @@ def _read_pyarrow_parquet_piece(fs, piece, columns, index_cols, is_series,
     if is_series:
         return df[df.columns[0]]
     else:
-        return df
+        return df[columns]
 
 
 _pyarrow_write_table_kwargs = {'row_group_size', 'version', 'use_dictionary',


Test patch:
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -1245,3 +1245,22 @@ def test_select_partitioned_column(tmpdir, engine):
 
     df_partitioned = dd.read_parquet(fn, engine=engine)
     df_partitioned[df_partitioned.fake_categorical1 == 'A'].compute()
+
+
+def test_arrow_partitioning(tmpdir):
+    # Issue #3518
+    pytest.importorskip('pyarrow')
+    path = str(tmpdir)
+    data = {
+        'p': np.repeat(np.arange(3), 2).astype(np.int8),
+        'b': np.repeat(-1, 6).astype(np.int16),
+        'c': np.repeat(-2, 6).astype(np.float32),
+        'd': np.repeat(-3, 6).astype(np.float64),
+    }
+    pdf = pd.DataFrame(data)
+    ddf = dd.from_pandas(pdf, npartitions=2)
+    ddf.to_parquet(path, engine='pyarrow', partition_on='p')
+
+    ddf = dd.read_parquet(path, engine='pyarrow')
+
+    ddf.astype({'b': np.float32}).compute()