Repository: databricks/koalas

Issue:
Cannot read parquet file from Pandas
I computed some data using pandas which I wrote to a file. The df looks like this https://i.imgur.com/pxCLrUh.png 
I wrote it to parquet using `plot_df.to_parquet("test.parquet", flavor='spark')`

Now when I read it in using spark, there is no issue, yet note there being one additional column which apparently pandas adds:
```
df = spark.read.parquet("test.parquet")
df.printSchema()
root
 |-- time: timestamp (nullable = true)
 |-- node: string (nullable = true)
 |-- power_usage: double (nullable = true)
 |-- temperature: double (nullable = true)
 |-- processes_running: double (nullable = true)
 |-- memory_used: double (nullable = true)
 |-- load1: double (nullable = true)
 |-- __index_level_0__: long (nullable = true)
```

Now if I run `kdf = df.to_koalas()` I get
```
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-29-4c46820eb041> in <module>
      1 df = spark.read.parquet("test.parquet")
      2 df.printSchema()
----> 3 kdf = df.to_koalas()

/var/scratch/lvs215/miniconda3/lib/python3.8/site-packages/databricks/koalas/frame.py in to_koalas(self, index_col)
   4449 
   4450             index_map = _get_index_map(self, index_col)
-> 4451             internal = InternalFrame(spark_frame=self, index_map=index_map)
   4452             return DataFrame(internal)
   4453 

/var/scratch/lvs215/miniconda3/lib/python3.8/site-packages/databricks/koalas/internal.py in __init__(self, spark_frame, index_map, column_labels, data_spark_columns, column_label_names)
    436 
    437         if index_map is None:
--> 438             assert not any(SPARK_INDEX_NAME_PATTERN.match(name) for name in spark_frame.columns), (
    439                 "Index columns should not appear in columns of the Spark DataFrame. Avoid "
    440                 "index column names [%s]." % SPARK_INDEX_NAME_PATTERN

AssertionError: Index columns should not appear in columns of the Spark DataFrame. Avoid index column names [re.compile('__index_level_[0-9]+__')].
```

The issue looks similar to https://github.com/xitongsys/parquet-go/issues/233

How to deal with this? Should I follow the suggestion in https://github.com/apache/arrow/issues/2244 and don't preserve the index?


Discussion:
You can specify the index column:

```py
kdf = df.to_koalas(index_col='__index_level_0__')
```

or

```py
kdf = ks.read_parquet("test.parquet", index_col='__index_level_0__')
```

See also: https://koalas.readthedocs.io/en/latest/user_guide/pandas_pyspark.html#pyspark
Forgot to comment.

Your suggestion indeed works. I think this is something that koalas (by design) cannot pickup upon automatically? It's not a big deal, but it in code it might appear a bit as a "magic" index name. On the other hand, the one that output the pandas dataframe could've exported it without that index as well, so that would be a counter argument.

Feel free to close this one if you'd like.
Actually parquet files output by pandas have metadata on which columns should be indices.
Do you want to work on picking up the metadata automatically in `read_parquet`?
For the case when you do `spark.read.parquet("test.parquet").to_koalas()`, we can't pick up automatically since Spark DataFrame doesn't have such metadata.

Resolving patch:
diff --git a/databricks/koalas/namespace.py b/databricks/koalas/namespace.py
--- a/databricks/koalas/namespace.py
+++ b/databricks/koalas/namespace.py
@@ -23,13 +23,17 @@
 from distutils.version import LooseVersion
 from functools import reduce
 from io import BytesIO
+import json
 
 import numpy as np
 import pandas as pd
 from pandas.api.types import is_datetime64_dtype, is_datetime64tz_dtype, is_list_like
+import pyarrow as pa
+import pyarrow.parquet as pq
 import pyspark
 from pyspark import sql as spark
 from pyspark.sql import functions as F
