Repository: dask/dask

Issue:
Support integer names in parquet
Currently we fail if we try to roundtrip a dataframe with integer column names to and from parquet

```python
import dask.dataframe as dd, pandas as pd
df = pd.DataFrame({0: range(5), 1: range(5)})
ddf = dd.from_pandas(df, npartitions=1)
ddf.to_parquet("foo.parquet")
dd.read_parquet("foo.parquet").head()
```

```python-traceback
KeyError                           Traceback (most recent call last)
<ipython-input-9-d16aa5855966> in <module>
----> 1 dd.read_parquet("foo.parquet").head()

~/workspace/dask/dask/dataframe/core.py in head(self, n, npartitions, compute)
   1068             Whether to compute the result, default is True.
   1069         """
-> 1070         return self._head(n=n, npartitions=npartitions, compute=compute, safe=True)
   1071 
   1072     def _head(self, n, npartitions, compute, safe):

~/workspace/dask/dask/dataframe/core.py in _head(self, n, npartitions, compute, safe)
   1101 
   1102         if compute:
-> 1103             result = result.compute()
   1104         return result
   1105 

~/workspace/dask/dask/base.py in compute(self, **kwargs)
    284         dask.base.compute
    285         """
--> 286         (result,) = compute(self, traverse=False, **kwargs)
    287         return result
    288 

~/workspace/dask/dask/base.py in compute(*args, **kwargs)
    566         postcomputes.append(x.__dask_postcompute__())
    567 
--> 568     results = schedule(dsk, keys, **kwargs)
    569     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    570 

~/workspace/dask/dask/threaded.py in get(dsk, result, cache, num_workers, pool, **kwargs)
     77             pool = MultiprocessingPoolExecutor(pool)
     78 
---> 79     results = get_async(
     80         pool.submit,
     81         pool._max_workers,

~/workspace/dask/dask/local.py in get_async(submit, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, chunksize, **kwargs)
    512                             _execute_task(task, data)  # Re-execute locally
    513                         else:
--> 514                             raise_exception(exc, tb)
    515                     res, worker_id = loads(res_info)
    516                     state["cache"][key] = res

~/workspace/dask/dask/local.py in reraise(exc, tb)
    323     if exc.__traceback__ is not tb:
    324         raise exc.with_traceback(tb)
--> 325     raise exc
    326 
    327 

~/workspace/dask/dask/local.py in execute_task(key, task_info, dumps, loads, get_id, pack_exception)
    221     try:
    222         task, data = loads(task_info)
--> 223         result = _execute_task(task, data)
    224         id = get_id()
    225         result = dumps((result, id))

~/workspace/dask/dask/core.py in _execute_task(arg, cache, dsk)
    119         # temporaries by their reference count and can execute certain
    120         # operations in-place.
--> 121         return func(*(_execute_task(a, cache) for a in args))
    122     elif not ishashable(arg):
    123         return arg

~/workspace/dask/dask/optimization.py in __call__(self, *args)
    967         if not len(args) == len(self.inkeys):
    968             raise ValueError("Expected %d args, got %d" % (len(self.inkeys), len(args)))
--> 969         return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))
    970 
    971     def __reduce__(self):

~/workspace/dask/dask/core.py in get(dsk, out, cache)
    149     for key in toposort(dsk):
    150         task = dsk[key]
--> 151         result = _execute_task(task, cache)
    152         cache[key] = result
    153     result = _execute_task(out, cache)

~/workspace/dask/dask/core.py in _execute_task(arg, cache, dsk)
    119         # temporaries by their reference count and can execute certain
    120         # operations in-place.
--> 121         return func(*(_execute_task(a, cache) for a in args))
    122     elif not ishashable(arg):
    123         return arg

~/workspace/dask/dask/dataframe/io/parquet/core.py in __call__(self, part)
     85             part = [part]
     86 
---> 87         return read_parquet_part(
     88             self.fs,
     89             self.engine,

~/workspace/dask/dask/dataframe/io/parquet/core.py in read_parquet_part(fs, engine, meta, part, columns, index, kwargs)
    406             # Part kwargs expected
    407             func = engine.read_partition
--> 408             dfs = [
    409                 func(fs, rg, columns.copy(), index, **toolz.merge(kwargs, kw))
    410                 for (rg, kw) in part

~/workspace/dask/dask/dataframe/io/parquet/core.py in <listcomp>(.0)
    407             func = engine.read_partition
    408             dfs = [
--> 409                 func(fs, rg, columns.copy(), index, **toolz.merge(kwargs, kw))
    410                 for (rg, kw) in part
    411             ]

~/workspace/dask/dask/dataframe/io/parquet/arrow.py in read_partition(cls, fs, pieces, columns, index, categories, partitions, filters, schema, **kwargs)
    723                 index = False
    724                 columns_and_parts = list(set(columns_and_parts) - set(df.index.names))
--> 725         df = df[list(columns_and_parts)]
    726 
    727         if index:

~/mambaforge/lib/python3.9/site-packages/pandas/core/frame.py in __getitem__(self, key)
   3459             if is_iterator(key):
   3460                 key = list(key)
-> 3461             indexer = self.loc._get_listlike_indexer(key, axis=1)[1]
   3462 
   3463         # take() does not accept boolean indexers

~/mambaforge/lib/python3.9/site-packages/pandas/core/indexing.py in _get_listlike_indexer(self, key, axis)
   1312             keyarr, indexer, new_indexer = ax._reindex_non_unique(keyarr)
   1313 
-> 1314         self._validate_read_indexer(keyarr, indexer, axis)
   1315 
   1316         if needs_i8_conversion(ax.dtype) or isinstance(

~/mambaforge/lib/python3.9/site-packages/pandas/core/indexing.py in _validate_read_indexer(self, key, indexer, axis)
   1372                 if use_interval_msg:
   1373                     key = list(key)
-> 1374                 raise KeyError(f"None of [{key}] are in the [{axis_name}]")
   1375 
   1376             not_found = list(ensure_index(key)[missing_mask.nonzero()[0]].unique())

KeyError: "None of [Index(['0', '1'], dtype='object')] are in the [columns]"
```

