Repository: modin-project/modin
Pull request: Fix parquet file reading when file was written with pandas Index
* Resolves #1366
* Add a check for the Index columns from the pandas metadata maintained
  by parquet.
* Do not assign work for the columns that are actually pandas Index
  objects.
* Add tests for this behavior.

Also test for MultiIndex object

Signed-off-by: Devin Petersohn <devin.petersohn@gmail.com>

<!--
Thank you for your contribution! 
Please review the contributing docs: https://modin.readthedocs.io/en/latest/contributing.html
if you have questions about contributing.
-->

## What do these changes do?

<!-- Please give a short brief about these changes. -->

- [x] passes `flake8 modin`
- [x] passes `black --check modin`
- [x] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [x] Resolves #1366 <!-- issue must be created for each patch -->
- [x] tests added and passing


Problem statement:
KeyError(f"None of [{key}] are in the [{axis_name}]")
### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
- **Modin version** (`modin.__version__`): 0.7.2
- **Python version**: 3.7
- **Code we can use to reproduce**:
```
import pandas as pd
import modin.pandas as mpd

df = pd.DataFrame({'idx': [1,4,5], 'A': [4,5,6], 'B':['a','b','c']})
df = df.set_index("idx")
# write out parquet with index using pandas
df.to_parquet("./data/tmp.parquet")
# read the same parquet using modin.pandas
mdf = mpd.read_parquet("./data/tmp.parquet")
```
returns 
>"None of [Index(['idx'], dtype='object')] are in the [columns]"
<!--
You can obtain the Modin version with

python -c "import modin; print(modin.__version__)"
-->

### Describe the problem
<!-- Describe the problem clearly here. -->
Parquet files that are written using Pandas, throws this error when read using modin's `read_parquet()`

### Source code / logs
<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->

> KeyError: "None of [Index(['idx'], dtype='object')] are in the [columns]"

Hints from issue discussion:
Hi @jadhosn thanks for the report!

This looks like a bug specific to setting the index, I will dig into this, thanks again for the report!
@devin-petersohn thank you! 
In the meantime, which of the following workarounds would you recommend to use to still gain an efficiency boost using modin:
* Write out the original parquet file with `idx` as a column, read the parquet file using modin, then `set_index` using modin
```
import pandas as pd
import modin.pandas as mpd

df = pd.DataFrame({'idx': [1,4,5], 'A': [4,5,6], 'B':['a','b','c']})
# write out pandas df with idx as column
df.to_parquet("./data/tmp.parquet")

# read the parquet file using modin
mdf = mpd.read_parquet("./data/tmp.parquet")

# set index
mdf = mdf.set_index("idx")
```
* Read the indexed parquet file using pandas, transform the pandas dataframe into a modin dataframe, then proceed with workflow
```
import pandas as pd
import modin.pandas as mpd

df = pd.DataFrame({'idx': [1,4,5], 'A': [4,5,6], 'B':['a','b','c']})
df = df.set_index("idx")

# write out pandas df with idx as column
df.to_parquet("./data/tmp.parquet")

# read the parquet file using pandas
df = pd.read_parquet("./data/tmp.parquet") 

# transform into modin dataframe
mdf = mpd.DataFrame(df)
```

Since I'm merging large data frames based on the index, would you recommend modin's `set_index` or `DataFrame` in terms of efficiency gains? 
`set_index` is probably the way to go, `DataFrame` will take 2-3 seconds per GB of data because it has to partition and communicate the data. `set_index` should take significantly less time. Let me know if you run into issues with this workaround!

Commit message: Fix parquet file reading when file was written with pandas Index

* Resolves #1366
* Add a check for the Index columns from the pandas metadata maintained
  by parquet.
* Do not assign work for the columns that are actually pandas Index
  objects.
* Add tests for this behavior.
* Also test for MultiIndex object

Signed-off-by: Devin Petersohn <devin.petersohn@gmail.com>

Patch to review:
diff --git a/modin/backends/pandas/parsers.py b/modin/backends/pandas/parsers.py
index caf6ca4b20b..d59716ec436 100644
--- a/modin/backends/pandas/parsers.py
+++ b/modin/backends/pandas/parsers.py
@@ -141,18 +141,21 @@ def parse(fname, **kwargs):
 class PandasParquetParser(PandasParser):
     @staticmethod
     def parse(fname, **kwargs):
-        import pyarrow.parquet as pq
-
         num_splits = kwargs.pop("num_splits", None)
         columns = kwargs.get("columns", None)
         if num_splits is None:
             return pandas.read_parquet(fname, **kwargs)
         kwargs["use_pandas_metadata"] = True
-        df = pq.read_table(fname, **kwargs).to_pandas()
+        df = pandas.read_parquet(fname, **kwargs)
+        if isinstance(df.index, pandas.RangeIndex):
+            idx = len(df.index)
+        else:
+            idx = df.index
+        columns = [c for c in columns if c not in df.index.names and c in df.columns]
         if columns is not None:
             df = df[columns]
         # Append the length of the index here to build it externally
-        return _split_result_for_readers(0, num_splits, df) + [len(df.index), df.dtypes]
+        return _split_result_for_readers(0, num_splits, df) + [idx, df.dtypes]
 
 
 class PandasHDFParser(PandasParser):  # pragma: no cover
