Repository: databricks/koalas

Issue:
Implement to_parquet
It's kind of embarrassing that we don't have any "big data" input and output methods, other than read_parquet. 


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
@@ -1509,6 +1509,42 @@ def to_koalas(self):
         else:
             return DataFrame(self)
 
+    def to_parquet(self, path: str, mode: str = 'error',
+                   partition_cols: Union[str, List[str], None] = None, compression=None):
+        """
+        Write the DataFrame out as a Parquet file or directory.
+
+        Parameters
+        ----------
+        path : str, required
+            Path to write to.
+        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, default None
+            Names of partitioning columns
+        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`.
+
+        See Also
+        --------
+        read_parquet
+
+        Examples
+        --------
+        >>> df.write.parquet('my_data.parquet', partition='date')  # doctest: +SKIP
+
+        >>> df.write.parquet('my_data.parquet', partition=['date', 'country'])  # doctest: +SKIP
+        """
+        self._sdf.write.parquet(path=path, mode=mode, partitionBy=partition_cols,
+                                compression=compression)
+
     def to_spark(self):
         """
         Return the current DataFrame as a Spark DataFrame.
diff --git a/databricks/koalas/missing/frame.py b/databricks/koalas/missing/frame.py
--- a/databricks/koalas/missing/frame.py
+++ b/databricks/koalas/missing/frame.py
@@ -148,7 +148,6 @@ class _MissingPandasLikeDataFrame(object):
     to_gbq = unsupported_function('to_gbq')
     to_hdf = unsupported_function('to_hdf')
     to_msgpack = unsupported_function('to_msgpack')
-    to_parquet = unsupported_function('to_parquet')
     to_period = unsupported_function('to_period')
     to_pickle = unsupported_function('to_pickle')
     to_sparse = unsupported_function('to_sparse')
diff --git a/databricks/koalas/namespace.py b/databricks/koalas/namespace.py
--- a/databricks/koalas/namespace.py
+++ b/databricks/koalas/namespace.py
@@ -264,6 +264,10 @@ def read_parquet(path, columns=None):
     -------
     DataFrame
 
+    See Also
+    --------
+    DataFrame.to_parquet
+
     Examples
     --------
     >>> ks.read_parquet('data.parquet', columns=['name', 'gender'])  # doctest: +SKIP


Test patch:
diff --git a/databricks/koalas/tests/test_dataframe_conversion.py b/databricks/koalas/tests/test_dataframe_conversion.py
--- a/databricks/koalas/tests/test_dataframe_conversion.py
+++ b/databricks/koalas/tests/test_dataframe_conversion.py
@@ -25,6 +25,7 @@
 
 
 class DataFrameConversionTest(ReusedSQLTestCase, SQLTestUtils, TestUtils):
+    """Test cases for "small data" conversion and I/O."""
 
     @property
     def pdf(self):
diff --git a/databricks/koalas/tests/test_parquet.py b/databricks/koalas/tests/test_dataframe_spark_io.py
similarity index 51%
rename from databricks/koalas/tests/test_parquet.py
rename to databricks/koalas/tests/test_dataframe_spark_io.py
--- a/databricks/koalas/tests/test_parquet.py
+++ b/databricks/koalas/tests/test_dataframe_spark_io.py
@@ -19,27 +19,37 @@
 import numpy as np
 import pandas as pd
 
-from databricks import koalas
+from databricks import koalas as ks
 from databricks.koalas.testing.utils import ReusedSQLTestCase, TestUtils
 
 
-class ParquetTest(ReusedSQLTestCase, TestUtils):
+class DataFrameSparkIOTest(ReusedSQLTestCase, TestUtils):
+    """Test cases for big data I/O using Spark."""
 
-    def test_local(self):
+    @property
+    def test_column_order(self):
+        return ['i32', 'i64', 'f', 'bhello']
+
+    @property
+    def test_pdf(self):
+        pdf = pd.DataFrame({
+            'i32': np.arange(20, dtype=np.int32) % 3,
+            'i64': np.arange(20, dtype=np.int64) % 5,
+            'f': np.arange(20, dtype=np.float64),
+            'bhello': np.random.choice(['hello', 'yo', 'people'], size=20).astype("O")},
+            columns=self.test_column_order)
+        return pdf
+
+    def test_parquet_read(self):
         with self.temp_dir() as tmp:
-            data = pd.DataFrame({
-                'i32': np.arange(1000, dtype=np.int32),
-                'i64': np.arange(1000, dtype=np.int64),
-                'f': np.arange(1000, dtype=np.float64),
-                'bhello': np.random.choice(['hello', 'yo', 'people'], size=1000).astype("O")})
-            data = data[['i32', 'i64', 'f', 'bhello']]
+            data = self.test_pdf
             self.spark.createDataFrame(data, 'i32 int, i64 long, f double, bhello string') \
                 .coalesce(1).write.parquet(tmp, mode='overwrite')
 
             def check(columns, expected):
                 if LooseVersion("0.21.1") <= LooseVersion(pd.__version__):
                     expected = pd.read_parquet(tmp, columns=columns)
-                actual = koalas.read_parquet(tmp, columns=columns)
+                actual = ks.read_parquet(tmp, columns=columns)
                 self.assertPandasEqual(expected, actual.toPandas())
 
             check(None, data)
@@ -57,5 +67,24 @@ def check(columns, expected):
                 expected = pd.read_parquet(tmp)
             else:
                 expected = data
-            actual = koalas.read_parquet(tmp)
+            actual = ks.read_parquet(tmp)
             self.assertPandasEqual(expected, actual.toPandas())
+
+    def test_parquet_write(self):
+        with self.temp_dir() as tmp:
+            pdf = self.test_pdf
+            expected = ks.DataFrame(pdf)
+
+            # 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
+            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
+            actual = ks.read_parquet(tmp)[self.test_column_order]
+            self.assert_eq(actual.sort_values(by='f'), expected.sort_values(by='f'))