cc @rjzamora @jorisvandenbossche @martindurant 


Discussion:
Out of interest, have you tried with fastparquet?

Note that the parquet format itself requires utf8 string column names.

On August 8, 2021 8:59:26 AM EDT, Matthew Rocklin ***@***.***> wrote:
>Currently we fail if we try to roundtrip a dataframe with integer column names to and from parquet
>
>```python
>import dask.dataframe as dd, pandas as pd
>df = pd.DataFrame({0: range(5), 1: range(5)})
>ddf = dd.from_pandas(df, npartitions=1)
>ddf.to_parquet("foo.parquet")
>dd.read_parquet("foo.parquet").head()
>```
>
>```python-traceback
>KeyError                           Traceback (most recent call last)
><ipython-input-9-d16aa5855966> in <module>
>----> 1 dd.read_parquet("foo.parquet").head()
>
>~/workspace/dask/dask/dataframe/core.py in head(self, n, npartitions, compute)
>   1068             Whether to compute the result, default is True.
>   1069         """
>-> 1070         return self._head(n=n, npartitions=npartitions, compute=compute, safe=True)
>   1071 
>   1072     def _head(self, n, npartitions, compute, safe):
>
>~/workspace/dask/dask/dataframe/core.py in _head(self, n, npartitions, compute, safe)
>   1101 
>   1102         if compute:
>-> 1103             result = result.compute()
>   1104         return result
>   1105 
>
>~/workspace/dask/dask/base.py in compute(self, **kwargs)
>    284         dask.base.compute
>    285         """
>--> 286         (result,) = compute(self, traverse=False, **kwargs)
>    287         return result
>    288 
>
>~/workspace/dask/dask/base.py in compute(*args, **kwargs)
>    566         postcomputes.append(x.__dask_postcompute__())
>    567 
>--> 568     results = schedule(dsk, keys, **kwargs)
>    569     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
>    570 
>
>~/workspace/dask/dask/threaded.py in get(dsk, result, cache, num_workers, pool, **kwargs)
>     77             pool = MultiprocessingPoolExecutor(pool)
>     78 
>---> 79     results = get_async(
>     80         pool.submit,
>     81         pool._max_workers,
>
>~/workspace/dask/dask/local.py in get_async(submit, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, chunksize, **kwargs)
>    512                             _execute_task(task, data)  # Re-execute locally
>    513                         else:
>--> 514                             raise_exception(exc, tb)
>    515                     res, worker_id = loads(res_info)
>    516                     state["cache"][key] = res
>
>~/workspace/dask/dask/local.py in reraise(exc, tb)
>    323     if exc.__traceback__ is not tb:
>    324         raise exc.with_traceback(tb)
>--> 325     raise exc
>    326 
>    327 
>
>~/workspace/dask/dask/local.py in execute_task(key, task_info, dumps, loads, get_id, pack_exception)
>    221     try:
>    222         task, data = loads(task_info)
>--> 223         result = _execute_task(task, data)
>    224         id = get_id()
>    225         result = dumps((result, id))
>
>~/workspace/dask/dask/core.py in _execute_task(arg, cache, dsk)
>    119         # temporaries by their reference count and can execute certain
>    120         # operations in-place.
>--> 121         return func(*(_execute_task(a, cache) for a in args))
>    122     elif not ishashable(arg):
>    123         return arg
>
>~/workspace/dask/dask/optimization.py in __call__(self, *args)
>    967         if not len(args) == len(self.inkeys):
>    968             raise ValueError("Expected %d args, got %d" % (len(self.inkeys), len(args)))
>--> 969         return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))
>    970 
>    971     def __reduce__(self):
>
>~/workspace/dask/dask/core.py in get(dsk, out, cache)
>    149     for key in toposort(dsk):
>    150         task = dsk[key]
>--> 151         result = _execute_task(task, cache)
>    152         cache[key] = result
>    153     result = _execute_task(out, cache)
>
>~/workspace/dask/dask/core.py in _execute_task(arg, cache, dsk)
>    119         # temporaries by their reference count and can execute certain
>    120         # operations in-place.
>--> 121         return func(*(_execute_task(a, cache) for a in args))
>    122     elif not ishashable(arg):
>    123         return arg
>
>~/workspace/dask/dask/dataframe/io/parquet/core.py in __call__(self, part)
>     85             part = [part]
>     86 
>---> 87         return read_parquet_part(
>     88             self.fs,
>     89             self.engine,
>
>~/workspace/dask/dask/dataframe/io/parquet/core.py in read_parquet_part(fs, engine, meta, part, columns, index, kwargs)
>    406             # Part kwargs expected
>    407             func = engine.read_partition
>--> 408             dfs = [
>    409                 func(fs, rg, columns.copy(), index, **toolz.merge(kwargs, kw))
>    410                 for (rg, kw) in part
>
>~/workspace/dask/dask/dataframe/io/parquet/core.py in <listcomp>(.0)
>    407             func = engine.read_partition
>    408             dfs = [
>--> 409                 func(fs, rg, columns.copy(), index, **toolz.merge(kwargs, kw))
>    410                 for (rg, kw) in part
>    411             ]
>
>~/workspace/dask/dask/dataframe/io/parquet/arrow.py in read_partition(cls, fs, pieces, columns, index, categories, partitions, filters, schema, **kwargs)
>    723                 index = False
>    724                 columns_and_parts = list(set(columns_and_parts) - set(df.index.names))
>--> 725         df = df[list(columns_and_parts)]
>    726 
>    727         if index:
>
>~/mambaforge/lib/python3.9/site-packages/pandas/core/frame.py in __getitem__(self, key)
>   3459             if is_iterator(key):
>   3460                 key = list(key)
>-> 3461             indexer = self.loc._get_listlike_indexer(key, axis=1)[1]
>   3462 
>   3463         # take() does not accept boolean indexers
>
>~/mambaforge/lib/python3.9/site-packages/pandas/core/indexing.py in _get_listlike_indexer(self, key, axis)
>   1312             keyarr, indexer, new_indexer = ax._reindex_non_unique(keyarr)
>   1313 
>-> 1314         self._validate_read_indexer(keyarr, indexer, axis)
>   1315 
>   1316         if needs_i8_conversion(ax.dtype) or isinstance(
>
>~/mambaforge/lib/python3.9/site-packages/pandas/core/indexing.py in _validate_read_indexer(self, key, indexer, axis)
>   1372                 if use_interval_msg:
>   1373                     key = list(key)
>-> 1374                 raise KeyError(f"None of [{key}] are in the [{axis_name}]")
>   1375 
>   1376             not_found = list(ensure_index(key)[missing_mask.nonzero()[0]].unique())
>
>KeyError: "None of [Index(['0', '1'], dtype='object')] are in the [columns]"
>```
>
>cc @rjzamora @jorisvandenbossche @martindurant 
>
>-- 
>You are receiving this because you were mentioned.
>Reply to this email directly or view it on GitHub:
>https://github.com/dask/dask/issues/8010
-- 
Sent from my Android device with K-9 Mail. Please excuse my brevity.
I haven't yet tried with fastparquet.  

