[REVIEW] Adding to_parquet and write_partition definitions to dask_cudf

CHANGELOG.md

@@ -156,6 +156,7 @@
 - PR #3340 Make all benchmarks use cudf base fixture to initialize RMM pool
 - PR #3337 Fix Java to pad validity buffers to 64-byte boundary
 - PR #3357 Disabling `column_view` iterators for non fixed-width types
+- PR #3369 Add write_partition to dask_cudf to fix to_parquet bug
 
 
 # cuDF 0.10.0 (16 Oct 2019)

python/dask_cudf/dask_cudf/core.py

@@ -263,6 +263,12 @@ def drop(df, keep_keys):
         results = [p for i, p in enumerate(parts) if uniques[i]]
         return from_delayed(results, meta=self._meta).reset_index()
 
+    def to_parquet(self, path, *args, **kwargs):
+        """ Calls dask.dataframe.io.to_parquet with CudfEngine backend """
+        from dask_cudf.io import to_parquet
+
+        to_parquet(self, path, *args, **kwargs)
+
     @derived_from(pd.DataFrame)
     def var(
         self,

python/dask_cudf/dask_cudf/io/__init__.py

@@ -3,6 +3,6 @@
 from .orc import read_orc
 
 try:
-    from .parquet import read_parquet
+    from .parquet import read_parquet, to_parquet
 except ImportError:
     pass

python/dask_cudf/dask_cudf/io/parquet.py

@@ -88,6 +88,56 @@ def read_partition(
 
         return df
 
+    @staticmethod
+    def write_partition(
+        df,
+        path,
+        fs,
+        filename,
+        partition_on,
+        return_metadata,
+        fmd=None,
+        compression=None,
+        index_cols=None,
+        **kwargs,
+    ):
+        # TODO: Replace `pq.write_table` with gpu-accelerated
+        #       write after cudf.io.to_parquet is supported.
+
+        md_list = []
+        preserve_index = False
+        if index_cols:
+            df = df.set_index(index_cols)
+            preserve_index = True
+
+        # NOTE: `to_arrow` does not accept `schema` argument
+        t = df.to_arrow(preserve_index=preserve_index)
+        if partition_on:
+            pq.write_to_dataset(
+                t,
+                path,
+                partition_cols=partition_on,
+                filesystem=fs,
+                metadata_collector=md_list,
+                **kwargs,
+            )
+        else:
+            with fs.open(fs.sep.join([path, filename]), "wb") as fil:
+                pq.write_table(
+                    t,
+                    fil,
+                    compression=compression,
+                    metadata_collector=md_list,
+                    **kwargs,
+                )
+            if md_list:
+                md_list[0].set_file_path(filename)
+        # Return the schema needed to write the metadata
+        if return_metadata:
+            return [{"schema": t.schema, "meta": md_list[0]}]
+        else:
+            return []
+
 
 def read_parquet(path, **kwargs):
     """ Read parquet files into a Dask DataFrame
@@ -112,3 +162,6 @@ class to support full functionality.
     if isinstance(columns, str):
         columns = [columns]
     return dd.read_parquet(path, columns=columns, engine=CudfEngine, **kwargs)
+
+
+to_parquet = partial(dd.to_parquet, engine=CudfEngine)

python/dask_cudf/dask_cudf/io/tests/test_parquet.py

@@ -60,6 +60,17 @@ def test_roundtrip_from_dask(tmpdir):
     assert_eq(ddf[["y"]], ddf2)
 
 
+def test_roundtrip_from_dask_cudf(tmpdir):
+    tmpdir = str(tmpdir)
+    gddf = dask_cudf.from_dask_dataframe(ddf)
+    gddf.to_parquet(tmpdir)
+
+    # NOTE: Need `.compute()` to resolve correct index
+    #       name after `from_dask_dataframe`
+    gddf2 = dask_cudf.read_parquet(tmpdir)
+    assert_eq(gddf.compute(), gddf2)
+
+
 def test_roundtrip_from_pandas(tmpdir):
     fn = str(tmpdir.join("test.parquet"))