Repository: databricks/koalas

Issue:
to_delta and to_parquet are dropping the index
Hi, I don't know if I made a mistake but `ks.DataFrame.to_delta` and `ks.read_delta` are dropping my index columns:
![grafik](https://user-images.githubusercontent.com/1200058/77863412-7e33cd00-7222-11ea-8980-304631d4fa18.png)

Is there a way to keep the index with the dataset?
Also, does this index actually give a performance boost compared to default Spark joins when I merge two large (terabyte range) datasets?


Discussion:
Can you try to specify the index columns via `ks.read_detla(index_col=["configName", "start", "emd"])`?
Yes I tried, the columns are just missing.

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -41,6 +41,7 @@
 from pyspark import sql as spark
 from pyspark.sql import functions as F, Column
 from pyspark.sql.functions import pandas_udf
+from pyspark.sql.readwriter import OptionUtils
 from pyspark.sql.types import (
     BooleanType,
     ByteType,
@@ -3937,6 +3938,7 @@ def to_table(
         format: Optional[str] = None,
         mode: str = "overwrite",
         partition_cols: Union[str, List[str], None] = None,
+        index_col: Optional[Union[str, List[str]]] = None,
         **options
     ):
         """
@@ -3966,6 +3968,9 @@ def to_table(
 
         partition_cols : str or list of str, optional, default None
             Names of partitioning columns
+        index_col: str or list of str, optional, default: None
+            Column names to be used in Spark to represent Koalas' index. The index name
+            in Koalas is ignored. By default, the index is always lost.
         options
             Additional options passed directly to Spark.
 
@@ -3989,7 +3994,7 @@ def to_table(
 
         >>> df.to_table('%s.my_table' % db, partition_cols='date')
         """
-        self.to_spark().write.saveAsTable(
+        self.to_spark(index_col=index_col).write.saveAsTable(
             name=name, format=format, mode=mode, partitionBy=partition_cols, **options
         )
 
@@ -3998,6 +4003,7 @@ def to_delta(
         path: str,
         mode: str = "overwrite",
         partition_cols: Union[str, List[str], None] = None,
+        index_col: Optional[Union[str, List[str]]] = None,
         **options
     ):
         """
@@ -4018,6 +4024,9 @@ def to_delta(
 
         partition_cols : str or list of str, optional, default None
             Names of partitioning columns
+        index_col: str or list of str, optional, default: None
+            Column names to be used in Spark to represent Koalas' index. The index name
+            in Koalas is ignored. By default, the index is always lost.
         options : dict
             All other options passed directly into Delta Lake.
 
@@ -4055,7 +4064,12 @@ def to_delta(
         ...             mode='overwrite', replaceWhere='date >= "2012-01-01"')
         """
         self.to_spark_io(
-            path=path, mode=mode, format="delta", partition_cols=partition_cols, **options
+            path=path,
+            mode=mode,
+            format="delta",
+            partition_cols=partition_cols,
+            index_col=index_col,
+            **options
         )
 
     def to_parquet(
@@ -4064,6 +4078,8 @@ def to_parquet(
         mode: str = "overwrite",
         partition_cols: Union[str, List[str], None] = None,
         compression: Optional[str] = None,
+        index_col: Optional[Union[str, List[str]]] = None,
+        **options
     ):
         """
         Write the DataFrame out as a Parquet file or directory.
@@ -4086,6 +4102,11 @@ def to_parquet(
         compression : str {'none', 'uncompressed', 'snappy', 'gzip', 'lzo', 'brotli', 'lz4', 'zstd'}
             Compression codec to use when saving to file. If None is set, it uses the
             value specified in `spark.sql.parquet.compression.codec`.
+        index_col: str or list of str, optional, default: None
+            Column names to be used in Spark to represent Koalas' index. The index name
+            in Koalas is ignored. By default, the index is always lost.
+        options : dict
+            All other options passed directly into Spark's data source.
 
         See Also
         --------
@@ -4113,9 +4134,11 @@ def to_parquet(
         ...     mode = 'overwrite',
         ...     partition_cols=['date', 'country'])
         """
-        self.to_spark().write.parquet(
-            path=path, mode=mode, partitionBy=partition_cols, compression=compression
+        builder = self.to_spark(index_col=index_col).write.mode(mode)
+        OptionUtils._set_opts(
+            builder, mode=mode, partitionBy=partition_cols, compression=compression
         )
+        builder.options(**options).format("parquet").save(path)
 
     def to_spark_io(
         self,
@@ -4123,6 +4146,7 @@ def to_spark_io(
         format: Optional[str] = None,
         mode: str = "overwrite",
         partition_cols: Union[str, List[str], None] = None,
+        index_col: Optional[Union[str, List[str]]] = None,
         **options
     ):
         """Write the DataFrame out to a Spark data source.
@@ -4148,6 +4172,9 @@ def to_spark_io(
             - 'error' or 'errorifexists': Throw an exception if data already exists.
         partition_cols : str or list of str, optional
             Names of partitioning columns
+        index_col: str or list of str, optional, default: None
+            Column names to be used in Spark to represent Koalas' index. The index name
+            in Koalas is ignored. By default, the index is always lost.
         options : dict
             All other options passed directly into Spark's data source.
 
@@ -4172,7 +4199,7 @@ def to_spark_io(
 
         >>> df.to_spark_io(path='%s/to_spark_io/foo.json' % path, format='json')
         """
-        self.to_spark().write.save(
+        self.to_spark(index_col=index_col).write.save(
             path=path, format=format, mode=mode, partitionBy=partition_cols, **options
         )
 
diff --git a/databricks/koalas/generic.py b/databricks/koalas/generic.py
--- a/databricks/koalas/generic.py
+++ b/databricks/koalas/generic.py
@@ -22,6 +22,7 @@
 from collections.abc import Iterable
 from distutils.version import LooseVersion
 from functools import reduce
+from typing import Optional, Union, List
 
 import numpy as np
 import pandas as pd
@@ -39,6 +40,7 @@
     scol_for,
     validate_axis,
     align_diff_frames,
+    name_like_string,
 )
 from databricks.koalas.window import Rolling, Expanding
 
@@ -554,6 +556,7 @@ def to_csv(
         date_format=None,
         escapechar=None,
         num_files=None,
+        index_col: Optional[Union[str, List[str]]] = None,
         **options
     ):
         r"""
@@ -589,6 +592,9 @@ def to_csv(
             when appropriate.
         num_files : the number of files to be written in `path` directory when
             this is a path.
+        index_col: str or list of str, optional, default: None
+            Column names to be used in Spark to represent Koalas' index. The index name
+            in Koalas is ignored. By default, the index is always lost.
         options: keyword arguments for additional options specific to PySpark.
             This kwargs are specific to PySpark's CSV options to pass. Check
             the options in PySpark's API documentation for spark.write.csv(...).
@@ -646,6 +652,22 @@ def to_csv(
         ... 2012-01-31 12:00:00
         ... 2012-02-29 12:00:00
         ... 2012-03-31 12:00:00
+
+        You can preserve the index in the roundtrip as below.
+
+        >>> df.set_index("country", append=True, inplace=True)
+        >>> df.date.to_csv(
+        ...     path=r'%s/to_csv/bar.csv' % path,
+        ...     num_files=1,
+        ...     index_col=["index1", "index2"])
+        >>> ks.read_csv(
+        ...     path=r'%s/to_csv/bar.csv' % path, index_col=["index1", "index2"]
+        ... ).sort_values(by="date")  # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
+                                     date
+        index1 index2
+        ...    ...    2012-01-31 12:00:00
+        ...    ...    2012-02-29 12:00:00
+        ...    ...    2012-03-31 12:00:00
         """
         if path is None:
             # If path is none, just collect and use pandas's to_csv.
@@ -654,7 +676,7 @@ def to_csv(
                 self, ks.Series
             ):
                 # 0.23 seems not having 'columns' parameter in Series' to_csv.
-                return kdf_or_ser.to_pandas().to_csv(
+                return kdf_or_ser.to_pandas().to_csv(  # type: ignore
                     None,
                     sep=sep,
                     na_rep=na_rep,
@@ -663,7 +685,7 @@ def to_csv(
                     index=False,
                 )
             else:
-                return kdf_or_ser.to_pandas().to_csv(
+                return kdf_or_ser.to_pandas().to_csv(  # type: ignore
                     None,
                     sep=sep,
                     na_rep=na_rep,
@@ -686,21 +708,34 @@ def to_csv(
         elif isinstance(columns, tuple):
             column_labels = [columns]
         else:
-            column_labels = [label if isinstance(label, tuple) else (label,) for label in columns]
+            column_labels = [
+                lb if isinstance(lb, tuple) else (lb,) for lb in columns  # type: ignore
+            ]
+
+        if isinstance(index_col, str):
+            index_cols = [index_col]
+        elif index_col is None:
+            index_cols = []
+        else:
+            index_cols = index_col
 
         if header is True and kdf._internal.column_labels_level > 1:
             raise ValueError("to_csv only support one-level index column now")
         elif isinstance(header, list):
-            sdf = kdf._sdf.select(
-                [
+            sdf = kdf.to_spark(index_col)  # type: ignore
+            sdf = sdf.select(
+                [scol_for(sdf, name_like_string(label)) for label in index_cols]
+                + [
                     self._internal.spark_column_for(label).alias(new_name)
                     for (label, new_name) in zip(column_labels, header)
                 ]
             )
             header = True
         else:
-            sdf = kdf._sdf.select(
-                [kdf._internal.spark_column_for(label) for label in column_labels]
+            sdf = kdf.to_spark(index_col)  # type: ignore
+            sdf = sdf.select(
+                [scol_for(sdf, name_like_string(label)) for label in index_cols]
+                + [kdf._internal.spark_column_for(label) for label in column_labels]
             )
 
         if num_files is not None:
@@ -719,7 +754,14 @@ def to_csv(
         )
         builder.options(**options).format("csv").save(path)
 
-    def to_json(self, path=None, compression="uncompressed", num_files=None, **options):
+    def to_json(
+        self,
+        path=None,
+        compression="uncompressed",
+        num_files=None,
+        index_col: Optional[Union[str, List[str]]] = None,
+        **options
+    ):
         """
         Convert the object to a JSON string.
 
@@ -748,6 +790,9 @@ def to_json(self, path=None, compression="uncompressed", num_files=None, **optio
             compression is inferred from the filename.
         num_files : the number of files to be written in `path` directory when
             this is a path.
+        index_col: str or list of str, optional, default: None
+            Column names to be used in Spark to represent Koalas' index. The index name
+            in Koalas is ignored. By default, the index is always lost.
         options: keyword arguments for additional options specific to PySpark.
             It is specific to PySpark's JSON options to pass. Check
             the options in PySpark's API documentation for `spark.write.json(...)`.
@@ -772,18 +817,19 @@ def to_json(self, path=None, compression="uncompressed", num_files=None, **optio
         0     a     b
         1     c     d
 
-        >>> df['col 1'].to_json(path=r'%s/to_json/foo.json' % path, num_files=1)
+        >>> df['col 1'].to_json(path=r'%s/to_json/foo.json' % path, num_files=1, index_col="index")
         >>> ks.read_json(
-        ...     path=r'%s/to_json/foo.json' % path
-        ... ).sort_values(by="col 1")
-          col 1
-        0     a
-        1     c
+        ...     path=r'%s/to_json/foo.json' % path, index_col="index"
+        ... ).sort_values(by="col 1")  # doctest: +NORMALIZE_WHITESPACE
+              col 1
+        index
+        0         a
+        1         c
         """
         if path is None:
             # If path is none, just collect and use pandas's to_json.
             kdf_or_ser = self
-            pdf = kdf_or_ser.to_pandas()
+            pdf = kdf_or_ser.to_pandas()  # type: ignore
             if isinstance(self, ks.Series):
                 pdf = pdf.to_frame()
             # To make the format consistent and readable by `read_json`, convert it to pandas' and
@@ -793,7 +839,7 @@ def to_json(self, path=None, compression="uncompressed", num_files=None, **optio
         kdf = self
         if isinstance(self, ks.Series):
             kdf = self.to_frame()
-        sdf = kdf.to_spark()
+        sdf = kdf.to_spark(index_col=index_col)  # type: ignore
 
         if num_files is not None:
             sdf = sdf.repartition(num_files)
diff --git a/databricks/koalas/namespace.py b/databricks/koalas/namespace.py
--- a/databricks/koalas/namespace.py
+++ b/databricks/koalas/namespace.py
@@ -384,6 +384,17 @@ def read_json(path: str, index_col: Optional[Union[str, List[str]]] = None, **op
       col 1 col 2
     0     a     b
     1     c     d
+
+    You can preserve the index in the roundtrip as below.
+
+    >>> df.to_json(path=r'%s/read_json/bar.json' % path, num_files=1, index_col="index")
+    >>> ks.read_json(
+    ...     path=r'%s/read_json/bar.json' % path, index_col="index"
+    ... ).sort_values(by="col 1")  # doctest: +NORMALIZE_WHITESPACE
+          col 1 col 2
+    index
+    0         a     b
+    1         c     d
     """
     return read_spark_io(path, format="json", index_col=index_col, **options)
 
@@ -446,6 +457,20 @@ def read_delta(
     >>> ks.read_delta('%s/read_delta/foo' % path, version=0)
        id
     0   0
+
+    You can preserve the index in the roundtrip as below.
+
+    >>> ks.range(10, 15, num_partitions=1).to_delta(
+    ...     '%s/read_delta/bar' % path, index_col="index")
+    >>> ks.read_delta('%s/read_delta/bar' % path, index_col="index")
+    ... # doctest: +NORMALIZE_WHITESPACE
+           id
+    index
+    0      10
+    1      11
+    2      12
+    3      13
+    4      14
     """
     if version is not None:
         options["versionAsOf"] = version
@@ -483,6 +508,12 @@ def read_table(name: str, index_col: Optional[Union[str, List[str]]] = None) ->
     >>> ks.read_table('%s.my_table' % db)
        id
     0   0
+
+    >>> ks.range(1).to_table('%s.my_table' % db, index_col="index")
+    >>> ks.read_table('%s.my_table' % db, index_col="index")  # doctest: +NORMALIZE_WHITESPACE
+           id
+    index
+    0       0
     """
     sdf = default_session().read.table(name)
     index_map = _get_index_map(sdf, index_col)
@@ -545,6 +576,21 @@ def read_spark_io(
     2  12
     3  13
     4  14
+
+    You can preserve the index in the roundtrip as below.
+
+    >>> ks.range(10, 15, num_partitions=1).to_spark_io('%s/read_spark_io/data.orc' % path,
+    ...                                                format='orc', index_col="index")
+    >>> ks.read_spark_io(
+    ...     path=r'%s/read_spark_io/data.orc' % path, format="orc", index_col="index")
+    ... # doctest: +NORMALIZE_WHITESPACE
+           id
+    index
+    0      10
+    1      11
+    2      12
+    3      13
+    4      14
     """
     sdf = default_session().read.load(path=path, format=format, schema=schema, **options)
     index_map = _get_index_map(sdf, index_col)
@@ -552,7 +598,7 @@ def read_spark_io(
     return DataFrame(_InternalFrame(spark_frame=sdf, index_map=index_map))
 
 
-def read_parquet(path, columns=None, index_col=None) -> DataFrame:
+def read_parquet(path, columns=None, index_col=None, **options) -> DataFrame:
     """Load a parquet object from the file path, returning a DataFrame.
 
     Parameters
@@ -563,6 +609,8 @@ def read_parquet(path, columns=None, index_col=None) -> 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.
+    options : dict
+        All other options passed directly into Spark's data source.
 
     Returns
     -------
@@ -581,24 +629,31 @@ def read_parquet(path, columns=None, index_col=None) -> DataFrame:
     >>> ks.read_parquet('%s/read_spark_io/data.parquet' % path, columns=['id'])
        id
     0   0
+
+    You can preserve the index in the roundtrip as below.
+
+    >>> ks.range(1).to_parquet('%s/read_spark_io/data.parquet' % path, index_col="index")
+    >>> ks.read_parquet('%s/read_spark_io/data.parquet' % path, columns=['id'], index_col="index")
+    ... # doctest: +NORMALIZE_WHITESPACE
+           id
+    index
+    0       0
     """
     if columns is not None:
         columns = list(columns)
-    if columns is None or len(columns) > 0:
-        sdf = default_session().read.parquet(path)
-        if columns is not None:
-            fields = [field.name for field in sdf.schema]
-            cols = [col for col in columns if col in fields]
-            if len(cols) > 0:
-                sdf = sdf.select(cols)
-            else:
-                sdf = default_session().createDataFrame([], schema=StructType())
-    else:
-        sdf = default_session().createDataFrame([], schema=StructType())
 
-    index_map = _get_index_map(sdf, index_col)
+    kdf = read_spark_io(path=path, format="parquet", options=options, index_col=index_col)
 
-    return DataFrame(_InternalFrame(spark_frame=sdf, index_map=index_map))
+    if columns is not None:
+        new_columns = [c for c in columns if c in kdf.columns]
+        if len(new_columns) > 0:
+            kdf = kdf[new_columns]
+        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))
+
+    return kdf
 
 
 def read_clipboard(sep=r"\s+", **kwargs):


Test patch: