Repository: mars-project/mars

Issue:
Add `DataFrame.stack` support
<!--
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.**

`DataFrame.stack` could be added support.


Discussion:

Resolving patch:
diff --git a/mars/dataframe/base/__init__.py b/mars/dataframe/base/__init__.py
--- a/mars/dataframe/base/__init__.py
+++ b/mars/dataframe/base/__init__.py
@@ -37,6 +37,7 @@
 from .select_dtypes import select_dtypes
 from .map_chunk import map_chunk
 from .rebalance import rebalance
+from .stack import stack
 
 
 def _install():
@@ -74,6 +75,7 @@ def _install():
         setattr(t, 'select_dtypes', select_dtypes)
         setattr(t, 'map_chunk', map_chunk)
         setattr(t, 'rebalance', rebalance)
+        setattr(t, 'stack', stack)
 
     for t in SERIES_TYPE:
         setattr(t, 'to_gpu', to_gpu)
diff --git a/mars/dataframe/base/stack.py b/mars/dataframe/base/stack.py
new file mode 100644
--- /dev/null
+++ b/mars/dataframe/base/stack.py
@@ -0,0 +1,304 @@
+# 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.
+
+from typing import Union, List
+
+import numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...serialize import KeyField, AnyField, BoolField
+from ...tiles import TilesError
+from ...utils import check_chunks_unknown_shape
+from ..operands import DataFrameOperand, DataFrameOperandMixin
+from ..utils import build_df, parse_index
+
+
+class DataFrameStack(DataFrameOperand, DataFrameOperandMixin):
+    _op_type_ = opcodes.STACK
+
+    _input_df = KeyField('input_df')
+    _level = AnyField('level')
+    _dropna = BoolField('dropna')
+
+    def __init__(self, input_df=None, level=None, dropna=None, **kw):
+        super().__init__(_input_df=input_df, _level=level,
+                         _dropna=dropna, **kw)
+
+    @property
+    def input_df(self):
+        return self._input_df
+
+    @property
+    def level(self):
+        return self._level
+
+    @property
+    def dropna(self):
+        return self._dropna
+
+    def _set_inputs(self, inputs):
+        super()._set_inputs(inputs)
+        self._input_df = self._inputs[0]
+
+    @classmethod
+    def _calc_size(cls,
+                   size: int,
+                   level: Union[List, int],
+                   dtypes: pd.Series):
+        index = dtypes.index
+
+        if not isinstance(index, pd.MultiIndex):
+            return size * len(index)
+
+        if isinstance(level, int):
+            level = [level]
+        return size * np.prod([index.levshape[lev]
+                               for lev in level]).item()
+
+    def __call__(self, input_df):
+        test_df = build_df(input_df)
+        test_df = test_df.stack(level=self._level, dropna=self._dropna)
+        if self._dropna:
+            size = np.nan
+        else:
+            size = self._calc_size(input_df.shape[0], self._level, input_df.dtypes)
+        if test_df.ndim == 1:
+            shape = (size,)
+            return self.new_series([input_df], shape=shape, dtype=test_df.dtype,
+                                   index_value=parse_index(test_df.index, input_df),
+                                   name=test_df.name)
+        else:
+            shape = (size, test_df.shape[1])
+            return self.new_dataframe([input_df], shape=shape, dtypes=test_df.dtypes,
+                                      index_value=parse_index(test_df.index, input_df),
+                                      columns_value=parse_index(test_df.columns, store_data=True))
+
+    @classmethod
+    def tile(cls, op: "DataFrameStack"):
+        input_df = op.input_df
+        out = op.outputs[0]
+        out_index = out.index_value.to_pandas()
+
+        if input_df.chunk_shape[1] > 1:
+            # rechunk into 1 chunk on axis 1
+            check_chunks_unknown_shape([input_df], TilesError)
+            input_df = input_df.rechunk({1: input_df.shape[1]})._inplace_tile()
+
+        out_chunks = []
+        for c in input_df.chunks:
+            chunk_op = op.copy().reset_key()
+            if op.dropna:
+                size = np.nan
+            else:
+                size = cls._calc_size(c.shape[0], op.level, c.dtypes)
+            if out.ndim == 1:
+                kw = {
+                    'shape': (size,),
+                    'index': (c.index[0],),
+                    'dtype': out.dtype,
+                    'index_value': parse_index(out_index, c),
+                    'name': out.name
+                }
+            else:
+                kw = {
+                    'shape': (size, out.shape[1]),
+                    'index': (c.index[0], 0),
+                    'dtypes': out.dtypes,
+                    'index_value': parse_index(out_index, c),
+                    'columns_value': out.columns_value
+                }
+            out_chunk = chunk_op.new_chunk([c], **kw)
+            out_chunks.append(out_chunk)
+
+        params = out.params
+        if out.ndim == 1:
+            params['nsplits'] = (tuple(out_c.shape[0] for out_c in out_chunks),)
+        else:
+            params['nsplits'] = (tuple(out_c.shape[0] for out_c in out_chunks),
+                                 (out.shape[1],))
+        params['chunks'] = out_chunks
+        new_op = op.copy()
+        return new_op.new_tileables(op.inputs, kws=[params])
+
+    @classmethod
+    def execute(cls, ctx, op: "DataFrameStack"):
+        inp: pd.DataFrame = ctx[op.input_df.key]
+        ctx[op.outputs[0].key] = inp.stack(level=op.level, dropna=op.dropna)
+
+
+def stack(df, level=-1, dropna=True):
+    """
+    Stack the prescribed level(s) from columns to index.
+
+    Return a reshaped DataFrame or Series having a multi-level
+    index with one or more new inner-most levels compared to the current
+    DataFrame. The new inner-most levels are created by pivoting the
+    columns of the current dataframe:
+
+      - if the columns have a single level, the output is a Series;
+      - if the columns have multiple levels, the new index
+        level(s) is (are) taken from the prescribed level(s) and
+        the output is a DataFrame.
+
+    Parameters
+    ----------
+    level : int, str, list, default -1
+        Level(s) to stack from the column axis onto the index
+        axis, defined as one index or label, or a list of indices
+        or labels.
+    dropna : bool, default True
+        Whether to drop rows in the resulting Frame/Series with
+        missing values. Stacking a column level onto the index
+        axis can create combinations of index and column values
+        that are missing from the original dataframe. See Examples
+        section.
+
+    Returns
+    -------
+    DataFrame or Series
+        Stacked dataframe or series.
+
+    See Also
+    --------
+    DataFrame.unstack : Unstack prescribed level(s) from index axis
+         onto column axis.
+    DataFrame.pivot : Reshape dataframe from long format to wide
+         format.
+    DataFrame.pivot_table : Create a spreadsheet-style pivot table
+         as a DataFrame.
+
+    Notes
+    -----
+    The function is named by analogy with a collection of books
+    being reorganized from being side by side on a horizontal
+    position (the columns of the dataframe) to being stacked
+    vertically on top of each other (in the index of the
+    dataframe).
+
+    Examples
+    --------
+    **Single level columns**
+
+    >>> import mars.dataframe as md
+    >>> df_single_level_cols = md.DataFrame([[0, 1], [2, 3]],
+    ...                                     index=['cat', 'dog'],
+    ...                                     columns=['weight', 'height'])
+
+    Stacking a dataframe with a single level column axis returns a Series:
+
+    >>> df_single_level_cols.execute()
+         weight height
+    cat       0      1
+    dog       2      3
+    >>> df_single_level_cols.stack().execute()
+    cat  weight    0
+         height    1
+    dog  weight    2
+         height    3
+    dtype: int64
+
+    **Multi level columns: simple case**
+
+    >>> multicol1 = pd.MultiIndex.from_tuples([('weight', 'kg'),
+    ...                                        ('weight', 'pounds')])
+    >>> df_multi_level_cols1 = md.DataFrame([[1, 2], [2, 4]],
+    ...                                     index=['cat', 'dog'],
+    ...                                     columns=multicol1)
+
+    Stacking a dataframe with a multi-level column axis:
+
+    >>> df_multi_level_cols1.execute()
+         weight
+             kg    pounds
+    cat       1        2
+    dog       2        4
+    >>> df_multi_level_cols1.stack().execute()
+                weight
+    cat kg           1
+        pounds       2
+    dog kg           2
+        pounds       4
+
+    **Missing values**
+
+    >>> multicol2 = pd.MultiIndex.from_tuples([('weight', 'kg'),
+    ...                                        ('height', 'm')])
+    >>> df_multi_level_cols2 = md.DataFrame([[1.0, 2.0], [3.0, 4.0]],
+    ...                                     index=['cat', 'dog'],
+    ...                                     columns=multicol2)
+
+    It is common to have missing values when stacking a dataframe
+    with multi-level columns, as the stacked dataframe typically
+    has more values than the original dataframe. Missing values
+    are filled with NaNs:
+
+    >>> df_multi_level_cols2.execute()
+        weight height
+            kg      m
+    cat    1.0    2.0
+    dog    3.0    4.0
+    >>> df_multi_level_cols2.stack().execute()
+            height  weight
+    cat kg     NaN     1.0
+        m      2.0     NaN
+    dog kg     NaN     3.0
+        m      4.0     NaN
+
+    **Prescribing the level(s) to be stacked**
+
+    The first parameter controls which level or levels are stacked:
+
+    >>> df_multi_level_cols2.stack(0).execute()
+                 kg    m
+    cat height  NaN  2.0
+        weight  1.0  NaN
+    dog height  NaN  4.0
+        weight  3.0  NaN
+    >>> df_multi_level_cols2.stack([0, 1]).execute()
+    cat  height  m     2.0
+         weight  kg    1.0
+    dog  height  m     4.0
+         weight  kg    3.0
+    dtype: float64
+
+    **Dropping missing values**
+
+    >>> df_multi_level_cols3 = md.DataFrame([[None, 1.0], [2.0, 3.0]],
+    ...                                     index=['cat', 'dog'],
+    ...                                     columns=multicol2)
+
+    Note that rows where all values are missing are dropped by
+    default but this behaviour can be controlled via the dropna
+    keyword parameter:
+
+    >>> df_multi_level_cols3.execute()
+        weight height
+            kg      m
+    cat    NaN    1.0
+    dog    2.0    3.0
+    >>> df_multi_level_cols3.stack(dropna=False).execute()
+            height  weight
+    cat kg     NaN     NaN
+        m      1.0     NaN
+    dog kg     NaN     2.0
+        m      3.0     NaN
+    >>> df_multi_level_cols3.stack(dropna=True).execute()
+            height  weight
+    cat m      1.0     NaN
+    dog kg     NaN     2.0
+        m      3.0     NaN
+    """
+    op = DataFrameStack(input_df=df, level=level, dropna=dropna)
+    return op(df)


