diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py index 76d3902a6e3..dc7f2fce17b 100644 --- a/modin/core/storage_formats/base/query_compiler.py +++ b/modin/core/storage_formats/base/query_compiler.py @@ -1956,7 +1956,7 @@ def searchsorted(self, **kwargs): # noqa: PR02 # END Abstract map partitions operations @doc_utils.add_refer_to("DataFrame.stack") - def stack(self, level, dropna): + def stack(self, level, dropna, sort): """ Stack the prescribed level(s) from columns to index. @@ -1964,13 +1964,17 @@ def stack(self, level, dropna): ---------- level : int or label dropna : bool + sort : bool Returns ------- BaseQueryCompiler """ return DataFrameDefault.register(pandas.DataFrame.stack)( - self, level=level, dropna=dropna + self, + level=level, + dropna=dropna, + sort=sort, ) # Abstract map partitions across select indices diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py index cd0e59ca56f..b62f2aa4474 100644 --- a/modin/core/storage_formats/pandas/query_compiler.py +++ b/modin/core/storage_formats/pandas/query_compiler.py @@ -1962,7 +1962,7 @@ def get_unique_level_values(index): result = result.reindex(0, new_index) return result - def stack(self, level, dropna): + def stack(self, level, dropna, sort): if not isinstance(self.columns, pandas.MultiIndex) or ( isinstance(self.columns, pandas.MultiIndex) and is_list_like(level) @@ -1974,7 +1974,9 @@ def stack(self, level, dropna): new_modin_frame = self._modin_frame.apply_full_axis( 1, - lambda df: pandas.DataFrame(df.stack(level=level, dropna=dropna)), + lambda df: pandas.DataFrame( + df.stack(level=level, dropna=dropna, sort=sort) + ), new_columns=new_columns, ) return self.__constructor__(new_modin_frame) diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py index e90837a3b6d..4f47c9374e9 100644 --- a/modin/pandas/dataframe.py +++ b/modin/pandas/dataframe.py @@ -2118,11 +2118,11 @@ def stack( is_multiindex and is_list_like(level) and len(level) == self.columns.nlevels ): return self._reduce_dimension( - query_compiler=self._query_compiler.stack(level, dropna) + query_compiler=self._query_compiler.stack(level, dropna, sort) ) else: return self.__constructor__( - query_compiler=self._query_compiler.stack(level, dropna) + query_compiler=self._query_compiler.stack(level, dropna, sort) ) def sub( diff --git a/modin/tests/pandas/dataframe/test_default.py b/modin/tests/pandas/dataframe/test_default.py index 71f49924c94..d40a30e38e9 100644 --- a/modin/tests/pandas/dataframe/test_default.py +++ b/modin/tests/pandas/dataframe/test_default.py @@ -1181,6 +1181,16 @@ def test_stack(data, is_multi_idx, is_multi_col): df_equals(modin_df.stack(level=[0, 1, 2]), pandas_df.stack(level=[0, 1, 2])) +@pytest.mark.parametrize("sort", [True, False]) +def test_stack_sort(sort): + # Example frame slightly modified from pandas docs to be unsorted + cols = pd.MultiIndex.from_tuples([("weight", "pounds"), ("weight", "kg")]) + modin_df, pandas_df = create_test_dfs( + [[1, 2], [2, 4]], index=["cat", "dog"], columns=cols + ) + df_equals(modin_df.stack(sort=sort), pandas_df.stack(sort=sort)) + + @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys) @pytest.mark.parametrize("axis1", [0, 1]) @pytest.mark.parametrize("axis2", [0, 1])