Repository: dask/dask

Issue:
Issue with reading parquet file output from Spark with partition columns using pyarrow-dataset engine
**What happened**: I am reading a parquet file output from sparklyr code and partitioned on column ['a', 'b', 'c'] using `pyarrow-dataset` engine, but encounter below error message: 
`ValueError: partition(s) should not exist in columns.`

I went to the source code and find that it's due to these lines of code. [https://github.com/dask/dask/blob/main/dask/dataframe/io/parquet/arrow.py#L1090-L1094](https://github.com/dask/dask/blob/main/dask/dataframe/io/parquet/arrow.py#L1090-L1094)
It's because the column names have overlap with partition keys. I think it's because the output from sparkly doesn't contain metadata file, dask didn't recognize it properly. There is no issue reading this file with `pyarrow` or `pyarrow-legacy` engine. If reading the file in using `pyarrow-legacy` engine, and write out using dask, then read back in with `pyarrow-dataset`, there is no issue. 

**What you expected to happen**: I want to be able to read the parquet file output from spark with partition columns directly using `pyarrow-dataset` engine with dd.read_parquet(),  becuase `pyarrow-dataset` is said to be faster and with more advanced filtering options. 

One more question is that will `pyarrow` or `pyarrow-legacy` engine deprecated some time? If so, what's the estimate time?

**Minimal Complete Verifiable Example**:

```python
# Put your MCVE code here
```

**Anything else we need to know?**:

**Environment**:

- Dask version: 2021.03.1
- Python version: 3.8
- Operating System:
- Install method (conda, pip, source):



Discussion:
From your description it seems like sparklyr is writing the file incorrectly. Is that also your interpretation? 
Ping @rjzamora for estimate on `pyarrow-legacy` deprecation plan.
No,  sparklyr is writing the file properly, we are using databricks to run the sparklyr code, which gives parquet files with no metadata. We are okay with the file created by the sparklyr code, I am able to consume it elsewhere, but not using pyarrow-dataset engine on dask. 
Ok thanks for that, there might be some issues with pyarrow 3, so you could try using pyarrow 2 instead.
>One more question is that will pyarrow or pyarrow-legacy engine deprecated some time? If so, what's the estimate time?

I don't personally have a great estimate at the moment. There are a few downstream libraries that have yet to move to `ArrowDatasetEngine`, because the legacy engine is still a bit faster in some cases (or at least they *were* - I'll need to check pyarrow-3).

>No, sparklyr is writing the file properly, we are using databricks to run the sparklyr code, which gives parquet files with no metadata. We are okay with the file created by the sparklyr code, I am able to consume it elsewhere, but not using pyarrow-dataset engine on dask.

You should be able to read a partitioned dataset without the "_metadata" file, so I agree that there is a problem in Dask if you have problems doing this.  With that said, is your dataset "partitioned" in the conventional (hive-partitioning) sense?  That is, if you read a single file from your dataset, does it include the partitioned column?

For example, if a file is in the `id=1/` directory, it should not include an "id" column if you read the file by itself:
```python
import pyarrow.parquet as pq

pq.ParquetFile('/tmp/test_parquet/id=1/part.0.parquet').read().to_pandas()
```
```
    x
0   1
1   4
2   7
3  10
4  13
```
In general: there are indeed other systems that write partitioned datasets with keeping the partition columns also as physical columns in the files (as opposed to dask which drops them before writing the files), so I think it would be good for dask to support reading such data (plain ``pyarrow.dataset`` does support that).

Trying to create a reproducible example:

```python
import tempfile
import pathlib
path = pathlib.Path(tempfile.mkdtemp())

df1 = pd.DataFrame({'part': 'a', 'col': range(5)})
df2 = pd.DataFrame({'part': 'b', 'col': range(5)})
(path / "part=a").mkdir()
(path / "part=b").mkdir()
df1.to_parquet(path / "part=a" / "data.parquet", index=False, engine="pyarrow")
df2.to_parquet(path / "part=b" / "data.parquet", index=False, engine="pyarrow")
```

However, the above example fails for me with both `pyarrow-legacy` as `pyarrow-dataset`:

```
In [21]: dd.read_parquet(path, engine="pyarrow-legacy")
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-21-0d5b08436f3e> in <module>
----> 1 dd.read_parquet(path, engine="pyarrow-legacy")

~/scipy/repos/dask/dask/dataframe/io/parquet/core.py in read_parquet(path, columns, filters, categories, index, storage_options, engine, gather_statistics, split_row_groups, read_from_paths, chunksize, **kwargs)
    305         index = [index]
    306 
--> 307     read_metadata_result = engine.read_metadata(
    308         fs,
    309         paths,

~/scipy/repos/dask/dask/dataframe/io/parquet/arrow.py in read_metadata(cls, fs, paths, categories, index, gather_statistics, filters, split_row_groups, read_from_paths, **kwargs)
    516 
    517         # Process metadata to define `meta` and `index_cols`
--> 518         meta, index_cols, categories, index = cls._generate_dd_meta(
    519             schema, index, categories, partition_info
    520         )

~/scipy/repos/dask/dask/dataframe/io/parquet/arrow.py in _generate_dd_meta(cls, schema, index, categories, partition_info)
   1094 
   1095         if set(column_names).intersection(partitions):
-> 1096             raise ValueError(
   1097                 "partition(s) should not exist in columns.\n"
   1098                 "categories: {} | partitions: {}".format(column_names, partitions)

ValueError: partition(s) should not exist in columns.
categories: ['part', 'col'] | partitions: ['part']

```