Test patch:
diff --git a/mars/dataframe/base/tests/test_base_execution.py b/mars/dataframe/base/tests/test_base_execution.py
--- a/mars/dataframe/base/tests/test_base_execution.py
+++ b/mars/dataframe/base/tests/test_base_execution.py
@@ -1630,3 +1630,32 @@ def testRebalanceExecution(self):
             result = self.executor.execute_dataframe(df2)
             self.assertEqual(len(result), 2)
             pd.testing.assert_frame_equal(pd.concat(result), raw)
+
+    def testStackExecution(self):
+        raw = pd.DataFrame(np.random.rand(10, 3), columns=list('abc'),
+                           index=[f's{i}' for i in range(10)])
+        for loc in [(5, 1), (8, 2), (1, 0)]:
+            raw.iloc[loc] = np.nan
+        df = from_pandas_df(raw, chunk_size=(5, 2))
+
+        for dropna in (True, False):
+            r = df.stack(dropna=dropna)
+            result = self.executor.execute_dataframe(r, concat=True)[0]
+            expected = raw.stack(dropna=dropna)
+            pd.testing.assert_series_equal(result, expected)
+
+        cols = pd.MultiIndex.from_tuples([
+            ('c1', 'cc1'), ('c1', 'cc2'), ('c2', 'cc3')])
+        raw2 = raw.copy()
+        raw2.columns = cols
+        df = from_pandas_df(raw2, chunk_size=(5, 2))
+
+        for level in [-1, 0, [0, 1]]:
+            for dropna in (True, False):
+                r = df.stack(level=level, dropna=dropna)
+                result = self.executor.execute_dataframe(r, concat=True)[0]
+                expected = raw2.stack(level=level, dropna=dropna)
+                assert_method = \
+                    pd.testing.assert_series_equal if expected.ndim == 1 \
+                        else pd.testing.assert_frame_equal
+                assert_method(result, expected)