Repository: dask/dask

Issue:
Error "column not in index" when using read_parquet() on hive-partitioned parquet dataset
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:
The `dask.dataframe.read_parquet()` method fails with error `KeyError: "['col1'] not in index"` when trying to read a partitioned parquet dataset (in my example generated by ` pyarrow`).

**Minimal Complete Verifiable Example**:

```python
import pandas as pd
import pyarrow as pa
import pyarrow.dataset as ds
import dask.dataframe as dd

df = pd.DataFrame({
    'col1': [1, 2],
    'col2': ['a', 'b']
})

table = pa.Table.from_pandas(df)

ds.write_dataset(data=table,
                 base_dir='foo',
                 basename_template="prefix-{i}.parquet",
                 format="parquet",
                 partitioning=["col1"],
                 existing_data_behavior="delete_matching"
                 )

dd1 = dd.read_parquet(path='foo', engine="pyarrow")
print(dd1)
dd1.compute()
```

```text
Dask DataFrame Structure:
                col1    col2
npartitions=3               
               int64  object
                 ...     ...
                 ...     ...
                 ...     ...
Dask Name: read-parquet, 3 tasks
```
```text
KeyError                                  Traceback (most recent call last)
~/dask-issue.py in <module>
     22 dd1 = dd.read_parquet(path='foo', engine="pyarrow")
     23 print(dd1)
---> 24 dd1.compute()
...
~/.local/share/virtualenvs/.../lib/python3.8/site-packages/pandas/core/indexing.py in _validate_read_indexer(self, key, indexer, axis, raise_missing)
   1314             if raise_missing:
   1315                 not_found = list(set(key) - set(ax))
-> 1316                 raise KeyError(f"{not_found} not in index")
   1317 
   1318             not_found = key[missing_mask]

KeyError: "['col1'] not in index"
```

**Anything else we need to know?**:
Please note that the lazy `print(dd1)` does include `col1`, it's only at `.compute()` we are failing.

I have previously experienced that `pyarrow` (without dask) needs the `partitioning` and `partition_base_dir` args to be set, in order to read the partition columns. Maybe this is the issue somewhere?
I have tested these `kwargs` without luck:

```python
dd1 = dd.read_parquet(path='foo', engine="pyarrow", 
    format="parquet", partitioning="hive", partition_base_dir="foo"
)
```

**Environment**:

- Dask version: 2022.1.1
- Pyarrow version: 6.0.1
- Python version: 3.8
- Operating System: Ubuntu 20.04.3 LTS 
- Install method (conda, pip, source): pipenv

<!-- If you are reporting an issue such as scale stability, cluster deadlock.
Please provide a cluster dump state with this issue, by running client.dump_cluster_state()

https://distributed.dask.org/en/stable/api.html?highlight=dump_cluster_state#distributed.Client.dump_cluster_state

-->

<details>
<summary>Cluster Dump State:</summary>

</details>


Discussion:
Hi @th0ger thanks for reporting this! @pavithraes has volunteered to take a look 😄 
@pavithraes a bit closer to a workaround:

```python
dd.read_parquet(path='foo', engine="pyarrow",
                dataset={
                    "partitioning": "hive",
                    "partition_base_dir": "foo"
                  }
                ).compute()
```
```text
...
~/.local/share/.../lib/python3.8/site-packages/dask/dataframe/io/parquet/core.py in read_parquet(path, columns, filters, categories, index, storage_options, engine, gather_statistics, ignore_metadata_file, metadata_task_size, split_row_groups, chunksize, aggregate_files, **kwargs)
    328         gather_statistics = True
    329 
--> 330     read_metadata_result = engine.read_metadata(
    331         fs,
    332         paths,

~/.local/share/.../lib/python3.8/site-packages/dask/dataframe/io/parquet/arrow.py in read_metadata(cls, fs, paths, categories, index, gather_statistics, filters, split_row_groups, chunksize, aggregate_files, ignore_metadata_file, metadata_task_size, **kwargs)
    369 
    370         # Stage 1: Collect general dataset information
--> 371         dataset_info = cls._collect_dataset_info(
    372             paths,
    373             fs,

~/.local/share/.../lib/python3.8/site-packages/dask/dataframe/io/parquet/arrow.py in _collect_dataset_info(cls, paths, fs, categories, index, gather_statistics, filters, split_row_groups, chunksize, aggregate_files, ignore_metadata_file, metadata_task_size, kwargs)
    911                 filesystem=fs,
    912                 format="parquet",
--> 913                 partitioning=partitioning["obj"].discover(
    914                     *partitioning.get("args", []),
    915                     **partitioning.get("kwargs", {}),

TypeError: string indices must be integers
```