Yes, I was able to work around this with text column names, but we should probably improve the UX here.  Either we need to fail hard on any non-string column names, or we need to convert to strings everywhere.  My guess is that we're converting to strings as we save the data, but that maybe we're also saving something like a pandas index of columns in the metadata and that no longer aligns?
The problems lies in a mismatch between dask's `meta` and the actual DataFrame:

```
In [4]: ddf2 = dd.read_parquet("foo.parquet")

In [5]: ddf2.columns
Out[5]: Index(['0', '1'], dtype='object')

In [6]: pd.read_parquet("foo55.parquet/").columns
Out[6]: Index([0, 1], dtype='object')
```

Pyarrow seems to do some effort here in restoring the integer column names based on the information stored in the pandas metadata (whether that good or not ...), but dask doesn't do the same, creating this mismatch. 
There is a `_get_pyarrow_dtypes` util that creates a dict of column name -> dtypes (from which `meta` is constructed afterwards) from the pyarrow schema. I suppose it's this function that needs to be updated to ensure the two paths match.
I think fastparquet doesn't try to restore the integer column names, so won't have this problem:

```
In [51]: pd.read_parquet("foo55.parquet/", engine="pyarrow").columns
Out[51]: Index([0, 1], dtype='object')

In [52]: pd.read_parquet("foo55.parquet/", engine="fastparquet").columns
Out[52]: Index(['0', '1'], dtype='object')
```