diff --git a/modin/engines/base/io/column_stores/column_store_reader.py b/modin/engines/base/io/column_stores/column_store_reader.py
index 4cac352431a..905dd0f126a 100644
--- a/modin/engines/base/io/column_stores/column_store_reader.py
+++ b/modin/engines/base/io/column_stores/column_store_reader.py
@@ -60,7 +60,11 @@ def build_index(cls, partition_ids):
         from modin.pandas import DEFAULT_NPARTITIONS
 
         index_len = cls.materialize(partition_ids[-2][0])
-        index = pandas.RangeIndex(index_len)
+        if isinstance(index_len, int):
+            index = pandas.RangeIndex(index_len)
+        else:
+            index = index_len
+            index_len = len(index)
         index_chunksize = compute_chunksize(
             pandas.DataFrame(index=index), DEFAULT_NPARTITIONS, axis=0
         )
diff --git a/modin/engines/base/io/column_stores/parquet_reader.py b/modin/engines/base/io/column_stores/parquet_reader.py
index 8599833fe2a..f5586327345 100644
--- a/modin/engines/base/io/column_stores/parquet_reader.py
+++ b/modin/engines/base/io/column_stores/parquet_reader.py
@@ -66,7 +66,18 @@ def read(cls, path, engine, columns, **kwargs):
                 pd = ParquetDataset(path)
                 column_names = pd.schema.names
             else:
-                pf = ParquetFile(path)
-                column_names = pf.metadata.schema.names
+                meta = ParquetFile(path).metadata
+                column_names = meta.schema.names
+                # This is how we convert the metadata from pyarrow to a python
+                # dictionary, from which we then get the index columns.
+                # We use these to filter out from the columns in the metadata since
+                # the pyarrow storage has no concept of row labels/index.
+                # This ensures that our metadata lines up with the partitions without
+                # extra communication steps once we `have done all the remote
+                # computation.
+                index_columns = eval(
+                    meta.metadata[b"pandas"].replace(b"null", b"None")
+                ).get("index_columns", [])
+                column_names = [c for c in column_names if c not in index_columns]
             columns = [name for name in column_names if not PQ_INDEX_REGEX.match(name)]
         return cls.build_query_compiler(path, columns, **kwargs)
diff --git a/modin/pandas/test/test_io.py b/modin/pandas/test/test_io.py
index c1c21944216..d50923f207d 100644
--- a/modin/pandas/test/test_io.py
+++ b/modin/pandas/test/test_io.py
@@ -428,6 +428,32 @@ def test_from_parquet_partitioned_columns_with_columns(make_parquet_file):
     df_equals(modin_df, pandas_df)
 
 
+def test_from_parquet_pandas_index():
+    # Ensure modin can read parquet files written by pandas with a non-RangeIndex object
+    pandas_df = pandas.DataFrame(
+        {
+            "idx": np.random.randint(0, 100_000, size=2000),
+            "A": np.random.randint(0, 100_000, size=2000),
+            "B": ["a", "b"] * 1000,
+            "C": ["c"] * 2000,
+        }
+    )
+    pandas_df.set_index("idx").to_parquet("tmp.parquet")
+    # read the same parquet using modin.pandas
+    df_equals(pd.read_parquet("tmp.parquet"), pandas.read_parquet("tmp.parquet"))
+
+    pandas_df = pandas.DataFrame(
+        {
+            "idx": np.random.randint(0, 100_000, size=2000),
+            "A": np.random.randint(0, 100_000, size=2000),
+            "B": ["a", "b"] * 1000,
+            "C": ["c"] * 2000,
+        }
+    )
+    pandas_df.set_index(["idx", "A"]).to_parquet("tmp.parquet")
+    df_equals(pd.read_parquet("tmp.parquet"), pandas.read_parquet("tmp.parquet"))
+
+
 def test_from_parquet_hdfs():
     path = "modin/pandas/test/data/hdfs.parquet"
     pandas_df = pandas.read_parquet(path)


Code review comments:
diff_hunk: @@ -428,6 +428,32 @@ def test_from_parquet_partitioned_columns_with_columns(make_parquet_file):
     df_equals(modin_df, pandas_df)
 
 
+def test_from_parquet_pandas_index():
+    # Ensure modin can read parquet files written by pandas with a non-RangeIndex object
+    pandas_df = pandas.DataFrame(
+        {
+            "idx": np.random.randint(0, 100_000, size=2000),
+            "A": np.random.randint(0, 100_000, size=2000),
+            "B": ["a", "b"] * 1000,
+            "C": ["c"] * 2000,
+        }
+    )
+    pandas_df.set_index("idx").to_parquet("tmp.parquet")
+    # read the same parquet using modin.pandas
+    df_equals(pd.read_parquet("tmp.parquet"), pandas.read_parquet("tmp.parquet"))
+
+    pandas_df = pandas.DataFrame(
+        {
+            "idx": np.random.randint(0, 100_000, size=2000),
+            "A": np.random.randint(0, 100_000, size=2000),
+            "B": ["a", "b"] * 1000,
+            "C": ["c"] * 2000,
+        }
+    )
path: modin/pandas/test/test_io.py
text: @user1:
```suggestion

```

@user1:
It looks like we can use the previous one `pandas_df`