The `partitioning="hive"` argument is valid pyarrow syntax that _should_ solve the issue. I belive pyarrow changed the behaviour, so you must provide this if the partion keys shall be returned. [Reading partitioned data](https://arrow.apache.org/docs/python/dataset.html#reading-partitioned-data)

I also tested this syntax:

```python
from pyarrow.dataset import HivePartitioning
partitioning = HivePartitioning(pa.schema([("col1", pa.int16()), ("col2", pa.string())]))
dd.read_parquet(path='foo', engine="pyarrow",
                dataset={
                    "partitioning": partitioning,
                    "partition_base_dir": "foo"
                  }
                ).compute()
```
```text
...
~/.local/share/virtualenvs/.../lib/python3.8/site-packages/dask/dataframe/io/parquet/core.py in read_parquet(path, columns, filters, categories, index, storage_options, engine, gather_statistics, ignore_metadata_file, metadata_task_size, split_row_groups, chunksize, aggregate_files, **kwargs)
    328         gather_statistics = True
    329 
--> 330     read_metadata_result = engine.read_metadata(
    331         fs,
    332         paths,

~/.local/share/virtualenvs/.../lib/python3.8/site-packages/dask/dataframe/io/parquet/arrow.py in read_metadata(cls, fs, paths, categories, index, gather_statistics, filters, split_row_groups, chunksize, aggregate_files, ignore_metadata_file, metadata_task_size, **kwargs)
    369 
    370         # Stage 1: Collect general dataset information
--> 371         dataset_info = cls._collect_dataset_info(
    372             paths,
    373             fs,

~/.local/share/virtualenvs/.../lib/python3.8/site-packages/dask/dataframe/io/parquet/arrow.py in _collect_dataset_info(cls, paths, fs, categories, index, gather_statistics, filters, split_row_groups, chunksize, aggregate_files, ignore_metadata_file, metadata_task_size, kwargs)
    911                 filesystem=fs,
    912                 format="parquet",
--> 913                 partitioning=partitioning["obj"].discover(
    914                     *partitioning.get("args", []),
    915                     **partitioning.get("kwargs", {}),

TypeError: 'pyarrow._dataset.HivePartitioning' object is not subscriptable
```
https://github.com/dask/dask/blob/b6713736533f98afda441e1d0319b7f5b45e53ad/dask/dataframe/io/parquet/arrow.py#L829-L838
@pavithraes, this special construct is giving some challenges in passing `partititioning="hive"` to `pyarrow.dataset.dataset()`:

https://github.com/dask/dask/blob/b6713736533f98afda441e1d0319b7f5b45e53ad/dask/dataframe/io/parquet/arrow.py#L857-L865

Here is a workaround that passes `"hive"`:

```python
class What:
    def discover(self, *args, **kwargs):
        return "hive"

partitioning = {"obj": What()}

assert partitioning["obj"].discover(
                        *partitioning.get("args", []),
                        **partitioning.get("kwargs", {}),
                    ) == "hive"

dd.read_parquet(path='foo', engine="pyarrow",
                dataset={
                    "partitioning": partitioning,
                    "partition_base_dir": "foo",
                  }
                ).compute()
```
However still getting error `["['col1'] not in index"]()`.

Wonder if this is because dask for python3.8 depends on pyarrow 4.0 (according to `environment-3.8.yaml`). I normally use pyarrow 6.0.1 when not using dask.
@bryanwweber @pavithraes could you kindly provide your comments on this?
@th0ger Thanks for your patience here!

@scharlottej13 and I looked at this, and we think this is related to pyarrow's `write_dataset`. We don't see `col1` while reading the parquet dataset `foo` using pyarrow's `read_table` either:

```python
import pandas as pd
import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.parquet as pq

df = pd.DataFrame({
    'col1': [1, 2],
    'col2': ['a', 'b']
})

table = pa.Table.from_pandas(df)

ds.write_dataset(data=table,
                 base_dir='foo',
                 basename_template="prefix-{i}.parquet",
                 format="parquet",
                 partitioning=["col1"],
                 existing_data_behavior="delete_matching")

print(pq.read_table('foo'))

# Output:
#
# pyarrow.Table
# col2: string
# ----
# col2: [["a"],["b"]]
```
which suggests this isn't an issue in Dask's `read_parquet`.

We're not super familiar with pyarrow, but maybe it has something to do with writing the `_metadata` files? We'd encourage you to share this question with the pyarrow community as well :)

Using Dask's `to_parquet` to write the file, then reading with `read_parquet` seems to work fine:

```python
import dask.dataframe as dd

ddf = dd.from_pandas(df, npartitions=2)

ddf.to_parquet(path='bar', engine='pyarrow', partition_on=["col1"])

dd1 = dd.read_parquet(path='bar', engine="pyarrow")
dd1.compute()
```

