Repository: databricks/koalas

Issue:
Add index_col for spark IO reads.
We should provide `index_col` argument for spark IO reads to avoid creating default indexes in the case index columns are known.

- `read_csv`
- `read_delta`
- `read_table`
- `read_spark_io`
- `read_parquet`


Discussion:
`DataFrame.to_koalas()` as well?
@itholic, are you interested in submitting a PR for this?
@HyukjinKwon Sure, i'll work on it from now on. :)

Resolving patch:
diff --git a/databricks/koalas/namespace.py b/databricks/koalas/namespace.py
--- a/databricks/koalas/namespace.py
+++ b/databricks/koalas/namespace.py
@@ -26,6 +26,7 @@
 import numpy as np
 import pandas as pd
 
+from pyspark import sql as spark
 from pyspark.sql import functions as F
 from pyspark.sql.types import ByteType, ShortType, IntegerType, LongType, FloatType, \
     DoubleType, BooleanType, TimestampType, DecimalType, StringType, DateType, StructType
@@ -123,8 +124,8 @@ def range(start: int,
     return DataFrame(sdf)
 
 
-def read_csv(path, sep=',', header='infer', names=None, usecols=None,
-             mangle_dupe_cols=True, parse_dates=False, comment=None):
+def read_csv(path, sep=',', header='infer', names=None, index_col=None,
+             usecols=None, mangle_dupe_cols=True, parse_dates=False, comment=None):
     """Read CSV (comma-separated) file into DataFrame.
 
     Parameters
@@ -145,6 +146,8 @@ def read_csv(path, sep=',', header='infer', names=None, usecols=None,
         explicitly pass `header=None`. Duplicates in this list will cause an error to be issued.
         If a string is given, it should be a DDL-formatted string in Spark SQL, which is
         preferred to avoid schema inference for better performance.
+    index_col: str or list of str, optional, default: None
+        Index column of table in Spark.
     usecols : list-like or callable, optional
         Return a subset of the columns. If list-like, all elements must either be
         positional (i.e. integer indices into the document columns) or strings that
@@ -241,11 +244,14 @@ def read_csv(path, sep=',', header='infer', names=None, usecols=None,
                 sdf = default_session().createDataFrame([], schema=StructType())
     else:
         sdf = default_session().createDataFrame([], schema=StructType())
-    return DataFrame(sdf)
+
+    index_map = _get_index_map(sdf, index_col)
+
+    return DataFrame(_InternalFrame(sdf=sdf, index_map=index_map))
 
 
 def read_delta(path: str, version: Optional[str] = None, timestamp: Optional[str] = None,
-               **options) -> DataFrame:
+               index_col: Optional[Union[str, List[str]]] = None, **options) -> DataFrame:
     """
     Read a Delta Lake table on some file system and return a DataFrame.
 
@@ -262,6 +268,8 @@ def read_delta(path: str, version: Optional[str] = None, timestamp: Optional[str
         Specifies the table version (based on timestamp) to read from,
         using Delta's time travel feature. This must be a valid date or timestamp string in Spark,
         and sets Delta's 'timestampAsOf' option.
+    index_col : str or list of str, optional, default: None
+        Index column of table in Spark.
     options
         Additional options that can be passed onto Delta.
 
@@ -287,7 +295,7 @@ def read_delta(path: str, version: Optional[str] = None, timestamp: Optional[str
         options['versionAsOf'] = version
     if timestamp is not None:
         options['timestampAsOf'] = timestamp
-    return read_spark_io(path, format='delta', options=options)
+    return read_spark_io(path, format='delta', index_col=index_col, options=options)
 
 
 def read_table(name: str, index_col: Optional[Union[str, List[str]]] = None) -> DataFrame:
@@ -321,22 +329,15 @@ def read_table(name: str, index_col: Optional[Union[str, List[str]]] = None) ->
     0   0
     """
     sdf = default_session().read.table(name)
-    index_map = None  # type: Optional[List[IndexMap]]
-
-    if index_col is not None:
-        if isinstance(index_col, str):
-            index_col = [index_col]
-        sdf_columns = set(sdf.columns)
-        for col in index_col:
-            if col not in sdf_columns:
-                raise KeyError(col)
-        index_map = [(col, col) for col in index_col]
+    index_map = _get_index_map(sdf, index_col)
 
     return DataFrame(_InternalFrame(sdf=sdf, index_map=index_map))
 
 
 def read_spark_io(path: Optional[str] = None, format: Optional[str] = None,
-                  schema: Union[str, 'StructType'] = None, **options) -> DataFrame:
+                  schema: Union[str, 'StructType'] = None,
+                  index_col: Optional[Union[str, List[str]]] = None,
+                  **options) -> DataFrame:
     """Load a DataFrame from a Spark data source.
 
     Parameters
@@ -355,6 +356,8 @@ def read_spark_io(path: Optional[str] = None, format: Optional[str] = None,
         Input schema. If none, Spark tries to infer the schema automatically.
         The schema can either be a Spark StructType, or a DDL-formatted string like
         `col0 INT, col1 DOUBLE`.
+    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.
 
@@ -374,10 +377,12 @@ def read_spark_io(path: Optional[str] = None, format: Optional[str] = None,
     0   0
     """
     sdf = default_session().read.load(path=path, format=format, schema=schema, options=options)
-    return DataFrame(sdf)
+    index_map = _get_index_map(sdf, index_col)
+
+    return DataFrame(_InternalFrame(sdf=sdf, index_map=index_map))
 
 
-def read_parquet(path, columns=None) -> DataFrame:
+def read_parquet(path, columns=None, index_col=None) -> DataFrame:
     """Load a parquet object from the file path, returning a DataFrame.
 
     Parameters
@@ -386,6 +391,8 @@ def read_parquet(path, columns=None) -> DataFrame:
         File path
     columns : list, default=None
         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.
 
     Returns
     -------
@@ -418,7 +425,10 @@ def read_parquet(path, columns=None) -> DataFrame:
                 sdf = default_session().createDataFrame([], schema=StructType())
     else:
         sdf = default_session().createDataFrame([], schema=StructType())
-    return DataFrame(sdf)
+
+    index_map = _get_index_map(sdf, index_col)
+
+    return DataFrame(_InternalFrame(sdf=sdf, index_map=index_map))
 
 
 def read_clipboard(sep=r'\s+', **kwargs):
@@ -818,16 +828,7 @@ def read_sql_table(table_name, con, schema=None, index_col=None, columns=None, *
         reader.schema(schema)
     reader.options(**options)
     sdf = reader.format("jdbc").load()
-    if index_col is not None:
-        if isinstance(index_col, str):
-            index_col = [index_col]
-        sdf_columns = set(sdf.columns)
-        for col in index_col:
-            if col not in sdf_columns:
-                raise KeyError(col)
-        index_map = [(col, col) for col in index_col]
-    else:
-        index_map = None
+    index_map = _get_index_map(sdf, index_col)
     kdf = DataFrame(_InternalFrame(sdf=sdf, index_map=index_map))
     if columns is not None:
         if isinstance(columns, str):
@@ -878,16 +879,7 @@ def read_sql_query(sql, con, index_col=None, **options):
     reader.option('url', con)
     reader.options(**options)
     sdf = reader.format("jdbc").load()
-    if index_col is not None:
-        if isinstance(index_col, str):
-            index_col = [index_col]
-        sdf_columns = set(sdf.columns)
-        for col in index_col:
-            if col not in sdf_columns:
-                raise KeyError(col)
-        index_map = [(col, col) for col in index_col]
-    else:
-        index_map = None
+    index_map = _get_index_map(sdf, index_col)
     return DataFrame(_InternalFrame(sdf=sdf, index_map=index_map))
 
 
@@ -1639,6 +1631,22 @@ def _to_datetime2(arg_year, arg_month, arg_day,
         infer_datetime_format=infer_datetime_format)
 
 
+def _get_index_map(sdf: spark.DataFrame,
+                   index_col: Optional[Union[str, List[str]]] = None):
+    if index_col is not None:
+        if isinstance(index_col, str):
+            index_col = [index_col]
+        sdf_columns = set(sdf.columns)
+        for col in index_col:
+            if col not in sdf_columns:
+                raise KeyError(col)
+        index_map = [(col, col) for col in index_col]  # type: Optional[List[IndexMap]]
+    else:
+        index_map = None
+
+    return index_map
+
+
 _get_dummies_default_accept_types = (
     DecimalType, StringType, DateType
 )


Test patch:
diff --git a/databricks/koalas/tests/test_csv.py b/databricks/koalas/tests/test_csv.py
--- a/databricks/koalas/tests/test_csv.py
+++ b/databricks/koalas/tests/test_csv.py
@@ -141,6 +141,11 @@ def check(header='infer', names=None, usecols=None):
             self.assertRaisesRegex(ValueError, expected_error_message,
                                    lambda: ks.read_csv(fn, usecols=[1, 'amount']))
 
+            # check with index_col
+            expected = pd.read_csv(fn).set_index('name')
+            actual = ks.read_csv(fn, index_col='name')
+            self.assertPandasAlmostEqual(expected, actual.toPandas())
+
     def test_read_with_spark_schema(self):
         with self.csv_file(self.csv_text_2) as fn:
             actual = ks.read_csv(fn, names="A string, B string, C long, D long, E long")
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
@@ -70,6 +70,16 @@ def check(columns, expected):
             actual = ks.read_parquet(tmp)
             self.assertPandasEqual(expected, actual.toPandas())
 
+            # When index columns are known
+            pdf = self.test_pdf
+            expected = ks.DataFrame(pdf)
+
+            expected_idx = expected.set_index('bhello')[['f', 'i32', 'i64']]
+            actual_idx = ks.read_parquet(tmp, index_col='bhello')[['f', 'i32', 'i64']]
+            self.assert_eq(
+                actual_idx.sort_values(by='f').to_spark().toPandas(),
+                expected_idx.sort_values(by='f').to_spark().toPandas())
+
     def test_parquet_write(self):
         with self.temp_dir() as tmp:
             pdf = self.test_pdf
@@ -158,3 +168,14 @@ def test_spark_io(self):
             self.assert_eq(
                 actual.sort_values(by='f').to_spark().toPandas(),
                 expected.sort_values(by='f').to_spark().toPandas())
+
+            # When index columns are known
+            pdf = self.test_pdf
+            expected = ks.DataFrame(pdf)
+            col_order = ['f', 'i32', 'i64']
+
+            expected_idx = expected.set_index('bhello')[col_order]
+            actual_idx = ks.read_spark_io(tmp, format='json', index_col='bhello')[col_order]
+            self.assert_eq(
+                actual_idx.sort_values(by='f').to_spark().toPandas(),
+                expected_idx.sort_values(by='f').to_spark().toPandas())