> but we should probably improve the UX here. Either we need to fail hard on any non-string column names, or we need to convert to strings everywhere.

Yes, note that pandas itself actually explicitly fails on this (but this check happens in the pandas layer, not in pyarrow, which dask uses directly in its `to_parquet`):

```
In [8]: df.to_parquet("foo_pandas.parquet")
...
~/miniconda3/envs/geo/lib/python3.9/site-packages/pandas/io/parquet.py in validate_dataframe(df)
    106         else:
    107             if df.columns.inferred_type not in {"string", "empty"}:
--> 108                 raise ValueError("parquet must have string column names")
    109 
    110         # index level names must be strings

ValueError: parquet must have string column names
```

Personally, I think raising an error for this is fine. But it should also be possible to fix the roundtrip with integers (either as I mentioned above by updating `_get_pyarrow_types`, or either by automatically casting to string names before converting to arrow / writing to parquet)

Resolving patch:
diff --git a/dask/dataframe/io/parquet/core.py b/dask/dataframe/io/parquet/core.py
--- a/dask/dataframe/io/parquet/core.py
+++ b/dask/dataframe/io/parquet/core.py
@@ -694,6 +694,9 @@ def to_parquet(
             "columns=%s" % (str(partition_on), str(list(df.columns)))
         )
 
+    if df.columns.inferred_type not in {"string", "empty"}:
+        raise ValueError("parquet doesn't support non-string column names")
+
     if isinstance(engine, str):
         engine = get_engine(engine, bool(kwargs))
 


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
@@ -3463,6 +3463,13 @@ def test_to_parquet_overwrite_raises(tmpdir, engine):
         dd.to_parquet(ddf, tmpdir, engine=engine, append=True, overwrite=True)
 
 
+def test_to_parquet_errors_non_string_column_names(tmpdir, engine):
+    df = pd.DataFrame({"x": range(10), 1: range(10)})
+    ddf = dd.from_pandas(df, npartitions=2)
+    with pytest.raises(ValueError, match="non-string column names"):
+        ddf.to_parquet(str(tmpdir.join("temp")), engine=engine)
+
+
 def test_dir_filter(tmpdir, engine):
     # github #6898
     df = pd.DataFrame.from_dict(