A side note, we noticed this warning with pyarrow:
`The pyarrow.dataset module is experimental (specifically the classes), and a stable API is not yet guaranteed.`
Hi @pavithraes and @scharlottej13, thank for looking into this.

First one regression from my end: I forgot to provide `partitioning_flavor="hive"` to `ds.write_dataset()`. This setting ensures that pyarrow names the partition folders as hive-style `/col1=<value>/` instead of just `/<value>/`. It seems to me that pyarrow does not write the `_metadata` and `_common_metadata` files that Dask does (at least by default). So without setting `partitioning_flavor="hive"`, the given name `"col1"` is never persisted, and theres no way that you could read `col1` in automatically.

We don't want to depend on using Dask solely for writing parquet, just to be able to use Dask to read it back. Dask must be able to read partitioned data written with pyarrow (or fastparquet etc.) - especially when we tell Dask to read it via the exact same pyarrow backend.

When it comes to the stability of the new pyarrow Dataset API, my impression is that it's getting quite stable. Also, let me remind you that Dask `engine="pyarrow"` now defaults to the new pyarrow Dataset API, so you already jumped in the water :-)
> engine : ... If 'pyarrow' or 'pyarrow-dataset' is
> specified, the ArrowDatasetEngine (which leverages the pyarrow.dataset
> API) will be used. If 'pyarrow-legacy' is specified, ArrowLegacyEngine
> will be used (which leverages the pyarrow.parquet.ParquetDataset API).
> NOTE: The 'pyarrow-legacy' option (ArrowLegacyEngine) is deprecated
> for pyarrow>=5.

Here is my updated test which shows that shows the interoperability between writing/reading with pyarrow/Dask:

```python
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import pyarrow.dataset as ds
import dask.dataframe as dd

df = pd.DataFrame({
    'col1': [1, 2],
    'col2': ['a', 'b']
})

## Write with dask
ddf = dd.from_pandas(df, npartitions=2)
ddf.to_parquet(path='foo-dask', engine='pyarrow', partition_on=["col1"])
# ls foo-dask/ 
# _common_metadata   _metadata  'col1=1'  'col1=2'

## Write with pyarrow
table = pa.Table.from_pandas(df)
ds.write_dataset(data=table,
                 base_dir='foo-pyarrow',
                 basename_template="part.{i}.parquet",
                 format="parquet",
                 partitioning=["col1"],
                 partitioning_flavor="hive",  # Set hive partition style /col1=<val>/
                 existing_data_behavior="delete_matching"
                 )
# ls foo-pyarrow/
# 'col1=1'  'col1=2'

## Read with pyarrow (old API)
print(pq.read_table('foo-dask').to_pandas())
print(pq.read_table('foo-pyarrow').to_pandas())
#                                   col2 col1
# __null_dask_index__          
# 0                                 a    1
# 1                                 b    2
#
#       col2 col1
# 0    a    1
# 1    b    2

## Read with pyarrow (new API)
print(ds.dataset("foo-dask", format="parquet", partitioning="hive").to_table().to_pandas())
print(ds.dataset("foo-pyarrow", format="parquet", partitioning="hive").to_table().to_pandas())
# Same output as above.

## Read with dask
print(dd.read_parquet(path='foo-dask', engine="pyarrow").compute())
try:
    print(dd.read_parquet(path='foo-pyarrow', engine="pyarrow").compute())
except Exception as e:
    print(e)
#     col2 col1
# 0    a    1
# 1    b    2
#
# "['col1'] not in index"
```

I still belive this is caused by the way Dask implementes its calls to pyarrow. But maybe some pyarrow API/behaviour has changed and someone from their team can advice.

I tried the workaround:
```python
dd.read_parquet(path='foo-pyarrow', engine="pyarrow", dataset={"partitioning": "hive"})
```
but this returns `[TypeError: string indices must be integers]()`. It seems this is not allowed due to Dask `kwargs` handling, similar to the discussion in #8677.
@th0ger Thanks for sharing a nice minimal example, I was able to reproduce this. :)

I'm wondering if this might be a bug in Dask, maybe @jcrist or @rjzamora can help confirm?

Because it seems to work as expected using pandas' `read_parquet`: `pd.read_parquet(path='foo-pyarrow', engine="pyarrow")`. With this information, here's a non-ideal workaround using `Delayed`:

```python
from dask import delayed

delayed_df = delayed(pd.read_parquet)(path='foo-pyarrow', engine="pyarrow")

ddf = dd.from_delayed(delayed_df)

ddf.compute()
```

