Repository: dask/dask

Issue:
dd.read_parquet no longer include index columns
When calling dask.dataframe.read_parquet(pq_file, index=False), dask 0.16.0 and earlier will include the MultiIndex columns as regular columns in the returned DataFrame, but dask 0.16.1 does not include them. The difference is due to dask, not fastparquet.

Here's an example:

```
from dask import dataframe as dd
import fastparquet as fpq
import pandas as pd
import numpy as np
df = pd.DataFrame(np.random.randn(10, 2), columns = ['a','b'], index=pd.MultiIndex.from_arrays([np.arange(10), np.arange(10) + 1], names=['x0','x1']))
df

fpq.write('/tmp/test.pq', df, write_index=True)

ddf = dd.read_parquet('/tmp/test.pq', index=False)
ddf
```

The df is like:

>               a         b
> x0 x1                    
> 0  1   1.088650  0.458507
> 1  2   1.271835 -0.840432
> 2  3  -0.979103 -0.240490
> 3  4   0.636582 -0.325154
> 4  5  -1.227770  0.602984
> 5  6  -0.707955  0.804685
> 6  7  -1.512000 -0.494987
> 7  8  -1.074822 -0.126040
> 8  9  -0.445532  0.682687
> 9  10  1.777594  1.127077
> 

For 0.16.0, the final ddf print as:

> Dask DataFrame Structure:
>                   x0     x1        a        b
> npartitions=1                                
>                int64  int64  float64  float64
>                  ...    ...      ...      ...
> Dask Name: read-parquet, 1 tasks
> 

For 0.16.1, the final ddf print as (index columns x0 and x1 are not included):

> Dask DataFrame Structure:
>                      a        b
> npartitions=1                  
>                float64  float64
>                    ...      ...
> Dask Name: read-parquet, 1 tasks
> 

Versions of pkgs:

> In [5]: dask.__version__, fpq.__version__, pd.__version__, np.__version__
> Out[5]: (u'0.16.0', u'0.1.3', u'0.21.0', '1.13.3')
> 

I'd think the older version's behavior is more convenient and intuitive, isn't it?

Thanks



Discussion:
This is a break of API behavior. Is it planned, or is it an accident?
Does this happen for both fastparquet and pyarrow engines?
Sorry for the slow replay.

I suspect this was in https://github.com/dask/dask/pull/2973, and was an unintentional change.

Personally, I don't find reading the indexes into regular columns intuitive, but perhaps better than dropping entirely :)
I disagree, @TomAugspurger : the data is stored as column(s), and I would take the metadata as just a hint that they should be made the index of the newly-read data-frame. If the user wishes to prevent setting the index, the data should be read as ordinary columns.
Thanks for the reply.

I only used fastparquet. The features of pyarrow seem to be more limited so I didn't try.

I think the problem here is API consistency. If I do `fpq.write(..., write_index=True)`, then I would expect to have some way to get the index columns back in `dd.read_parquet`. I understand that MultiIndex is not supported in parquet for now, so I have to do `dd.read_parquet(..., index=False)`. If the 'index=False' here means that I don't want index columns, then what's the API in read_parquet that let me get the index column data?

Yes, I agree @xwang777  - I'll look into reverting this change for `index=False`.
Interestingly, if you explicitly include the list of columns and include the index name(s) there, you do get them back:

```
[ ] dd.read_parquet('/tmp/test.pq', index=False, columns=['x0', 'x1', 'a', 'b'])
   x0  x1         a         b
0   0   1  1.231914 -1.044352
1   1   2 -0.499282  0.436183
2   2   3  0.552027 -0.265289
3   3   4  0.345140  0.485580
4   4   5  0.937267 -0.604615
5   5   6 -1.707770  1.143857
6   6   7  1.103741  0.411966
7   7   8  0.648006 -2.070996
8   8   9  0.900816 -1.107321
9   9  10 -0.619174  0.095268
```
The following fixes it. @jcrist @TomAugspurger - does this seem like the right thing to do?
```
--- a/dask/dataframe/io/parquet.py
+++ b/dask/dataframe/io/parquet.py
@@ -234,6 +234,7 @@ def _read_fastparquet(fs, fs_token, paths, columns=None, filters=None,
         index_names, column_names, storage_name_mapping, column_index_names = (
             _parse_pandas_metadata(json.loads(pandas_md[0]))
         )
+        column_names.extend([n for n in index_names if n != (index or "")])
     else:
         raise ValueError("File has multiple entries for 'pandas' metadata")
```
i.e., if `index=False` or `index='a'` (one of the columns or putative indexes), the data from what were previously indexes get loaded as normal columns.
Note that if you specify an `index=` with any name that is not a column/index at all, you get the standard range index (0, 1, ..) but with that name applied. I don't know if that is intentional.
I would push it into `_normalize_index` / `_normalize_columns`, which are for reconciling user-input with the data. Putting it there will also fix things for Pyarrow, which I assume suffers from the same bug?

