Repository: databricks/koalas

Issue:
DataFrame.stack
Implement `DataFrame.stack`.

https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.stack.html


Discussion:

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -17,7 +17,7 @@
 """
 A wrapper class for Spark DataFrame to behave similar to pandas DataFrame.
 """
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
 from distutils.version import LooseVersion
 import re
 import warnings
@@ -7965,6 +7965,186 @@ def melt(self, id_vars=None, value_vars=None, var_name=None, value_name="value")
 
         return DataFrame(exploded_df)
 
+    def stack(self):
+        """
+        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.
+
+        The new index levels are sorted.
+
+        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**
+
+        >>> df_single_level_cols = ks.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
+             weight  height
+        cat       0       1
+        dog       2       3
+        >>> df_single_level_cols.stack().sort_index()
+        cat  height    1
+             weight    0
+        dog  height    3
+             weight    2
+        Name: 0, dtype: int64
+
+        **Multi level columns: simple case**
+
+        >>> multicol1 = pd.MultiIndex.from_tuples([('weight', 'kg'),
+        ...                                        ('weight', 'pounds')])
+        >>> df_multi_level_cols1 = ks.DataFrame([[1, 2], [2, 4]],
+        ...                                     index=['cat', 'dog'],
+        ...                                     columns=multicol1)
+
+        Stacking a dataframe with a multi-level column axis:
+
+        >>> df_multi_level_cols1  # doctest: +NORMALIZE_WHITESPACE
+            weight
+                kg pounds
+        cat      1      2
+        dog      2      4
+        >>> df_multi_level_cols1.stack().sort_index()
+                    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 = ks.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
+            weight height
+                kg      m
+        cat    1.0    2.0
+        dog    3.0    4.0
+        >>> df_multi_level_cols2.stack().sort_index()  # doctest: +SKIP
+                height  weight
+        cat kg     NaN     1.0
+            m      2.0     NaN
+        dog kg     NaN     3.0
+            m      4.0     NaN
+        """
+        from databricks.koalas.series import _col
+
+        if len(self._internal.column_labels) == 0:
+            return DataFrame(self._internal.with_filter(F.lit(False)))
+
+        column_labels = defaultdict(dict)
+        index_values = set()
+        should_returns_series = False
+        for label in self._internal.column_labels:
+            new_label = label[:-1]
+            if len(new_label) == 0:
+                new_label = ("0",)
+                should_returns_series = True
+            value = label[-1]
+
+            scol = self._internal.scol_for(label)
+            column_labels[new_label][value] = scol
+
+            index_values.add(value)
+
+        column_labels = OrderedDict(sorted(column_labels.items(), key=lambda x: x[0]))
+
+        if self._internal.column_label_names is None:
+            column_label_names = None
+            index_name = None
+        else:
+            column_label_names = self._internal.column_label_names[:-1]
+            if self._internal.column_label_names[-1] is None:
+                index_name = None
+            else:
+                index_name = (self._internal.column_label_names[-1],)
+
+        index_column = SPARK_INDEX_NAME_FORMAT(len(self._internal.index_map))
+        index_map = self._internal.index_map + [(index_column, index_name)]
+        data_columns = [name_like_string(label) for label in column_labels]
+
+        structs = [
+            F.struct(
+                [F.lit(value).alias(index_column)]
+                + [
+                    (
+                        column_labels[label][value]
+                        if value in column_labels[label]
+                        else F.lit(None)
+                    ).alias(name)
+                    for label, name in zip(column_labels, data_columns)
+                ]
+            ).alias(value)
+            for value in index_values
+        ]
+
+        pairs = F.explode(F.array(structs))
+
+        sdf = self._sdf.withColumn("pairs", pairs)
+        sdf = sdf.select(
+            self._internal.index_scols
+            + [sdf["pairs"][index_column].alias(index_column)]
+            + [sdf["pairs"][name].alias(name) for name in data_columns]
+        )
+
+        internal = _InternalFrame(
+            sdf=sdf,
+            index_map=index_map,
+            column_labels=list(column_labels),
+            column_scols=[scol_for(sdf, col) for col in data_columns],
+            column_label_names=column_label_names,
+        )
+        kdf = DataFrame(internal)
+
+        if should_returns_series:
+            return _col(kdf)
+        else:
+            return kdf
+
     def unstack(self):
         """
         Pivot the (necessarily hierarchical) index labels.
@@ -7984,6 +8164,7 @@ def unstack(self):
         See Also
         --------
         DataFrame.pivot : Pivot a table based on column values.
+        DataFrame.stack : Pivot a level of the column labels (inverse operation from unstack).
 
         Examples
         --------
diff --git a/databricks/koalas/missing/frame.py b/databricks/koalas/missing/frame.py
--- a/databricks/koalas/missing/frame.py
+++ b/databricks/koalas/missing/frame.py
@@ -69,7 +69,6 @@ class _MissingPandasLikeDataFrame(object):
     set_axis = unsupported_function("set_axis")
     slice_shift = unsupported_function("slice_shift")
     squeeze = unsupported_function("squeeze")
-    stack = unsupported_function("stack")
     swapaxes = unsupported_function("swapaxes")
     swaplevel = unsupported_function("swaplevel")
     tail = unsupported_function("tail")


Test patch:
diff --git a/databricks/koalas/tests/test_dataframe.py b/databricks/koalas/tests/test_dataframe.py
--- a/databricks/koalas/tests/test_dataframe.py
+++ b/databricks/koalas/tests/test_dataframe.py
@@ -18,6 +18,7 @@
 from distutils.version import LooseVersion
 import inspect
 import sys
+import unittest
 
 import numpy as np
 import pandas as pd
@@ -1852,6 +1853,55 @@ def test_pivot_table_and_index(self):
         self.assert_eq(ktable.index, ptable.index)
         self.assert_eq(repr(ktable.index), repr(ptable.index))
 
+    @unittest.skipIf(
+        LooseVersion(pyspark.__version__) < LooseVersion("2.4"),
+        "stack won't work property with PySpark<2.4",
+    )
+    def test_stack(self):
+        pdf_single_level_cols = pd.DataFrame(
+            [[0, 1], [2, 3]], index=["cat", "dog"], columns=["weight", "height"]
+        )
+        kdf_single_level_cols = ks.from_pandas(pdf_single_level_cols)
+
+        self.assert_eq(
+            kdf_single_level_cols.stack().sort_index(), pdf_single_level_cols.stack().sort_index()
+        )
+
+        multicol1 = pd.MultiIndex.from_tuples(
+            [("weight", "kg"), ("weight", "pounds")], names=["x", "y"]
+        )
+        pdf_multi_level_cols1 = pd.DataFrame(
+            [[1, 2], [2, 4]], index=["cat", "dog"], columns=multicol1
+        )
+        kdf_multi_level_cols1 = ks.from_pandas(pdf_multi_level_cols1)
+
+        self.assert_eq(
+            kdf_multi_level_cols1.stack().sort_index(), pdf_multi_level_cols1.stack().sort_index()
+        )
+
+        multicol2 = pd.MultiIndex.from_tuples([("weight", "kg"), ("height", "m")])
+        pdf_multi_level_cols2 = pd.DataFrame(
+            [[1.0, 2.0], [3.0, 4.0]], index=["cat", "dog"], columns=multicol2
+        )
+        kdf_multi_level_cols2 = ks.from_pandas(pdf_multi_level_cols2)
+
+        self.assert_eq(
+            kdf_multi_level_cols2.stack().sort_index(), pdf_multi_level_cols2.stack().sort_index()
+        )
+
+        pdf = pd.DataFrame(
+            {
+                ("y", "c"): [True, True],
+                ("x", "b"): [False, False],
+                ("x", "c"): [True, False],
+                ("y", "a"): [False, True],
+            }
+        )
+        kdf = ks.from_pandas(pdf)
+
+        self.assert_eq(kdf.stack().sort_index(), pdf.stack().sort_index(), almost=True)
+        self.assert_eq(kdf[[]].stack().sort_index(), pdf[[]].stack().sort_index(), almost=True)
+
     def test_unstack(self):
         pdf = pd.DataFrame(
             np.random.randn(3, 3),