I also came across [this comment](https://github.com/dask/dask/issues/7512#issuecomment-852381361), which looks similar, but the [suggested workaround](https://github.com/dask/dask/issues/7512#issuecomment-889640208) doesn't seem to help in our example. :/

Thanks for the reporting this @th0ger - My understading is that dask is able to round-trip the partitioned dataset, but fails when the data was written directly by pyarrow (please do correct me if I am wrong). It is likely that this is a bug in Dask and I will be happy to help investigate.

Some other notes:
- The pyarrow-dataset engine was implemented before the pyarrow.dataset API had write support, so it is no suprise that this case is not tested in CI.
- The `pyarrow.dataset` API is listed as "experimental" by pyarrow, but they have also started deprecating the legacy API (so we can't really rely on pyarrow-legacy for stability here).  Since the two APIs are quite different, we recently decided [to deprecate](https://github.com/dask/dask/pull/8758) the "pyarrow-legacy" engine in Dask.  I am hoping this will allow us to simplify a lot of code and focus our efforts moving forward.
- As you already discovered, you cannot pass `dataset={"partitioning": "hive"}` to `read_parquet` yet. I don't think there is any reason for this ugliness in newer versions of dask/pyarrow, so #8765 will likely resolve that issue soon.

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
@@ -1046,6 +1046,7 @@ def _create_dd_meta(cls, dataset_info):
         categories = dataset_info["categories"]
         partition_obj = dataset_info["partitions"]
         partitions = dataset_info["partition_names"]
+        physical_column_names = dataset_info.get("physical_schema", schema).names
         columns = None
 
         # Set index and column names using
@@ -1068,7 +1069,7 @@ def _create_dd_meta(cls, dataset_info):
         else:
             # No pandas metadata implies no index, unless selected by the user
             index_names = []
-            column_names = dataset_info.get("physical_schema", schema).names
+            column_names = physical_column_names
             storage_name_mapping = {k: k for k in column_names}
             column_index_names = [None]
         if index is None and index_names:
@@ -1078,7 +1079,7 @@ def _create_dd_meta(cls, dataset_info):
         # 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]
+            _partitions = [p for p in partitions if p not in physical_column_names]
             if not _partitions:
                 partitions = []
                 dataset_info["partitions"] = None
@@ -1088,7 +1089,9 @@ def _create_dd_meta(cls, dataset_info):
                 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)
+                    "physical columns: {} | partitions: {}".format(
+                        physical_column_names, partitions
+                    )
                 )
 
         column_names, index_names = _normalize_index_columns(


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
@@ -3945,3 +3945,49 @@ def test_custom_filename_with_partition(tmpdir, engine):
     assert_eq(
         pdf, actual, check_index=False, check_dtype=False, check_categorical=False
     )
+
+
+@PYARROW_MARK
+@pytest.mark.skipif(
+    pa_version < parse_version("5.0"),
+    reason="pyarrow write_dataset was added in version 5.0",
+)
+def test_roundtrip_partitioned_pyarrow_dataset(tmpdir, engine):
+    # See: https://github.com/dask/dask/issues/8650
+
+    import pyarrow.parquet as pq
+    from pyarrow.dataset import HivePartitioning, write_dataset
+
+    # Sample data
+    df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})
+
+    # Write partitioned dataset with dask
+    dask_path = tmpdir.mkdir("foo-dask")
+    ddf = dd.from_pandas(df, npartitions=2)
+    ddf.to_parquet(dask_path, engine=engine, partition_on=["col1"], write_index=False)
+
+    # Write partitioned dataset with pyarrow
+    pa_path = tmpdir.mkdir("foo-pyarrow")
+    table = pa.Table.from_pandas(df)
+    write_dataset(
+        data=table,
+        base_dir=pa_path,
+        basename_template="part.{i}.parquet",
+        format="parquet",
+        partitioning=HivePartitioning(pa.schema([("col1", pa.int32())])),
+    )
+
+    # Define simple function to ensure results should
+    # be comparable (same column and row order)
+    def _prep(x):
+        return x.sort_values("col2")[["col1", "col2"]]
+
+    # Check that reading dask-written data is the same for pyarrow and dask
+    df_read_dask = dd.read_parquet(dask_path, engine=engine)
+    df_read_pa = pq.read_table(dask_path).to_pandas()
+    assert_eq(_prep(df_read_dask), _prep(df_read_pa), check_index=False)
+
+    # Check that reading pyarrow-written data is the same for pyarrow and dask
+    df_read_dask = dd.read_parquet(pa_path, engine=engine)
+    df_read_pa = pq.read_table(pa_path).to_pandas()
+    assert_eq(_prep(df_read_dask), _prep(df_read_pa), check_index=False)