Repository: mars-project/mars

Issue:
Support `md.to_parquet`
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**

`to_parquet` could be added to Mars DataFrame.


Discussion:
The package is simply not usable in the real world for dataframe manipulation without this support. Supporting multiple engines including `pyarrow` can also matter.
> The package is simply not usable in the real world for dataframe manipulation without this support. Supporting multiple engines including `pyarrow` can also matter.

Thanks, we'd love to hear the voice from users, I have elevated the priority from P1 to P0, we will implement this feature ASAP.

Resolving patch:
diff --git a/mars/dataframe/datasource/read_parquet.py b/mars/dataframe/datasource/read_parquet.py
--- a/mars/dataframe/datasource/read_parquet.py
+++ b/mars/dataframe/datasource/read_parquet.py
@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import pickle
 
 import numpy as np
 import pandas as pd
@@ -33,7 +35,7 @@
 from ...config import options
 from ...filesystem import open_file, glob
 from ...serialize import AnyField, BoolField, DictField, ListField,\
-    StringField, Int32Field
+    StringField, Int32Field, BytesField
 from ..arrays import ArrowStringDtype
 from ..operands import DataFrameOperandMixin, DataFrameOperand, OutputType
 from ..utils import parse_index, to_arrow_dtypes, contain_arrow_dtype, \
@@ -138,15 +140,22 @@ class DataFrameReadParquet(DataFrameOperand, DataFrameOperandMixin):
     _incremental_index = BoolField('incremental_index')
     _storage_options = DictField('storage_options')
 
+    # for chunk
+    _partitions = BytesField('partitions')
+    _partition_keys = ListField('partition_keys')
+
     def __init__(self, path=None, engine=None, columns=None, use_arrow_dtype=None,
                  groups_as_chunks=None, group_index=None, incremental_index=None,
-                 read_kwargs=None, storage_options=None, **kw):
+                 read_kwargs=None, partitions=None, partition_keys=None,
+                 storage_options=None, **kw):
         super().__init__(_path=path, _engine=engine, _columns=columns,
                          _use_arrow_dtype=use_arrow_dtype,
                          _groups_as_chunks=groups_as_chunks,
                          _group_index=group_index,
                          _read_kwargs=read_kwargs,
                          _incremental_index=incremental_index,
+                         _partitions=partitions,
+                         _partition_keys=partition_keys,
                          _storage_options=storage_options,
                          _output_types=[OutputType.dataframe], **kw)
 
@@ -182,22 +191,62 @@ def read_kwargs(self):
     def incremental_index(self):
         return self._incremental_index
 
+    @property
+    def partitions(self):
+        return self._partitions
+
+    @property
+    def partition_keys(self):
+        return self._partition_keys
+
     @property
     def storage_options(self):
         return self._storage_options
 
     @classmethod
-    def tile(cls, op):
-        chunk_index = 0
-        out_chunks = []
-        out_df = op.outputs[0]
-
-        dtypes = out_df.dtypes
+    def _to_arrow_dtypes(cls, dtypes, op):
         if op.use_arrow_dtype is None and not op.gpu and \
                 options.dataframe.use_arrow_dtype:  # pragma: no cover
             # check if use_arrow_dtype set on the server side
-            dtypes = to_arrow_dtypes(out_df.dtypes)
+            dtypes = to_arrow_dtypes(dtypes)
+        return dtypes
+
+    @classmethod
+    def _tile_partitioned(cls, op):
+        out_df = op.outputs[0]
+        shape = (np.nan, out_df.shape[1])
+        dtypes = cls._to_arrow_dtypes(out_df.dtypes, op)
+        dataset = pq.ParquetDataset(op.path)
+
+        chunk_index = 0
+        out_chunks = []
+        for piece in dataset.pieces:
+            chunk_op = op.copy().reset_key()
+            chunk_op._path = piece.path
+            chunk_op._partitions = pickle.dumps(dataset.partitions)
+            chunk_op._partition_keys = piece.partition_keys
+            new_chunk = chunk_op.new_chunk(
+                None, shape=shape, index=(chunk_index, 0),
+                index_value=out_df.index_value,
+                columns_value=out_df.columns_value,
+                dtypes=dtypes)
+            out_chunks.append(new_chunk)
+            chunk_index += 1
+
+        new_op = op.copy()
+        nsplits = ((np.nan,) * len(out_chunks), (out_df.shape[1],))
+        return new_op.new_dataframes(None, out_df.shape, dtypes=dtypes,
+                                     index_value=out_df.index_value,
+                                     columns_value=out_df.columns_value,
+                                     chunks=out_chunks, nsplits=nsplits)
+
+    @classmethod
+    def _tile_no_partitioned(cls, op):
+        chunk_index = 0
+        out_chunks = []
+        out_df = op.outputs[0]
 
