Repository: databricks/koalas

Issue:
Implement generic functionality to read / write Spark data source tables
Basically a wrapper around SparkSession.read and spark.DataFrame.write.

We need to decide on a name here ... Some names I can think of:

- read_spark_io / to_spark_io
- read_data / to_data
- read_spark_ds / to_spark_ds
- read_data_source / to_data_source
...

Any other thoughts?



Discussion:

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -1510,7 +1510,8 @@ def to_koalas(self):
             return DataFrame(self)
 
     def to_parquet(self, path: str, mode: str = 'error',
-                   partition_cols: Union[str, List[str], None] = None, compression=None):
+                   partition_cols: Union[str, List[str], None] = None,
+                   compression: Optional[str] = None):
         """
         Write the DataFrame out as a Parquet file or directory.
 
@@ -1545,6 +1546,47 @@ def to_parquet(self, path: str, mode: str = 'error',
         self._sdf.write.parquet(path=path, mode=mode, partitionBy=partition_cols,
                                 compression=compression)
 
+    def to_spark_io(self, path: Optional[str] = None, format: Optional[str] = None,
+                    mode: str = 'error', partition_cols: Union[str, List[str], None] = None,
+                    **options):
+        """Write the DataFrame out to a Spark data source.
+
+        Parameters
+        ----------
+        path : string, optional
+            Path to the data source.
+        format : string, optional
+            Specifies the output data source format. Some common ones are:
+
+            - 'delta'
+            - 'parquet'
+            - 'orc'
+            - 'json'
+            - 'csv'
+        mode : str {'append', 'overwrite', 'ignore', 'error', 'errorifexists'}, default 'error'.
+            Specifies the behavior of the save operation when data already.
+
+            - 'append': Append the new data to existing data.
+            - 'overwrite': Overwrite existing data.
+            - 'ignore': Silently ignore this operation if data already exists.
+            - 'error' or 'errorifexists': Throw an exception if data already exists.
+        partition_cols : str or list of str, optional
+            Names of partitioning columns
+        options : dict
+            All other options passed directly into Spark's data source.
+
+        See Also
+        --------
+        read_spark_io
+        to_parquet
+
+        Examples
+        --------
+        >>> df.to_spark_io(path='data.json', format='json')  # doctest: +SKIP
+        """
+        self._sdf.write.save(path=path, format=format, mode=mode, partitionBy=partition_cols,
+                             options=options)
+
     def to_spark(self):
         """
         Return the current DataFrame as a Spark DataFrame.
diff --git a/databricks/koalas/namespace.py b/databricks/koalas/namespace.py
--- a/databricks/koalas/namespace.py
+++ b/databricks/koalas/namespace.py
@@ -250,6 +250,42 @@ def read_csv(path, header='infer', names=None, usecols=None,
     return DataFrame(sdf)
 
 
+def read_spark_io(path: Optional[str] = None, format: Optional[str] = None,
+                  schema: Union[str, 'StructType'] = None, **options) -> DataFrame:
+    """Load a DataFrame from a Spark data source.
+
+    Parameters
+    ----------
+    path : string, optional
+        Path to the data source.
+    format : string, optional
+        Specifies the output data source format. Some common ones are:
+
+        - 'delta'
+        - 'parquet'
+        - 'orc'
+        - 'json'
+        - 'csv'
+    schema : string or StructType, optional
+        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`.
+    options : dict
+        All other options passed directly into Spark's data source.
+
+    See Also
+    --------
+    DataFrame.to_spark_io
+    DataFrame.read_parquet
+
+    Examples
+    --------
+    >>> ks.read_spark_io('data.parquet', format='parquet', schema='name string')  # doctest: +SKIP
+    """
+    sdf = default_session().read.load(path=path, format=format, schema=schema, options=options)
+    return DataFrame(sdf)
+
+
 def read_parquet(path, columns=None):
     """Load a parquet object from the file path, returning a DataFrame.
 


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
@@ -77,14 +77,34 @@ def test_parquet_write(self):
 
             # Write out partitioned by one column
             expected.to_parquet(tmp, mode='overwrite', partition_cols='i32')
-            # reset column order, as once the data is written out, Spark rearranges partition
-            # columns to appear first
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
             actual = ks.read_parquet(tmp)[self.test_column_order]
             self.assert_eq(actual.sort_values(by='f'), expected.sort_values(by='f'))
 
             # Write out partitioned by two columns
             expected.to_parquet(tmp, mode='overwrite', partition_cols=['i32', 'bhello'])
-            # reset column order, as once the data is written out, Spark rearranges partition
-            # columns to appear first
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
             actual = ks.read_parquet(tmp)[self.test_column_order]
             self.assert_eq(actual.sort_values(by='f'), expected.sort_values(by='f'))
+
+    def test_spark_io(self):
+        with self.temp_dir() as tmp:
+            pdf = self.test_pdf
+            expected = ks.DataFrame(pdf)
+
+            # Write out partitioned by one column
+            expected.to_spark_io(tmp, format='json', mode='overwrite', partition_cols='i32')
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
+            actual = ks.read_spark_io(tmp, format='json')[self.test_column_order]
+            self.assert_eq(actual.sort_values(by='f'), expected.sort_values(by='f'))
+
+            # Write out partitioned by two columns
+            expected.to_spark_io(tmp, format='json', mode='overwrite',
+                                 partition_cols=['i32', 'bhello'])
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
+            actual = ks.read_spark_io(path=tmp, format='json')[self.test_column_order]
+            self.assert_eq(actual.sort_values(by='f'), expected.sort_values(by='f'))