Repository: dask/dask

Issue:
dask dataframe to_parquet method with engine=pyarrow fails trying to write metadata
(version 0.16)

dask dataframe to_parquet method seems to fail because at some point the compression argument becomes part of kwargs.

    def _write_partition_pyarrow(df, open_with, filename, write_index, compression, metadata_path=None, **kwargs):
        import pyarrow as pa
        from pyarrow import parquet
        t = pa.Table.from_pandas(df, preserve_index=write_index)

        with open_with(filename, 'wb') as fil:
            parquet.write_table(t, fil, **kwargs)

        if metadata_path is not None:
            with open_with(metadata_path, 'wb') as fil:
                parquet.write_metadata(t.schema, fil, **kwargs)

The method above belongs to io/parquet.py (from dask project module), and as you can see, **kwargs is used in parquet.write_table(t, fil, **kwargs), and  parquet.write_metadata(t.schema, fil, **kwargs). However, parquet.write_metadata does not expect a compression argument. In fact the method signature is the following:

    def write_metadata(schema, where, version='1.0', use_deprecated_int96_timestamps=False,
                   coerce_timestamps=None)

the signature comes from pyarrow/parquet.py module (from pyarrow project)


Discussion:

Resolving patch:
diff --git a/dask/dataframe/io/parquet.py b/dask/dataframe/io/parquet.py
--- a/dask/dataframe/io/parquet.py
+++ b/dask/dataframe/io/parquet.py
@@ -15,13 +15,12 @@
 from ...delayed import delayed
 from ...bytes.core import get_fs_paths_myopen
 
-
 __all__ = ('read_parquet', 'to_parquet')
 
 
 def _meta_from_dtypes(to_read_columns, file_columns, file_dtypes, index_cols):
     meta = pd.DataFrame({c: pd.Series([], dtype=d)
-                        for (c, d) in file_dtypes.items()},
+                         for (c, d) in file_dtypes.items()},
                         columns=[c for c in file_columns
                                  if c in file_dtypes])
     df = meta[list(to_read_columns)]
@@ -64,8 +63,8 @@ def _read_fastparquet(fs, paths, myopen, columns=None, filters=None,
     name = 'read-parquet-' + tokenize(pf, columns, categories)
 
     rgs = [rg for rg in pf.row_groups if
-           not(fastparquet.api.filter_out_stats(rg, filters, pf.schema)) and
-           not(fastparquet.api.filter_out_cats(rg, filters))]
+           not (fastparquet.api.filter_out_stats(rg, filters, pf.schema)) and
+           not (fastparquet.api.filter_out_cats(rg, filters))]
 
     if index is False:
         index_col = None
@@ -96,7 +95,7 @@ def _read_fastparquet(fs, paths, myopen, columns=None, filters=None,
     for cat in categories:
         if cat in meta:
             meta[cat] = pd.Series(pd.Categorical([],
-                                  categories=[UNKNOWN_CATEGORIES]),
+                                                 categories=[UNKNOWN_CATEGORIES]),
                                   index=meta.index)
 
     if out_type == Series:
@@ -185,8 +184,7 @@ def _write_fastparquet(df, path, write_index=None, append=False,
     sep = fs.sep
 
     object_encoding = kwargs.pop('object_encoding', 'utf8')
-    if object_encoding == 'infer' or (isinstance(object_encoding, dict) and
-                                      'infer' in object_encoding.values()):
+    if object_encoding == 'infer' or (isinstance(object_encoding, dict) and 'infer' in object_encoding.values()):
         raise ValueError('"infer" not allowed as object encoding, '
                          'because this required data in memory.')
 
@@ -203,8 +201,7 @@ def _write_fastparquet(df, path, write_index=None, append=False,
         if pf.file_scheme not in ['hive', 'empty', 'flat']:
             raise ValueError('Requested file scheme is hive, '
                              'but existing file scheme is not.')
-        elif ((set(pf.columns) != set(df.columns) - set(partition_on)) or
-              (set(partition_on) != set(pf.cats))):
+        elif ((set(pf.columns) != set(df.columns) - set(partition_on)) or (set(partition_on) != set(pf.cats))):
             raise ValueError('Appended columns not the same.\n'
                              'New: {} | Previous: {}'
                              .format(pf.columns, list(df.columns)))
@@ -358,7 +355,7 @@ def _get_pyarrow_dtypes(schema):
 def _read_pyarrow_parquet_piece(open_file_func, piece, columns, index_cols,
                                 is_series, partitions):
     with open_file_func(piece.path, mode='rb') as f:
-        table = piece.read(columns=columns,  partitions=partitions,
+        table = piece.read(columns=columns, partitions=partitions,
                            use_pandas_metadata=True,
                            file=f)
     df = table.to_pandas()
@@ -416,6 +413,7 @@ def _write_partition_pyarrow(df, open_with, filename, write_index,
 
     if metadata_path is not None:
         with open_with(metadata_path, 'wb') as fil:
+            kwargs.pop('compression', None)
             parquet.write_metadata(t.schema, fil, **kwargs)
 
 


Test patch:
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -23,7 +23,6 @@
 except ImportError:
     pq = False
 
-
 df = pd.DataFrame({'x': [6, 2, 3, 4, 5],
                    'y': [1.0, 2.0, 1.0, 2.0, 1.0]},
                   index=pd.Index([10, 20, 30, 40, 50], name='myindex'))
@@ -402,9 +401,9 @@ def test_read_parquet_custom_columns(tmpdir, engine):
     (pd.DataFrame({'x': [3, 1, 2]}, index=[3, 2, 1]), {}, {}),
     (pd.DataFrame({'x': [3, 1, 5]}, index=pd.Index([1, 2, 3], name='foo')), {}, {}),
     (pd.DataFrame({'x': [1, 2, 3],
-                  'y': [3, 2, 1]}), {}, {}),
+                   'y': [3, 2, 1]}), {}, {}),
     (pd.DataFrame({'x': [1, 2, 3],
-                  'y': [3, 2, 1]}, columns=['y', 'x']), {}, {}),
+                   'y': [3, 2, 1]}, columns=['y', 'x']), {}, {}),
     (pd.DataFrame({'0': [3, 2, 1]}), {}, {}),
     (pd.DataFrame({'x': [3, 2, None]}), {}, {}),
     (pd.DataFrame({'-': [3., 2., None]}), {}, {}),
@@ -597,3 +596,19 @@ def test_parquet_select_cats(tmpdir):
     assert list(rddf.columns) == ['ints']
     rddf = dd.read_parquet(fn)
     assert list(rddf.columns) == list(df)
+
+
+@pytest.mark.parametrize('compression,', ['default', None, 'gzip', 'snappy'])
+def test_writing_parquet_with_compression(tmpdir, compression, engine):
+    fn = str(tmpdir)
+
+    if engine == 'fastparquet' and compression == 'snappy':
+        pytest.importorskip('snappy')
+
+    df = pd.DataFrame({'x': ['a', 'b', 'c'] * 10,
+                       'y': [1, 2, 3] * 10})
+    ddf = dd.from_pandas(df, npartitions=3)
+
+    ddf.to_parquet(fn, compression=compression, engine=engine)
+    out = dd.read_parquet(fn, engine=engine)
+    assert_eq(out, df, check_index=(engine != 'fastparquet'))