I guess that `_normalize_columns` will need to take `index, index_names` too now, and if `index=False`, then `index_names` will all be included in the output of `_normalize_columns`.
```
In [9]: dd.read_parquet('out.parq', index='a', engine='pyarrow').compute()
Out[9]:
                  b
a
 1.231914 -1.044352
-0.499282  0.436183
 0.552027 -0.265289
 0.345140  0.485580
 0.937267 -0.604615
-1.707770  1.143857
 1.103741  0.411966
 0.648006 -2.070996
 0.900816 -1.107321
-0.619174  0.095268

In [10]: dd.read_parquet('out.parq', index='blah', engine='pyarrow').compute()
Out[10]:
              a         b
x0 x1
0  1   1.231914 -1.044352
1  2  -0.499282  0.436183
2  3   0.552027 -0.265289
3  4   0.345140  0.485580
4  5   0.937267 -0.604615
5  6  -1.707770  1.143857
6  7   1.103741  0.411966
7  8   0.648006 -2.070996
8  9   0.900816 -1.107321
9  10 -0.619174  0.095268
```
You get the same latter result whether index=None, 'blah', 'x0', ['x0'], []...
I'm not sure that the behaviour should be.
In fact, for the case that works, I thought that dask dataframes could not handle multi-indexes anyway.

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
@@ -143,54 +143,71 @@ def _meta_from_dtypes(to_read_columns, file_dtypes, index_cols,
     return df
 
 
-def _normalize_columns(user_columns, data_columns):
-    """Normalize user- and file-provided column names
+def _normalize_index_columns(user_columns, data_columns, user_index, data_index):
+    """Normalize user and file-provided column and index names
 
     Parameters
     ----------
     user_columns : None, str or list of str
     data_columns : list of str
+    user_index : None, str, or list of str
+    data_index : list of str
 
     Returns
     -------
     column_names : list of str
+    index_names : list of str
     out_type : {pd.Series, pd.DataFrame}
     """
-    out_type = DataFrame
-    if user_columns is not None:
-        if isinstance(user_columns, string_types):
-            columns = [user_columns]
-            out_type = Series
-        else:
-            columns = list(user_columns)
-        column_names = columns
-    else:
-        column_names = list(data_columns)
-    return column_names, out_type
+    specified_columns = user_columns is not None
+    specified_index = user_index is not None
 
+    out_type = DataFrame
 
-def _normalize_index(user_index, data_index):
-    """Normalize user- and file-provided index names
-
-    Parameters
-    ----------
-    user_index : None, str, or list of str
-    data_index : list of str
-
-    Returns
-    -------
-    index_names : list of str
-    """
-    if user_index is not None:
-        if user_index is False:
-            index_names = []
-        elif isinstance(user_index, string_types):
-            index_names = [user_index]
-        else:
-            index_names = list(user_index)
+    if user_columns is None:
+        user_columns = list(data_columns)
+    elif isinstance(user_columns, string_types):
+        user_columns = [user_columns]
+        out_type = Series
+    else:
+        user_columns = list(user_columns)
+
+    if user_index is None:
+        user_index = data_index
+    elif user_index is False:
+        # When index is False, use no index and all fields should be treated as
+        # columns (unless `columns` provided).
+        user_index = []
+        data_columns = data_index + data_columns
+    elif isinstance(user_index, string_types):
+        user_index = [user_index]
+    else:
+        user_index = list(user_index)
+
+    if specified_index and not specified_columns:
+        # Only `index` provided. Use specified index, and all column fields
+        # that weren't specified as indices
+        index_names = user_index
+        column_names = [x for x in data_columns if x not in index_names]
+    elif specified_columns and not specified_index:
+        # Only `columns` provided. Use specified columns, and all index fields
+        # that weren't specified as columns
+        column_names = user_columns
+        index_names = [x for x in data_index if x not in column_names]
+    elif specified_index and specified_columns:
+        # Both `index` and `columns` provided. Use as specified, but error if
+        # they intersect.
+        column_names = user_columns
+        index_names = user_index
+        if set(column_names).intersection(index_names):
+            raise ValueError("Specified index and column names must not "
+                             "intersect")
     else:
+        # Use default columns and index from the metadata
+        column_names = data_columns
         index_names = data_index
-    return index_names
+
+    return column_names, index_names, out_type
 
 
 # ----------------------------------------------------------------------
@@ -242,8 +259,8 @@ def _read_fastparquet(fs, fs_token, paths, columns=None, filters=None,
     if filters is None:
         filters = []
 
-    column_names, out_type = _normalize_columns(columns, column_names)
-    index_names = _normalize_index(index, index_names)
+    column_names, index_names, out_type = _normalize_index_columns(columns, column_names,
+                                                                   index, index_names)
 
     if categories is None:
         categories = pf.categories
@@ -268,7 +285,7 @@ def _read_fastparquet(fs, fs_token, paths, columns=None, filters=None,
     dtypes = {storage_name_mapping.get(k, k): v for k, v in dtypes.items()}
 
     meta = _meta_from_dtypes(all_columns, dtypes, index_names, [None])
-    # fastparquet / dask don't handle multiindex
+    # fastparquet doesn't handle multiindex
     if len(index_names) > 1:
         raise ValueError("Cannot read DataFrame with MultiIndex.")
     elif len(index_names) == 0:
@@ -528,8 +545,8 @@ 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, out_type = _normalize_columns(columns, column_names)
-    index_names = _normalize_index(index, index_names)
+    column_names, index_names, out_type = _normalize_index_columns(columns, column_names,
+                                                                   index, index_names)
 
     all_columns = index_names + column_names
 
@@ -550,7 +567,7 @@ def _read_pyarrow(fs, fs_token, paths, columns=None, filters=None,
             (task_name, i): (_read_pyarrow_parquet_piece,
                              fs,
                              piece,
-                             all_columns,
+                             column_names,
                              index_names,
                              out_type == Series,
                              dataset.partitions)
@@ -576,17 +593,27 @@ def _get_pyarrow_dtypes(schema):
 def _read_pyarrow_parquet_piece(fs, piece, columns, index_cols, is_series,
                                 partitions):
     with fs.open(piece.path, mode='rb') as f:
-        table = piece.read(columns=columns, partitions=partitions,
+        table = piece.read(columns=index_cols + columns,
+                           partitions=partitions,
                            use_pandas_metadata=True,
                            file=f)
     df = table.to_pandas()
-    if (index_cols and df.index.name is None and
-            len(df.columns.intersection(index_cols))):
+    has_index = not isinstance(df.index, pd.RangeIndex)
+
+    if not has_index and index_cols:
         # Index should be set, but it isn't
         df = df.set_index(index_cols)
-    elif not index_cols and df.index.name is not None:
-        # Index shouldn't be set, but it is
+    elif has_index and df.index.names != index_cols:
+        # Index is set, but isn't correct
+        # This can happen when reading in not every column in a multi-index
         df = df.reset_index(drop=False)
+        if index_cols:
+            df = df.set_index(index_cols)
+        drop = list(set(df.columns).difference(columns))
+        if drop:
+            df = df.drop(drop, axis=1)
+        # Ensure proper ordering
+        df = df.reindex(columns=columns, copy=False)
 
     if is_series:
         return df[df.columns[0]]
@@ -725,17 +752,21 @@ def read_parquet(path, columns=None, filters=None, categories=None, index=None,
         filesystems. To read from multiple files you can pass a globstring or a
         list of paths, with the caveat that they must all have the same
         protocol.
-    columns: list or None
-        List of column names to load
-    filters: list
+    columns : string, list or None (default)
+        Field name(s) to read in as columns in the output. By default all
+        non-index fields will be read (as determined by the pandas parquet
+        metadata, if present). Provide a single field name instead of a list to
+        read in the data as a Series.
+    filters : list
         List of filters to apply, like ``[('x', '>', 0), ...]``. This implements
         row-group (partition) -level filtering only, i.e., to prevent the
         loading of some chunks of the data, and only if relevant statistics
         have been included in the metadata.
-    index: string or None (default) or False
-        Name of index column to use if that column is sorted;
-        False to force dask to not use any column as the index
-    categories: list, dict or None
+    index : string, list, False or None (default)
+        Field name(s) to use as the output frame index. By default will be
+        inferred from the pandas parquet file metadata (if present). Use False
+        to read all fields as columns.
+    categories : list, dict or None
         For any fields listed here, if the parquet encoding is Dictionary,
         the column will be created with dtype category. Use only if it is
         guaranteed that the column is encoded as dictionary in all row-groups.


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
@@ -163,29 +163,86 @@ def test_read_list(tmpdir, write_engine, read_engine):
 
 
 @write_read_engines_xfail
-def test_auto_add_index(tmpdir, write_engine, read_engine):
+def test_columns_index(tmpdir, write_engine, read_engine):
     fn = str(tmpdir)
     ddf.to_parquet(fn, engine=write_engine)
-    ddf2 = dd.read_parquet(fn, columns=['x'], index='myindex', engine=read_engine)
-    assert_eq(df[['x']], ddf2)
+    df2 = df.reset_index()
+
+    assert_eq(dd.read_parquet(fn, columns=[], engine=read_engine), df[[]])
+    assert_eq(dd.read_parquet(fn, columns=['x'], engine=read_engine), df[['x']])
+    assert_eq(dd.read_parquet(fn, index='myindex', columns=['x'], engine=read_engine), df[['x']])
+    assert_eq(dd.read_parquet(fn, index='myindex', columns=['x', 'y'], engine=read_engine), df)
+
+    assert_eq(dd.read_parquet(fn, index=False, engine=read_engine),
+              df2, check_index=False)
+    assert_eq(dd.read_parquet(fn, index=False, columns=['x', 'y'], engine=read_engine),
+              df, check_index=False)
+    assert_eq(dd.read_parquet(fn, index=False, columns=['myindex', 'x'], engine=read_engine),
+              df2[['myindex', 'x']], check_index=False)
+
+
+def test_columns_index_with_multi_index(tmpdir, engine):
+    fn = os.path.join(str(tmpdir), 'test.parquet')
+    index = pd.MultiIndex.from_arrays([np.arange(10), np.arange(10) + 1],
+                                      names=['x0', 'x1'])
+    df = pd.DataFrame(np.random.randn(10, 2), columns=['a', 'b'], index=index)
+    df2 = df.reset_index(drop=False)
+
+    if engine == 'fastparquet':
+        fastparquet.write(fn, df, write_index=True)
+
+        # fastparquet doesn't support multi-index
+        with pytest.raises(ValueError):
+            ddf = dd.read_parquet(fn, engine=engine)
+    else:
+        import pyarrow as pa
+        pq.write_table(pa.Table.from_pandas(df), fn)
 
+        # Pyarrow supports multi-index reads
+        ddf = dd.read_parquet(fn, engine=engine)
+        assert_eq(ddf, df)
 
-@write_read_engines_xfail
-def test_index_column_false_index(tmpdir, write_engine, read_engine):
-    fn = str(tmpdir)
-    ddf.to_parquet(fn, engine=write_engine)
-    ddf2 = dd.read_parquet(fn, columns=['myindex'], index=False, engine=read_engine)
-    assert_eq(pd.DataFrame(df.index), ddf2, check_index=False)
+        d = dd.read_parquet(fn, columns='a', engine=engine)
+        assert_eq(d, df['a'])
 
+        d = dd.read_parquet(fn, index=['a', 'b'], columns=['x0', 'x1'], engine=engine)
+        assert_eq(d, df2.set_index(['a', 'b'])[['x0', 'x1']])
 
-@pytest.mark.parametrize("columns", [['myindex'], []])
-@pytest.mark.parametrize("index", ['myindex', None])
-@write_read_engines_xfail
-def test_columns_index(tmpdir, write_engine, read_engine, columns, index):
-    fn = str(tmpdir)
-    ddf.to_parquet(fn, engine=write_engine)
-    ddf2 = dd.read_parquet(fn, columns=columns, index=index, engine=read_engine)
-    assert_eq(df[[]], ddf2)
+    # Just index
+    d = dd.read_parquet(fn, index=False, engine=engine)
+    assert_eq(d, df2)
+
+    d = dd.read_parquet(fn, index=['a'], engine=engine)
+    assert_eq(d, df2.set_index('a')[['b']])
+
+    d = dd.read_parquet(fn, index=['x0'], engine=engine)
+    assert_eq(d, df2.set_index('x0')[['a', 'b']])
+
+    # Just columns
+    d = dd.read_parquet(fn, columns=['x0', 'a'], engine=engine)
+    assert_eq(d, df2.set_index('x1')[['x0', 'a']])
+
+    # Both index and columns
+    d = dd.read_parquet(fn, index=False, columns=['x0', 'b'], engine=engine)
+    assert_eq(d, df2[['x0', 'b']])
+
+    for index in ['x1', 'b']:
+        d = dd.read_parquet(fn, index=index, columns=['x0', 'a'], engine=engine)
+        assert_eq(d, df2.set_index(index)[['x0', 'a']])
+
+    # Columns and index intersect
+    for index in ['a', 'x0']:
+        with pytest.raises(ValueError):
+            d = dd.read_parquet(fn, index=index, columns=['x0', 'a'], engine=engine)
+
+    # Series output
+    for ind, col, sol_df in [(None, 'x0', df2.set_index('x1')),
+                             (False, 'b', df2),
+                             (False, 'x0', df2),
+                             ('a', 'x0', df2.set_index('a')),
+                             ('a', 'b', df2.set_index('a'))]:
+        d = dd.read_parquet(fn, index=ind, columns=col, engine=engine)
+        assert_eq(d, sol_df[col])
 
 
 @write_read_engines_xfail