Repository: databricks/koalas

Issue:
df.to_parquet doesn't create partitions
My environment: Databrick platform: runtime 7.0 ML
Koalas: 1.0.1

I'm trying to write parquets from a koalas dataframe to S3 with partitions. The partitions are not created (I tried with a single or multiple partition cols).
If I'm using the pyspark API, the partitions are created.

code:
no partition created
df.to_parquet(path='s3://{bucket}/{Path_to_data}, mode='overwrite', compression='gzip', partition_cols=['year','month','day'])

Partition created
df.to_spark().write.mode('overwrite').partitionBy('year', 'month', 'day').parquet('s3://{bucket}/{Path_to_data}')

I would like to have the possibilities to use the partitions in koalas in the same way I'm doing it in spark.



Discussion:
Seems like it's a bug.
Could you use `DataFrame.spark.to_spark_io` as a workaround?

```py
df.spark.to_spark_io(path='s3://{bucket}/{Path_to_data}, format='parquet', mode='overwrite', compression='gzip', partition_cols=['year','month','day'])
```

I'll work on fixing this.

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -43,7 +43,6 @@
 from pyspark import sql as spark
 from pyspark.sql import functions as F, Column
 from pyspark.sql.functions import pandas_udf, PandasUDFType
-from pyspark.sql.readwriter import OptionUtils
 from pyspark.sql.types import (
     BooleanType,
     DoubleType,
@@ -4026,7 +4025,7 @@ def to_table(
         name: str,
         format: Optional[str] = None,
         mode: str = "overwrite",
-        partition_cols: Union[str, List[str], None] = None,
+        partition_cols: Optional[Union[str, List[str]]] = None,
         index_col: Optional[Union[str, List[str]]] = None,
         **options
     ):
@@ -4038,7 +4037,7 @@ def to_delta(
         self,
         path: str,
         mode: str = "overwrite",
-        partition_cols: Union[str, List[str], None] = None,
+        partition_cols: Optional[Union[str, List[str]]] = None,
         index_col: Optional[Union[str, List[str]]] = None,
         **options
     ):
@@ -4115,7 +4114,7 @@ def to_parquet(
         self,
         path: str,
         mode: str = "overwrite",
-        partition_cols: Union[str, List[str], None] = None,
+        partition_cols: Optional[Union[str, List[str]]] = None,
         compression: Optional[str] = None,
         index_col: Optional[Union[str, List[str]]] = None,
         **options
@@ -4177,9 +4176,9 @@ def to_parquet(
             options = options.get("options")  # type: ignore
 
         builder = self.to_spark(index_col=index_col).write.mode(mode)
-        OptionUtils._set_opts(
-            builder, mode=mode, partitionBy=partition_cols, compression=compression
-        )
+        if partition_cols is not None:
+            builder.partitionBy(partition_cols)
+        builder._set_opts(compression=compression)
         builder.options(**options).format("parquet").save(path)
 
     def to_spark_io(
@@ -4187,7 +4186,7 @@ def to_spark_io(
         path: Optional[str] = None,
         format: Optional[str] = None,
         mode: str = "overwrite",
-        partition_cols: Union[str, List[str], None] = None,
+        partition_cols: Optional[Union[str, List[str]]] = None,
         index_col: Optional[Union[str, List[str]]] = None,
         **options
     ):
diff --git a/databricks/koalas/generic.py b/databricks/koalas/generic.py
--- a/databricks/koalas/generic.py
+++ b/databricks/koalas/generic.py
@@ -30,7 +30,6 @@
 
 from pyspark import sql as spark
 from pyspark.sql import functions as F
-from pyspark.sql.readwriter import OptionUtils
 from pyspark.sql.types import DataType, DoubleType, FloatType
 
 from databricks import koalas as ks  # For running doctests and reference resolution in PyCharm.
@@ -576,6 +575,8 @@ def to_csv(
         date_format=None,
         escapechar=None,
         num_files=None,
+        mode: str = "overwrite",
+        partition_cols: Optional[Union[str, List[str]]] = None,
         index_col: Optional[Union[str, List[str]]] = None,
         **options
     ):
@@ -612,6 +613,17 @@ def to_csv(
             when appropriate.
         num_files : the number of files to be written in `path` directory when
             this is a path.
+        mode : str {'append', 'overwrite', 'ignore', 'error', 'errorifexists'},
+            default 'overwrite'. Specifies the behavior of the save operation when the
+            destination exists 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
         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.
@@ -769,9 +781,10 @@ def to_csv(
         if num_files is not None:
             sdf = sdf.repartition(num_files)
 
-        builder = sdf.write.mode("overwrite")
-        OptionUtils._set_opts(
-            builder,
+        builder = sdf.write.mode(mode)
+        if partition_cols is not None:
+            builder.partitionBy(partition_cols)
+        builder._set_opts(
             path=path,
             sep=sep,
             nullValue=na_rep,
@@ -787,6 +800,8 @@ def to_json(
         path=None,
         compression="uncompressed",
         num_files=None,
+        mode: str = "overwrite",
+        partition_cols: Optional[Union[str, List[str]]] = None,
         index_col: Optional[Union[str, List[str]]] = None,
         **options
     ):
@@ -818,6 +833,17 @@ def to_json(
             compression is inferred from the filename.
         num_files : the number of files to be written in `path` directory when
             this is a path.
+        mode : str {'append', 'overwrite', 'ignore', 'error', 'errorifexists'},
+            default 'overwrite'. Specifies the behavior of the save operation when the
+            destination exists 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
         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.
@@ -875,8 +901,10 @@ def to_json(
         if num_files is not None:
             sdf = sdf.repartition(num_files)
 
-        builder = sdf.write.mode("overwrite")
-        OptionUtils._set_opts(builder, compression=compression)
+        builder = sdf.write.mode(mode)
+        if partition_cols is not None:
+            builder.partitionBy(partition_cols)
+        builder._set_opts(compression=compression)
         builder.options(**options).format("json").save(path)
 
     def to_excel(
diff --git a/databricks/koalas/spark/accessors.py b/databricks/koalas/spark/accessors.py
--- a/databricks/koalas/spark/accessors.py
+++ b/databricks/koalas/spark/accessors.py
@@ -23,8 +23,7 @@
 
 import pyspark
 from pyspark import StorageLevel
-from pyspark.sql import Column
-from pyspark.sql import DataFrame as SparkDataFrame
+from pyspark.sql import Column, DataFrame as SparkDataFrame
 
 if TYPE_CHECKING:
     import databricks.koalas as ks
@@ -569,7 +568,7 @@ def to_table(
         name: str,
         format: Optional[str] = None,
         mode: str = "overwrite",
-        partition_cols: Union[str, List[str], None] = None,
+        partition_cols: Optional[Union[str, List[str]]] = None,
         index_col: Optional[Union[str, List[str]]] = None,
         **options
     ):
@@ -640,7 +639,7 @@ def to_spark_io(
         path: Optional[str] = None,
         format: Optional[str] = None,
         mode: str = "overwrite",
-        partition_cols: Union[str, List[str], None] = None,
+        partition_cols: Optional[Union[str, List[str]]] = None,
         index_col: Optional[Union[str, List[str]]] = None,
         **options
     ):
@@ -683,7 +682,6 @@ def to_spark_io(
         DataFrame.to_spark_io
         DataFrame.spark.to_spark_io
 
-
         Examples
         --------
         >>> df = ks.DataFrame(dict(


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
@@ -359,3 +359,25 @@ def test_to_csv_with_path_and_pyspark_options(self):
         output_path = "%s/%s" % (self.tmp_dir, output_paths[0])
         with open(output_path) as f:
             self.assertEqual(f.read(), expected)
+
+    def test_to_csv_with_partition_cols(self):
+        pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]})
+        kdf = ks.DataFrame(pdf)
+
+        kdf.to_csv(self.tmp_dir, partition_cols="b", num_files=1)
+
+        partition_paths = [path for path in os.listdir(self.tmp_dir) if path.startswith("b=")]
+        assert len(partition_paths) > 0
+        for partition_path in partition_paths:
+            column, value = partition_path.split("=")
+            expected = pdf[pdf[column] == value].drop("b", axis=1).to_csv(index=False)
+
+            output_paths = [
+                path
+                for path in os.listdir("%s/%s" % (self.tmp_dir, partition_path))
+                if path.startswith("part-")
+            ]
+            assert len(output_paths) > 0
+            output_path = "%s/%s/%s" % (self.tmp_dir, partition_path, output_paths[0])
+            with open(output_path) as f:
+                self.assertEqual(f.read(), expected)
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
@@ -153,6 +153,28 @@ def test_to_json_with_path(self):
         output_path = "%s/%s" % (self.tmp_dir, output_paths[0])
         self.assertEqual("[%s]" % open(output_path).read().strip(), expected)
 
+    def test_to_json_with_partition_cols(self):
+        pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]})
+        kdf = ks.DataFrame(pdf)
+
+        kdf.to_json(self.tmp_dir, partition_cols="b", num_files=1)
+
+        partition_paths = [path for path in os.listdir(self.tmp_dir) if path.startswith("b=")]
+        assert len(partition_paths) > 0
+        for partition_path in partition_paths:
+            column, value = partition_path.split("=")
+            expected = pdf[pdf[column] == value].drop("b", axis=1).to_json(orient="records")
+
+            output_paths = [
+                path
+                for path in os.listdir("%s/%s" % (self.tmp_dir, partition_path))
+                if path.startswith("part-")
+            ]
+            assert len(output_paths) > 0
+            output_path = "%s/%s/%s" % (self.tmp_dir, partition_path, output_paths[0])
+            with open(output_path) as f:
+                self.assertEqual("[%s]" % open(output_path).read().strip(), expected)
+
     def test_to_clipboard(self):
         pdf = self.pdf
         kdf = self.kdf
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
@@ -94,7 +94,9 @@ def test_parquet_write(self):
             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]
+            actual = ks.read_parquet(tmp)
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
             self.assert_eq(
                 actual.sort_values(by="f").to_spark().toPandas(),
                 expected.sort_values(by="f").to_spark().toPandas(),
@@ -104,7 +106,9 @@ def test_parquet_write(self):
             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]
+            actual = ks.read_parquet(tmp)
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
             self.assert_eq(
                 actual.sort_values(by="f").to_spark().toPandas(),
                 expected.sort_values(by="f").to_spark().toPandas(),
@@ -119,7 +123,9 @@ def test_table(self):
             expected.spark.to_table("test_table", 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_table("test_table")[self.test_column_order]
+            actual = ks.read_table("test_table")
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
             self.assert_eq(
                 actual.sort_values(by="f").to_spark().toPandas(),
                 expected.sort_values(by="f").to_spark().toPandas(),
@@ -129,7 +135,9 @@ def test_table(self):
             expected.to_table("test_table", 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_table("test_table")[self.test_column_order]
+            actual = ks.read_table("test_table")
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
             self.assert_eq(
                 actual.sort_values(by="f").to_spark().toPandas(),
                 expected.sort_values(by="f").to_spark().toPandas(),
@@ -137,21 +145,21 @@ def test_table(self):
 
             # When index columns are known
             expected_idx = expected.set_index("bhello")[["f", "i32", "i64"]]
-            actual_idx = ks.read_table("test_table", "bhello")[["f", "i32", "i64"]]
+            actual_idx = ks.read_table("test_table", 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(),
             )
 
             expected_idx = expected.set_index(["bhello"])[["f", "i32", "i64"]]
-            actual_idx = ks.read_table("test_table", ["bhello"])[["f", "i32", "i64"]]
+            actual_idx = ks.read_table("test_table", 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(),
             )
 
             expected_idx = expected.set_index(["i32", "bhello"])[["f", "i64"]]
-            actual_idx = ks.read_table("test_table", ["i32", "bhello"])[["f", "i64"]]
+            actual_idx = ks.read_table("test_table", index_col=["i32", "bhello"])[["f", "i64"]]
             self.assert_eq(
                 actual_idx.sort_values(by="f").to_spark().toPandas(),
                 expected_idx.sort_values(by="f").to_spark().toPandas(),
@@ -166,7 +174,9 @@ def test_spark_io(self):
             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]
+            actual = ks.read_spark_io(tmp, format="json")
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
             self.assert_eq(
                 actual.sort_values(by="f").to_spark().toPandas(),
                 expected.sort_values(by="f").to_spark().toPandas(),
@@ -178,7 +188,9 @@ def test_spark_io(self):
             )
             # 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]
+            actual = ks.read_spark_io(path=tmp, format="json")
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
             self.assert_eq(
                 actual.sort_values(by="f").to_spark().toPandas(),
                 expected.sort_values(by="f").to_spark().toPandas(),