+        dtypes = cls._to_arrow_dtypes(out_df.dtypes, op)
         shape = (np.nan, out_df.shape[1])
 
         paths = op.path if isinstance(op.path, (tuple, list)) else \
@@ -237,11 +286,30 @@ def tile(cls, op):
                                      columns_value=out_df.columns_value,
                                      chunks=out_chunks, nsplits=nsplits)
 
+    @classmethod
+    def tile(cls, op):
+        if os.path.isdir(op.path):
+            return cls._tile_partitioned(op)
+        else:
+            return cls._tile_no_partitioned(op)
+
+    @classmethod
+    def _execute_partitioned(cls, ctx, op):
+        out = op.outputs[0]
+        partitions = pickle.loads(op.partitions)
+        piece = pq.ParquetDatasetPiece(op.path, partition_keys=op.partition_keys,
+                                       open_file_func=open_file)
+        df = piece.read(partitions=partitions).to_pandas()
+        ctx[out.key] = df
+
     @classmethod
     def execute(cls, ctx, op):
         out = op.outputs[0]
         path = op.path
 
+        if op.partitions is not None:
+            return cls._execute_partitioned(ctx, op)
+
         engine = get_engine(op.engine)
         with open_file(path, storage_options=op.storage_options) as f:
             use_arrow_dtype = contain_arrow_dtype(out.dtypes)
