Repository: modin-project/modin (python)

Issue:
Pass sort parameter in stack to `query_compiler` from modin/pandas/dataframe.py

Pull request:
FIX-#7327: Use sort parameter of DataFrame.stack
<!--
Thank you for your contribution!
Please review the contributing docs: https://modin.readthedocs.io/en/latest/development/contributing.html
if you have questions about contributing.
-->

## What do these changes do?

The `sort` parameter of `DataFrame.stack` was exposed in pandas docs in 2.1.0, but we didn't link it from the frontend to the query compiler.

- [x] first commit message and PR title follow format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
  > **_NOTE:_**  If you edit the PR title to match this format, you need to add another commit (even if it's empty) or amend your last commit for the CI job that checks the PR title to pick up the new PR title.
- [x] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] passes `black --check modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [x] Resolves #7327 <!-- issue must be created for each patch -->
- [x] tests added and passing
- [x] module layout described at `docs/development/architecture.rst` is up-to-date <!-- if you have added, renamed or removed files or directories please update the documentation accordingly -->


Patch:
diff --git a/modin/core/dataframe/pandas/partitioning/partition_manager.py b/modin/core/dataframe/pandas/partitioning/partition_manager.py
index cb207f64..6cf22b5f 100644
--- a/modin/core/dataframe/pandas/partitioning/partition_manager.py
+++ b/modin/core/dataframe/pandas/partitioning/partition_manager.py
@@ -915,18 +915,22 @@ class PandasDataframePartitionManager(
         # step cannot be less than 1
         step = max(partitions.shape[0] // column_splits, 1)
         preprocessed_map_func = cls.preprocess_func(map_func)
-        kw = {
-            "num_splits": step,
-        }
         result = np.empty(partitions.shape, dtype=object)
         for i in range(
             0,
             partitions.shape[0],
             step,
         ):
-            joined_column_partitions = cls.column_partitions(partitions[i : i + step])
+            partitions_subset = partitions[i : i + step]
+            # This is necessary when ``partitions.shape[0]`` is not divisible
+            # by `column_splits` without a remainder.
+            actual_step = len(partitions_subset)
+            kw = {
+                "num_splits": actual_step,
+            }
+            joined_column_partitions = cls.column_partitions(partitions_subset)
             for j in range(partitions.shape[1]):
-                result[i : i + step, j] = joined_column_partitions[j].apply(
+                result[i : i + actual_step, j] = joined_column_partitions[j].apply(
                     preprocessed_map_func,
                     *map_func_args if map_func_args is not None else (),
                     **kw,
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
index 965fb98e..dc7f2fce 100644
--- a/modin/core/storage_formats/base/query_compiler.py
+++ b/modin/core/storage_formats/base/query_compiler.py
@@ -22,7 +22,7 @@ from __future__ import annotations
 import abc
 import warnings
 from functools import cached_property
-from typing import TYPE_CHECKING, Hashable, List, Optional
+from typing import TYPE_CHECKING, Hashable, List, Literal, Optional
 
 import numpy as np
 import pandas
@@ -1956,7 +1956,7 @@ class BaseQueryCompiler(
     # 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 @@ class BaseQueryCompiler(
         ----------
         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
@@ -4270,6 +4274,24 @@ class BaseQueryCompiler(
         """
         return self.index if axis == 0 else self.columns
 
+    def get_axis_len(self, axis: Literal[0, 1]) -> int:
+        """
+        Return the length of the specified axis.
+
+        A query compiler may choose to override this method if it has a more efficient way
+        of computing the length of an axis without materializing it.
+
+        Parameters
+        ----------
+        axis : {0, 1}
+            Axis to return labels on.
+
+        Returns
+        -------
+        int
+        """
+        return len(self.get_axis(axis))
+
     def take_2d_labels(
         self,
         index,
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
index 655d427d..b62f2aa4 100644
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -25,7 +25,7 @@ import hashlib
 import re
 import warnings
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Hashable, List, Optional
+from typing import TYPE_CHECKING, Hashable, List, Literal, Optional
 
 import numpy as np
 import pandas
@@ -395,6 +395,24 @@ class PandasQueryCompiler(BaseQueryCompiler, QueryCompilerCaster):
     index: pandas.Index = property(_get_axis(0), _set_axis(0))
     columns: pandas.Index = property(_get_axis(1), _set_axis(1))
 
+    def get_axis_len(self, axis: Literal[0, 1]) -> int:
+        """
+        Return the length of the specified axis.
+
+        Parameters
+        ----------
+        axis : {0, 1}
+            Axis to return labels on.
+
+        Returns
+        -------
+        int
+        """
+        if axis == 0:
+            return len(self._modin_frame)
+        else:
+            return sum(self._modin_frame.column_widths)
+
     @property
     def dtypes(self) -> pandas.Series:
         return self._modin_frame.dtypes
@@ -1944,7 +1962,7 @@ class PandasQueryCompiler(BaseQueryCompiler, QueryCompilerCaster):
             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)
@@ -1956,7 +1974,9 @@ class PandasQueryCompiler(BaseQueryCompiler, QueryCompilerCaster):
 
         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/base.py b/modin/pandas/base.py
index 1b893743..298a376f 100644
--- a/modin/pandas/base.py
+++ b/modin/pandas/base.py
@@ -288,7 +288,9 @@ class BasePandasDataset(ClassLogger):
             A pandas dataset with `num_rows` or fewer rows and `num_cols` or fewer columns.
         """
         # Fast track for empty dataframe.
-        if len(self.index) == 0 or (self._is_dataframe and len(self.columns) == 0):
+        if len(self) == 0 or (
+            self._is_dataframe and self._query_compiler.get_axis_len(1) == 0
+        ):
             return pandas.DataFrame(
                 index=self.index,
                 columns=self.columns if self._is_dataframe else None,
@@ -1004,7 +1006,7 @@ class BasePandasDataset(ClassLogger):
                 return result._query_compiler
             return result
         elif isinstance(func, dict):
-            if len(self.columns) != len(set(self.columns)):
+            if self._query_compiler.get_axis_len(1) != len(set(self.columns)):
                 warnings.warn(
                     "duplicate column names not supported with apply().",
                     FutureWarning,
@@ -2860,7 +2862,7 @@ class BasePandasDataset(ClassLogger):
             axis_length = len(axis_labels)
         else:
             # Getting rows requires indices instead of labels. RangeIndex provides this.
-            axis_labels = pandas.RangeIndex(len(self.index))
+            axis_labels = pandas.RangeIndex(len(self))
             axis_length = len(axis_labels)
         if weights is not None:
             # Index of the weights Series should correspond to the index of the
@@ -3217,7 +3219,7 @@ class BasePandasDataset(ClassLogger):
         """
         if n != 0:
             return self.iloc[-n:]
-        return self.iloc[len(self.index) :]
+        return self.iloc[len(self) :]
 
     def take(self, indices, axis=0, **kwargs) -> Self:  # noqa: PR01, RT01, D200
         """
@@ -4149,7 +4151,7 @@ class BasePandasDataset(ClassLogger):
         -------
         int
         """
-        return len(self.index)
+        return self._query_compiler.get_axis_len(0)
 
     @_doc_binary_op(
         operation="less than comparison",
diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py
index de96ea0a..4f47c937 100644
--- a/modin/pandas/dataframe.py
+++ b/modin/pandas/dataframe.py
@@ -268,13 +268,15 @@ class DataFrame(BasePandasDataset):
         -------
         str
         """
-        num_rows = pandas.get_option("display.max_rows") or len(self.index)
-        num_cols = pandas.get_option("display.max_columns") or len(self.columns)
+        num_rows = pandas.get_option("display.max_rows") or len(self)
+        num_cols = pandas.get_option(
+            "display.max_columns"
+        ) or self._query_compiler.get_axis_len(1)
         result = repr(self._build_repr_df(num_rows, num_cols))
-        if len(self.index) > num_rows or len(self.columns) > num_cols:
+        if len(self) > num_rows or self._query_compiler.get_axis_len(1) > num_cols:
             # The split here is so that we don't repr pandas row lengths.
             return result.rsplit("\n\n", 1)[0] + "\n\n[{0} rows x {1} columns]".format(
-                len(self.index), len(self.columns)
+                *self.shape
             )
         else:
             return result
@@ -293,13 +295,11 @@ class DataFrame(BasePandasDataset):
         # We use pandas _repr_html_ to get a string of the HTML representation
         # of the dataframe.
         result = self._build_repr_df(num_rows, num_cols)._repr_html_()
-        if len(self.index) > num_rows or len(self.columns) > num_cols:
+        if len(self) > num_rows or self._query_compiler.get_axis_len(1) > num_cols:
             # We split so that we insert our correct dataframe dimensions.
             return result.split("<p>")[
                 0
-            ] + "<p>{0} rows x {1} columns</p>\n</div>".format(
-                len(self.index), len(self.columns)
-            )
+            ] + "<p>{0} rows x {1} columns</p>\n</div>".format(*self.shape)
         else:
             return result
 
@@ -365,7 +365,7 @@ class DataFrame(BasePandasDataset):
         """
         Indicate whether ``DataFrame`` is empty.
         """
-        return len(self.columns) == 0 or len(self.index) == 0
+        return self._query_compiler.get_axis_len(1) == 0 or len(self) == 0
 
     @property
     def axes(self) -> list[pandas.Index]:  # noqa: RT01, D200
@@ -379,7 +379,7 @@ class DataFrame(BasePandasDataset):
         """
         Return a tuple representing the dimensionality of the ``DataFrame``.
         """
-        return len(self.index), len(self.columns)
+        return len(self), self._query_compiler.get_axis_len(1)
 
     def add_prefix(self, prefix, axis=None) -> DataFrame:  # noqa: PR01, RT01, D200
         """
@@ -781,7 +781,9 @@ class DataFrame(BasePandasDataset):
         """
         if isinstance(other, BasePandasDataset):
             common = self.columns.union(other.index)
-            if len(common) > len(self.columns) or len(common) > len(other.index):
+            if len(common) > self._query_compiler.get_axis_len(1) or len(common) > len(
+                other
+            ):
                 raise ValueError("Matrices are not aligned")
 
             qc = other.reindex(index=common)._query_compiler
@@ -1084,7 +1086,7 @@ class DataFrame(BasePandasDataset):
                     + f"{len(value.columns)} columns instead."
                 )
             value = value.squeeze(axis=1)
-        if not self._query_compiler.lazy_row_count and len(self.index) == 0:
+        if not self._query_compiler.lazy_row_count and len(self) == 0:
             if not hasattr(value, "index"):
                 try:
                     value = pandas.Series(value)
@@ -1099,7 +1101,7 @@ class DataFrame(BasePandasDataset):
             new_query_compiler = self.__constructor__(
                 value, index=new_index, columns=new_columns
             )._query_compiler
-        elif len(self.columns) == 0 and loc == 0:
+        elif self._query_compiler.get_axis_len(1) == 0 and loc == 0:
             new_index = self.index
             new_query_compiler = self.__constructor__(
                 data=value,
@@ -1110,18 +1112,19 @@ class DataFrame(BasePandasDataset):
             if (
                 is_list_like(value)
                 and not isinstance(value, (pandas.Series, Series))
-                and len(value) != len(self.index)
+                and len(value) != len(self)
             ):
                 raise ValueError(
                     "Length of values ({}) does not match length of index ({})".format(
-                        len(value), len(self.index)
+                        len(value), len(self)
                     )
                 )
             if allow_duplicates is not True and column in self.columns:
                 raise ValueError(f"cannot insert {column}, already exists")
-            if not -len(self.columns) <= loc <= len(self.columns):
+            columns_len = self._query_compiler.get_axis_len(1)
+            if not -columns_len <= loc <= columns_len:
                 raise IndexError(
-                    f"index {loc} is out of bounds for axis 0 with size {len(self.columns)}"
+                    f"index {loc} is out of bounds for axis 0 with size {columns_len}"
                 )
             elif loc < 0:
                 raise ValueError("unbounded slice")
@@ -2074,12 +2077,14 @@ class DataFrame(BasePandasDataset):
         Squeeze 1 dimensional axis objects into scalars.
         """
         axis = self._get_axis_number(axis) if axis is not None else None
-        if axis is None and (len(self.columns) == 1 or len(self.index) == 1):
+        if axis is None and (
+            self._query_compiler.get_axis_len(1) == 1 or len(self) == 1
+        ):
             return Series(query_compiler=self._query_compiler).squeeze()
-        if axis == 1 and len(self.columns) == 1:
+        if axis == 1 and self._query_compiler.get_axis_len(1) == 1:
             self._query_compiler._shape_hint = "column"
             return Series(query_compiler=self._query_compiler)
-        if axis == 0 and len(self.index) == 1:
+        if axis == 0 and len(self) == 1:
             qc = self.T._query_compiler
             qc._shape_hint = "column"
             return Series(query_compiler=qc)
@@ -2113,11 +2118,11 @@ class DataFrame(BasePandasDataset):
             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(
@@ -2671,7 +2676,7 @@ class DataFrame(BasePandasDataset):
             return self._setitem_slice(key, value)
 
         if hashable(key) and key not in self.columns:
-            if isinstance(value, Series) and len(self.columns) == 0:
+            if isinstance(value, Series) and self._query_compiler.get_axis_len(1) == 0:
                 # Note: column information is lost when assigning a query compiler
                 prev_index = self.columns
                 self._query_compiler = value._query_compiler.copy()
@@ -2680,7 +2685,9 @@ class DataFrame(BasePandasDataset):
                 self.columns = prev_index.insert(0, key)
                 return
             # Do new column assignment after error checks and possible value modifications
-            self.insert(loc=len(self.columns), column=key, value=value)
+            self.insert(
+                loc=self._query_compiler.get_axis_len(1), column=key, value=value
+            )
             return
 
         if not hashable(key):
@@ -2756,7 +2763,7 @@ class DataFrame(BasePandasDataset):
 
                 new_qc = self._query_compiler.insert_item(
                     axis=1,
-                    loc=len(self.columns),
+                    loc=self._query_compiler.get_axis_len(1),
                     value=value._query_compiler,
                     how="left",
                 )
@@ -2783,7 +2790,7 @@ class DataFrame(BasePandasDataset):
             if not isinstance(value, (Series, Categorical, np.ndarray, list, range)):
                 value = list(value)
 
-        if not self._query_compiler.lazy_row_count and len(self.index) == 0:
+        if not self._query_compiler.lazy_row_count and len(self) == 0:
             new_self = self.__constructor__({key: value}, columns=self.columns)
             self._update_inplace(new_self._query_compiler)
         else:
diff --git a/modin/pandas/series.py b/modin/pandas/series.py
index 20825343..11188e85 100644
--- a/modin/pandas/series.py
+++ b/modin/pandas/series.py
@@ -445,8 +445,8 @@ class Series(BasePandasDataset):
             name_str = "Name: {}, ".format(str(self.name))
         else:
             name_str = ""
-        if len(self.index) > num_rows:
-            len_str = "Length: {}, ".format(len(self.index))
+        if len(self) > num_rows:
+            len_str = "Length: {}, ".format(len(self))
         else:
             len_str = ""
         dtype_str = "dtype: {}".format(
@@ -966,7 +966,7 @@ class Series(BasePandasDataset):
         """
         if isinstance(other, BasePandasDataset):
             common = self.index.union(other.index)
-            if len(common) > len(self.index) or len(common) > len(other.index):
+            if len(common) > len(self) or len(common) > len(other):
                 raise ValueError("Matrices are not aligned")
 
             qc = other.reindex(index=common)._query_compiler
@@ -1761,7 +1761,7 @@ class Series(BasePandasDataset):
             name = 0 if self.name is None else self.name
 
         if drop and level is None:
-            new_idx = pandas.RangeIndex(len(self.index))
+            new_idx = pandas.RangeIndex(len(self))
             if inplace:
                 self.index = new_idx
             else:
@@ -1989,7 +1989,7 @@ class Series(BasePandasDataset):
         if axis is not None:
             # Validate `axis`
             pandas.Series._get_axis_number(axis)
-        if len(self.index) == 1:
+        if len(self) == 1:
             return self._reduce_dimension(self._query_compiler)
         else:
             return self.copy()
@@ -2307,7 +2307,7 @@ class Series(BasePandasDataset):
         """
         Indicate whether Series is empty.
         """
-        return len(self.index) == 0
+        return len(self) == 0
 
     @property
     def hasnans(self) -> bool:  # noqa: RT01, D200
@@ -2648,7 +2648,7 @@ class Series(BasePandasDataset):
         if is_bool_indexer(key):
             return self.__constructor__(
                 query_compiler=self._query_compiler.getitem_row_array(
-                    pandas.RangeIndex(len(self.index))[key]
+                    pandas.RangeIndex(len(self))[key]
                 )
             )
         # TODO: More efficiently handle `tuple` case for `Series.__getitem__`


Test patch:
diff --git a/modin/tests/core/storage_formats/pandas/test_internals.py b/modin/tests/core/storage_formats/pandas/test_internals.py
index b030fe7b..5d3e99c0 100644
--- a/modin/tests/core/storage_formats/pandas/test_internals.py
+++ b/modin/tests/core/storage_formats/pandas/test_internals.py
@@ -2677,8 +2677,9 @@ def test_dynamic_partitioning(partitioning_scheme, expected_map_approach):
             expected_method.assert_called()
 
 
-def test_map_partitions_joined_by_column():
-    with context(NPartitions=CpuCount.get() * 2):
+@pytest.mark.parametrize("npartitions", [7, CpuCount.get() * 2])
+def test_map_partitions_joined_by_column(npartitions):
+    with context(NPartitions=npartitions):
         ncols = MinColumnPartitionSize.get()
         nrows = MinRowPartitionSize.get() * CpuCount.get() * 2
         data = {f"col{i}": np.ones(nrows) for i in range(ncols)}
diff --git a/modin/tests/pandas/dataframe/test_default.py b/modin/tests/pandas/dataframe/test_default.py
index 71f49924..d40a30e3 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])