>In general: there are indeed other systems that write partitioned datasets with keeping the partition columns also as physical columns in the files (as opposed to dask which drops them before writing the files), so I think it would be good for dask to support reading such data (plain pyarrow.dataset does support that).

Ah - Interesting.  This actually makes a lot of sense to me, but my only exposure to "partitioned" datasets has been through Dask :)

In the case that the partition column actually exists in the data, I assume that we should always **use** that data (rather than the value embedded in the partition name)?  I can imagine a scenario where the data type is coarsened for the partitioning (e.g. rounding to the nearest day for a timestamp), but the parquet file includes the fully-resolved data type.
Are there systems that use mixed behavior here while partitioning on multiple columns?  That is, is it possible that one partition column will be included in the file while others are not?  I ask, because it is slightly easier to support all or nothing.

Resolving patch:
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
--- a/dask/dataframe/io/parquet/arrow.py
+++ b/dask/dataframe/io/parquet/arrow.py
@@ -515,7 +515,7 @@ def read_metadata(
         )
 
         # Process metadata to define `meta` and `index_cols`
-        meta, index_cols, categories, index = cls._generate_dd_meta(
+        meta, index_cols, categories, index, partition_info = cls._generate_dd_meta(
             schema, index, categories, partition_info
         )
 
@@ -1092,11 +1092,21 @@ def _generate_dd_meta(cls, schema, index, categories, partition_info):
         if index is None and index_names:
             index = index_names
 
-        if set(column_names).intersection(partitions):
-            raise ValueError(
-                "partition(s) should not exist in columns.\n"
-                "categories: {} | partitions: {}".format(column_names, partitions)
-            )
+        # Ensure that there is no overlap between partition columns
+        # and explicit column storage
+        if partitions:
+            _partitions = [p for p in partitions if p not in column_names]
+            if not _partitions:
+                partitions = []
+                partition_info["partitions"] = None
+                partition_info["partition_keys"] = {}
+                partition_info["partition_names"] = partitions
+            elif len(_partitions) != len(partitions):
+                raise ValueError(
+                    "No partition-columns should be written in the \n"
+                    "file unless they are ALL written in the file.\n"
+                    "columns: {} | partitions: {}".format(column_names, partitions)
+                )
 
         column_names, index_names = _normalize_index_columns(
             columns, column_names + partitions, index, index_names
@@ -1142,7 +1152,7 @@ def _generate_dd_meta(cls, schema, index, categories, partition_info):
                         index=meta.index,
                     )
 
-        return meta, index_cols, categories, index
+        return meta, index_cols, categories, index, partition_info
 
     @classmethod
     def _construct_parts(
diff --git a/dask/dataframe/io/parquet/fastparquet.py b/dask/dataframe/io/parquet/fastparquet.py
--- a/dask/dataframe/io/parquet/fastparquet.py
+++ b/dask/dataframe/io/parquet/fastparquet.py
@@ -180,6 +180,19 @@ def _determine_pf_parts(fs, paths, gather_statistics, **kwargs):
             paths[0], open_with=fs.open, sep=fs.sep, **kwargs.get("file", {})
         )
 
+    # Ensure that there is no overlap between partition columns
+    # and explicit columns in `pf`
+    if pf.cats:
+        _partitions = [p for p in pf.cats if p not in pf.columns]
+        if not _partitions:
+            pf.cats = {}
+        elif len(_partitions) != len(pf.cats):
+            raise ValueError(
+                "No partition-columns should be written in the \n"
+                "file unless they are ALL written in the file.\n"
+                "columns: {} | partitions: {}".format(pf.columns, pf.cats.keys())
+            )
+
     return parts, pf, gather_statistics, base
 
 


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
@@ -2906,6 +2906,42 @@ def _arrow_table_to_pandas(
     dd.read_parquet(str(tmpdir), engine=ArrowEngineWithTimestampClamp).compute()
 
 
+@pytest.mark.parametrize(
+    "write_cols",
+    [["part", "col"], ["part", "kind", "col"]],
+)
+def test_partitioned_column_overlap(tmpdir, engine, write_cols):
+
+    tmpdir.mkdir("part=a")
+    tmpdir.mkdir("part=b")
+    path0 = str(tmpdir.mkdir("part=a/kind=x"))
+    path1 = str(tmpdir.mkdir("part=b/kind=x"))
+    path0 = os.path.join(path0, "data.parquet")
+    path1 = os.path.join(path1, "data.parquet")
+
+    _df1 = pd.DataFrame({"part": "a", "kind": "x", "col": range(5)})
+    _df2 = pd.DataFrame({"part": "b", "kind": "x", "col": range(5)})
+    df1 = _df1[write_cols]
+    df2 = _df2[write_cols]
+    df1.to_parquet(path0, index=False)
+    df2.to_parquet(path1, index=False)
+
+    if engine == "fastparquet":
+        path = [path0, path1]
+    else:
+        path = str(tmpdir)
+
+    if write_cols == ["part", "kind", "col"]:
+        result = dd.read_parquet(path, engine=engine)
+        expect = pd.concat([_df1, _df2], ignore_index=True)
+        assert_eq(result, expect, check_index=False)
+    else:
+        # For now, partial overlap between partition columns and
+        # real columns is not allowed
+        with pytest.raises(ValueError):
+            dd.read_parquet(path, engine=engine)
+
+
 @fp_pandas_xfail
 def test_partitioned_preserve_index(tmpdir, write_engine, read_engine):