@@ -305,26 +373,37 @@ def read_parquet(path, engine: str = "auto", columns=None,
     Mars DataFrame
     """
 
-    if not isinstance(path, list):
-        file_path = glob(path, storage_options=storage_options)[0]
-    else:
-        file_path = path[0]
-
     engine_type = check_engine(engine)
     engine = get_engine(engine_type)
-    with open_file(file_path, storage_options=storage_options) as f:
-        dtypes = engine.read_dtypes(f)
 
-    if columns:
-        dtypes = dtypes[columns]
+    if os.path.isdir(path):
+        # If path is a directory, we will read as a partitioned datasets.
+        if engine_type != 'pyarrow':
+            raise TypeError('Only support pyarrow engine when reading from'
+                            'partitioned datasets.')
+        dataset = pq.ParquetDataset(path)
+        dtypes = dataset.schema.to_arrow_schema().empty_table().to_pandas().dtypes
+        for partition in dataset.partitions:
+            dtypes[partition.name] = pd.CategoricalDtype()
+    else:
+        if not isinstance(path, list):
+            file_path = glob(path, storage_options=storage_options)[0]
+        else:
+            file_path = path[0]
+
+        with open_file(file_path, storage_options=storage_options) as f:
+            dtypes = engine.read_dtypes(f)
+
+        if columns:
+            dtypes = dtypes[columns]
+
+        if use_arrow_dtype is None:
+            use_arrow_dtype = options.dataframe.use_arrow_dtype
+        if use_arrow_dtype:
+            dtypes = to_arrow_dtypes(dtypes)
 
     index_value = parse_index(pd.RangeIndex(-1))
     columns_value = parse_index(dtypes.index, store_data=True)
-    if use_arrow_dtype is None:
-        use_arrow_dtype = options.dataframe.use_arrow_dtype
-    if use_arrow_dtype:
-        dtypes = to_arrow_dtypes(dtypes)
-
     op = DataFrameReadParquet(path=path, engine=engine_type, columns=columns,
                               groups_as_chunks=groups_as_chunks,
                               use_arrow_dtype=use_arrow_dtype,
diff --git a/mars/dataframe/datastore/__init__.py b/mars/dataframe/datastore/__init__.py
--- a/mars/dataframe/datastore/__init__.py
+++ b/mars/dataframe/datastore/__init__.py
@@ -17,11 +17,13 @@ def _install():
     from .to_csv import to_csv
     from .to_sql import to_sql
     from .to_vineyard import to_vineyard
+    from .to_parquet import to_parquet
     from ..operands import DATAFRAME_TYPE, SERIES_TYPE
 
     for cls in DATAFRAME_TYPE:
         setattr(cls, 'to_csv', to_csv)
         setattr(cls, 'to_sql', to_sql)
+        setattr(cls, 'to_parquet', to_parquet)
         setattr(cls, 'to_vineyard', to_vineyard)
 
     for cls in SERIES_TYPE:
diff --git a/mars/dataframe/datastore/to_parquet.py b/mars/dataframe/datastore/to_parquet.py
new file mode 100644
--- /dev/null
+++ b/mars/dataframe/datastore/to_parquet.py
@@ -0,0 +1,211 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pandas as pd
+
+from ... import opcodes as OperandDef
+from ...filesystem import open_file
+from ...serialize import KeyField, AnyField, StringField, ListField, \
+    BoolField, DictField
+from ..operands import DataFrameOperand, DataFrameOperandMixin
+from ..utils import parse_index
+from ..datasource.read_parquet import check_engine
+
+try:
+    import pyarrow.parquet as pq
+    import pyarrow as pa
+except ImportError:  # pragma: no cover
+    pq = None
+    pa = None
+
+
+class DataFrameToParquet(DataFrameOperand, DataFrameOperandMixin):
+    _op_type_ = OperandDef.TO_PARQUET
+
+    _input = KeyField('input')
+    _path = AnyField('path')
+    _engine = StringField('engine')
+    _index = BoolField('index')
+    _compression = AnyField('compression')
+    _partition_cols = ListField('partition_cols')
+    _additional_kwargs = DictField('additional_kwargs')
+    _storage_options = DictField('storage_options')
+
+    def __init__(self, path=None, engine=None, index=None, compression=None,
+                 partition_cols=None, storage_options=None,
+                 additional_kwargs=None, **kw):
+        super().__init__(_path=path, _engine=engine, _index=index,
+                         _compression=compression, _partition_cols=partition_cols,
+                         _storage_options=storage_options,
+                         _additional_kwargs=additional_kwargs,
+                         **kw)
+
+    @property
+    def input(self):
+        return self._input
+
+    @property
+    def path(self):
+        return self._path
+
+    @property
+    def engine(self):
+        return self._engine
+
+    @property
+    def index(self):
+        return self._index
+
+    @property
+    def compression(self):
+        return self._compression
+
+    @property
+    def partition_cols(self):
+        return self._partition_cols
+
+    @property
+    def storage_options(self):
+        return self._storage_options
+
+    @property
+    def additional_kwargs(self):
+        return self._additional_kwargs
+
+    def _set_inputs(self, inputs):
+        super()._set_inputs(inputs)
+        self._input = self._inputs[0]
+
+    @classmethod
+    def _get_path(cls, path, i):
+        if '*' not in path:
+            return path
+        return path.replace('*', str(i))
+
+    @classmethod
+    def tile(cls, op):
+        in_df = op.input
+        out_df = op.outputs[0]
+
+        # make sure only 1 chunk on the column axis
+        in_df = in_df.rechunk({1: in_df.shape[1]})._inplace_tile()
+
+        out_chunks = []
+        for chunk in in_df.chunks:
+            chunk_op = op.copy().reset_key()
+            index_value = parse_index(chunk.index_value.to_pandas()[:0], chunk)
+            out_chunk = chunk_op.new_chunk([chunk], shape=(0, 0),
+                                           index_value=index_value,
+                                           columns_value=out_df.columns_value,
+                                           dtypes=out_df.dtypes,
+                                           index=chunk.index)
+            out_chunks.append(out_chunk)
+
+        new_op = op.copy()
+        params = out_df.params.copy()
+        params.update(dict(chunks=out_chunks, nsplits=((0,) * in_df.chunk_shape[0], (0,))))
+        return new_op.new_tileables([in_df], **params)
+
+    @classmethod
+    def execute(cls, ctx, op):
+        df = ctx[op.input.key]
+        out = op.outputs[0]
+        path = cls._get_path(op.path, op.outputs[0].index[0])
+        if op.partition_cols is None:
+            if op.engine == 'fastparquet':
+                df.to_parquet(path, engine=op.engine, compression=op.compression,
+                              index=op.index, open_with=open_file, **op.additional_kwargs)
+            else:
+                with open_file(path, mode='wb', storage_options=op.storage_options) as f:
+                    df.to_parquet(f, engine=op.engine, compression=op.compression,
+                                  index=op.index, **op.additional_kwargs or dict())
+        else:
+            if op.engine == 'pyarrow':
+                pq.write_to_dataset(pa.Table.from_pandas(df), path,
+                                    partition_cols=op.partition_cols)
+            else:  # pragma: no cover
+                raise NotImplementedError('Only support pyarrow engine when '
+                                          'specify `partition_cols`.')
+
+        ctx[out.key] = pd.DataFrame()
+
+    def __call__(self, df):
+        index_value = parse_index(df.index_value.to_pandas()[:0], df)
+        columns_value = parse_index(df.columns_value.to_pandas()[:0], store_data=True)
+        return self.new_dataframe([df], shape=(0, 0), dtypes=df.dtypes[:0],
+                                  index_value=index_value, columns_value=columns_value)
+
+
+def to_parquet(df, path, engine='auto', compression='snappy', index=None,
+               partition_cols=None, **kwargs):
+    """
+    Write a DataFrame to the binary parquet format, each chunk will be
+    written to a Parquet file.
+
+    Parameters
+    ----------
+    path : str or file-like object
+        If path is a string with wildcard e.g. '/to/path/out-*.parquet',
+        `to_parquet` will try to write multiple files, for instance,
+        chunk (0, 0) will write data into '/to/path/out-0.parquet'.
+        If path is a string without wildcard, we will treat it as a directory.
+
+    engine : {'auto', 'pyarrow', 'fastparquet'}, default 'auto'
+        Parquet library to use. The default behavior is to try 'pyarrow',
+        falling back to 'fastparquet' if 'pyarrow' is unavailable.
+
+    compression : {'snappy', 'gzip', 'brotli', None}, default 'snappy'
+        Name of the compression to use. Use ``None`` for no compression.
+
+    index : bool, default None
+        If ``True``, include the dataframe's index(es) in the file output.
+        If ``False``, they will not be written to the file.
+        If ``None``, similar to ``True`` the dataframe's index(es)
+        will be saved. However, instead of being saved as values,
+        the RangeIndex will be stored as a range in the metadata so it
+        doesn't require much space and is faster. Other indexes will
+        be included as columns in the file output.
+
+    partition_cols : list, optional, default None
+        Column names by which to partition the dataset.
+        Columns are partitioned in the order they are given.
+        Must be None if path is not a string.
+
+    **kwargs
+        Additional arguments passed to the parquet library.
+
+    Examples
+    --------
+    >>> import mars.dataframe as md
+    >>> df = md.DataFrame(data={'col1': [1, 2], 'col2': [3, 4]})
+    >>> df.to_parquet('*.parquet.gzip',
+    ...               compression='gzip').execute()  # doctest: +SKIP
+    >>> md.read_parquet('*.parquet.gzip').execute()  # doctest: +SKIP
+       col1  col2
+    0     1     3
+    1     2     4
+
+    >>> import io
+    >>> f = io.BytesIO()
+    >>> df.to_parquet(f).execute()
+    >>> f.seek(0)
+    0
+    >>> content = f.read()
+    """
+    engine = check_engine(engine)
+    op = DataFrameToParquet(path=path, engine=engine, compression=compression, index=index,
+                            partition_cols=partition_cols, additional_kwargs=kwargs)
+    return op(df)


Test patch:
diff --git a/mars/dataframe/datastore/tests/test_datastore_execute.py b/mars/dataframe/datastore/tests/test_datastore_execute.py
--- a/mars/dataframe/datastore/tests/test_datastore_execute.py
+++ b/mars/dataframe/datastore/tests/test_datastore_execute.py
@@ -34,6 +34,14 @@
     import sqlalchemy
 except ImportError:
     sqlalchemy = None
+try:
+    import pyarrow as pa
+except ImportError:
+    pa = None
+try:
+    import fastparquet
+except ImportError:
+    fastparquet = None
 
 
 class Test(TestBase):
@@ -163,3 +171,44 @@ def testWithGivenSession(session):
 
             with new_session('http://' + cluster._web_endpoint).as_default() as web_session:
                 testWithGivenSession(web_session)
+
+    @unittest.skipIf(pa is None or fastparquet is None, 'pyarrow or fastparquet not installed')
+    def testToParquetExecution(self):
+        raw = pd.DataFrame({
+            'col1': np.random.rand(100),
+            'col2': np.arange(100),
+            'col3': np.random.choice(['a', 'b', 'c'], (100,)),
+        })
+        df = DataFrame(raw, chunk_size=33)
+
+        with tempfile.TemporaryDirectory() as base_path:
+            # DATAFRAME TESTS
+            path = os.path.join(base_path, 'out-*.parquet')
+            r = df.to_parquet(path)
+            self.executor.execute_dataframe(r)
+
+            read_df = md.read_parquet(path)
+            result = self.executor.execute_dataframe(read_df, concat=True)[0]
+            result = result.sort_index()
+            pd.testing.assert_frame_equal(result, raw)
+
+            # test fastparquet
+            path = os.path.join(base_path, 'out-fastparquet-*.parquet')
+            r = df.to_parquet(path, engine='fastparquet', compression='gzip')
+            self.executor.execute_dataframe(r)
+
+            read_df = md.read_parquet(path)
+            result = self.executor.execute_dataframe(read_df, concat=True)[0]
+            result = result.sort_index()
+            pd.testing.assert_frame_equal(result, raw)
+
+            # test partition_cols
+            path = os.path.join(base_path, 'out-partitioned')
+            r = df.to_parquet(path, partition_cols=['col3'])
+            self.executor.execute_dataframe(r)
+
+            read_df = md.read_parquet(path)
+            result = self.executor.execute_dataframe(read_df, concat=True)[0]
+            result['col3'] = result['col3'].astype('object')
+            pd.testing.assert_frame_equal(result.sort_values('col1').reset_index(drop=True),
+                                          raw.sort_values('col1').reset_index(drop=True))