+from pyspark.sql.functions import pandas_udf, PandasUDFType
 from pyspark.sql.types import (
     ByteType,
     ShortType,
@@ -625,7 +629,7 @@ def read_spark_io(
     return DataFrame(InternalFrame(spark_frame=sdf, index_map=index_map))
 
 
-def read_parquet(path, columns=None, index_col=None, **options) -> DataFrame:
+def read_parquet(path, columns=None, index_col=None, pandas_metadata=False, **options) -> DataFrame:
     """Load a parquet object from the file path, returning a DataFrame.
 
     Parameters
@@ -636,6 +640,8 @@ def read_parquet(path, columns=None, index_col=None, **options) -> DataFrame:
         If not None, only these columns will be read from the file.
     index_col : str or list of str, optional, default: None
         Index column of table in Spark.
+    pandas_metadata : bool, default: False
+        If True, try to respect the metadata if the Parquet file is written from pandas.
     options : dict
         All other options passed directly into Spark's data source.
 
@@ -672,6 +678,44 @@ def read_parquet(path, columns=None, index_col=None, **options) -> DataFrame:
     if columns is not None:
         columns = list(columns)
 
+    index_names = None
+
+    if index_col is None and pandas_metadata:
+        if LooseVersion(pyspark.__version__) < LooseVersion("3.0.0"):
+            raise ValueError("pandas_metadata is not supported with Spark < 3.0.")
+
+        # Try to read pandas metadata
+
+        @pandas_udf("index_col array<string>, index_names array<string>", PandasUDFType.SCALAR)
+        def read_index_metadata(pser):
+            binary = pser.iloc[0]
+            metadata = pq.ParquetFile(pa.BufferReader(binary)).metadata.metadata
+            if b"pandas" in metadata:
+                pandas_metadata = json.loads(metadata[b"pandas"].decode("utf8"))
+                if all(isinstance(col, str) for col in pandas_metadata["index_columns"]):
+                    index_col = []
+                    index_names = []
+                    for col in pandas_metadata["index_columns"]:
+                        index_col.append(col)
+                        for column in pandas_metadata["columns"]:
+                            if column["field_name"] == col:
+                                index_names.append(column["name"])
+                                break
+                        else:
+                            index_names.append(None)
+                    return pd.DataFrame({"index_col": [index_col], "index_names": [index_names]})
+            return pd.DataFrame({"index_col": [None], "index_names": [None]})
+
+        index_col, index_names = (
+            default_session()
+            .read.format("binaryFile")
+            .load(path)
+            .limit(1)
+            .select(read_index_metadata("content").alias("index_metadata"))
+            .select("index_metadata.*")
+            .head()
+        )
+
     kdf = read_spark_io(path=path, format="parquet", options=options, index_col=index_col)
 
     if columns is not None:
@@ -681,7 +725,10 @@ def read_parquet(path, columns=None, index_col=None, **options) -> DataFrame:
         else:
             sdf = default_session().createDataFrame([], schema=StructType())
             index_map = _get_index_map(sdf, index_col)
-            return DataFrame(InternalFrame(spark_frame=sdf, index_map=index_map))
+            kdf = DataFrame(InternalFrame(spark_frame=sdf, index_map=index_map))
+
+    if index_names is not None:
+        kdf.index.names = index_names
 
     return kdf
 


Test patch:
diff --git a/databricks/koalas/tests/test_dataframe_spark_io.py b/databricks/koalas/tests/test_dataframe_spark_io.py
--- a/databricks/koalas/tests/test_dataframe_spark_io.py
+++ b/databricks/koalas/tests/test_dataframe_spark_io.py
@@ -15,6 +15,7 @@
 #
 
 from distutils.version import LooseVersion
+import unittest
 
 import numpy as np
 import pandas as pd
@@ -91,6 +92,33 @@ def check(columns, expected):
                 expected_idx.sort_values(by="f").to_spark().toPandas(),
             )
 
+    @unittest.skipIf(
+        LooseVersion(pyspark.__version__) < LooseVersion("3.0.0"),
+        "The test only works with Spark>=3.0",
+    )
+    def test_parquet_read_with_pandas_metadata(self):
+        with self.temp_dir() as tmp:
+            expected1 = self.test_pdf
+
+            path1 = "{}/file1.parquet".format(tmp)
+            expected1.to_parquet(path1)
+
+            self.assert_eq(ks.read_parquet(path1, pandas_metadata=True), expected1)
+
+            expected2 = expected1.reset_index()
+
+            path2 = "{}/file2.parquet".format(tmp)
+            expected2.to_parquet(path2)
+
+            self.assert_eq(ks.read_parquet(path2, pandas_metadata=True), expected2)
+
+            expected3 = expected2.set_index("index", append=True)
+
+            path3 = "{}/file3.parquet".format(tmp)
+            expected3.to_parquet(path3)
+
+            self.assert_eq(ks.read_parquet(path3, pandas_metadata=True), expected3)
+
     def test_parquet_write(self):
         with self.temp_dir() as tmp:
             pdf = self.test_pdf