Repository: modin-project/modin

Issue:
Defer metadata computation until really needed for the result of `.sort_values`
We currently trigger the computations here: https://github.com/modin-project/modin/blob/94b2a0e86ea631451f44d4af64ab5424e42f7764/modin/core/dataframe/pandas/dataframe/dataframe.py#L2154-L2157

We may want to avoid it.

_Originally posted by @RehanSD in https://github.com/modin-project/modin/pull/5780#discussion_r1137835188_


Discussion:

Resolving patch:
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -451,6 +451,20 @@ def axes(self):
         """
         return [self.index, self.columns]
 
+    def get_axis(self, axis: int = 0) -> pandas.Index:
+        """
+        Get index object for the requested axis.
+
+        Parameters
+        ----------
+        axis : {0, 1}, default: 0
+
+        Returns
+        -------
+        pandas.Index
+        """
+        return self.index if axis == 0 else self.columns
+
     def _compute_axis_labels_and_lengths(self, axis: int, partitions=None):
         """
         Compute the labels for specific `axis`.
@@ -2170,22 +2184,20 @@ def sort_function(df):
             shuffling_functions,
             sort_function,
         )
-        new_axes = self.axes
+        new_axes = (
+            [None, self.copy_columns_cache()]
+            if axis == Axis.ROW_WISE
+            else [self.copy_index_cache(), None]
+        )
         new_lengths = [None, None]
         if kwargs.get("ignore_index", False):
-            new_axes[axis.value] = RangeIndex(len(new_axes[axis.value]))
-        else:
-            (
-                new_axes[axis.value],
-                new_lengths[axis.value],
-            ) = self._compute_axis_labels_and_lengths(axis.value, new_partitions)
+            new_axes[axis.value] = RangeIndex(len(self.get_axis(axis.value)))
 
-        new_axes[axis.value] = new_axes[axis.value].set_names(
-            self.axes[axis.value].names
-        )
         # We perform the final steps of the sort on full axis partitions, so we know that the
         # length of each partition is the full length of the dataframe.
-        new_lengths[axis.value ^ 1] = [len(self.columns)]
+        new_lengths[axis.value ^ 1] = (
+            [len(self.columns)] if self.has_materialized_columns else None
+        )
         # Since the strategy to pick our pivots involves random sampling
         # we could end up picking poor pivots, leading to skew in our partitions.
         # We should add a fix to check if there is skew in the partitions and rebalance
@@ -2193,7 +2205,7 @@ def sort_function(df):
         # resolves the case where there isn't the right amount of partitions - not where
         # there is skew across the lengths of partitions.
         new_modin_frame = self.__constructor__(
-            new_partitions, *new_axes, *new_lengths, self.dtypes
+            new_partitions, *new_axes, *new_lengths, self._dtypes
         )
         if kwargs.get("ignore_index", False):
             new_modin_frame._propagate_index_objs(axis=0)
diff --git a/modin/core/dataframe/pandas/dataframe/utils.py b/modin/core/dataframe/pandas/dataframe/utils.py
--- a/modin/core/dataframe/pandas/dataframe/utils.py
+++ b/modin/core/dataframe/pandas/dataframe/utils.py
@@ -274,7 +274,8 @@ def split_partitions_using_pivots_for_sort(
             groupby_col = len(pivots) - groupby_col
     if len(non_na_rows) == 1:
         groups = [
-            pandas.DataFrame(columns=df.columns).astype(df.dtypes)
+            # taking an empty slice for an index's metadata
+            pandas.DataFrame(index=df.index[:0], columns=df.columns).astype(df.dtypes)
             if i != groupby_col
             else non_na_rows
             for i in range(len(pivots) + 1)
@@ -282,9 +283,11 @@ def split_partitions_using_pivots_for_sort(
     else:
         grouped = non_na_rows.groupby(groupby_col)
         groups = [
-            grouped.get_group(i)
-            if i in grouped.keys
-            else pandas.DataFrame(columns=df.columns).astype(df.dtypes)
+            grouped.get_group(i) if i in grouped.keys
+            # taking an empty slice for an index's metadata
+            else pandas.DataFrame(index=df.index[:0], columns=df.columns).astype(
+                df.dtypes
+            )
             for i in range(len(pivots) + 1)
         ]
     index_to_insert_na_vals = -1 if kwargs.get("na_position", "last") == "last" else 0


Test patch:
diff --git a/modin/pandas/test/dataframe/test_join_sort.py b/modin/pandas/test/dataframe/test_join_sort.py
--- a/modin/pandas/test/dataframe/test_join_sort.py
+++ b/modin/pandas/test/dataframe/test_join_sort.py
@@ -643,6 +643,35 @@ def test_sort_values_descending_with_only_two_bins():
     )
 
 
+@pytest.mark.skipif(
+    StorageFormat.get() == "Hdk",
+    reason="https://github.com/modin-project/modin/issues/3941",
+)
+@pytest.mark.parametrize("ignore_index", [True, False])
+def test_sort_values_preserve_index_names(ignore_index):
+    modin_df, pandas_df = create_test_dfs(
+        np.random.choice(128, 128, replace=False).reshape((128, 1))
+    )
+
+    pandas_df.index.names, pandas_df.columns.names = ["custom_name"], ["custom_name"]
+    modin_df.index.names, modin_df.columns.names = ["custom_name"], ["custom_name"]
+    # workaround for #1618 to actually propagate index change
+    modin_df.index = modin_df.index
+    modin_df.columns = modin_df.columns
+
+    def comparator(df1, df2):
+        assert df1.index.names == df2.index.names
+        assert df1.columns.names == df2.columns.names
+        df_equals(df1, df2)
+
+    eval_general(
+        modin_df,
+        pandas_df,
+        lambda df: df.sort_values(df.columns[0], ignore_index=ignore_index),
+        comparator=comparator,
+    )
+
+
 @pytest.mark.parametrize("ascending", [True, False])
 def test_sort_values_with_one_partition(ascending):
     # Test case from https://github.com/modin-project/modin/issues/5859
diff --git a/modin/test/storage_formats/pandas/test_internals.py b/modin/test/storage_formats/pandas/test_internals.py
--- a/modin/test/storage_formats/pandas/test_internals.py
+++ b/modin/test/storage_formats/pandas/test_internals.py
@@ -903,6 +903,37 @@ def test_shuffle_partitions_with_empty_pivots(ascending):
     assert ref.equals(res)
 
 
+@pytest.mark.parametrize("ascending", [True, False])
+def test_split_partition_preserve_names(ascending):
+    """
+    This test verifies that the dataframes being split by ``split_partitions_using_pivots_for_sort``
+    preserve their index/column names.
+    """
+    from modin.core.dataframe.pandas.dataframe.utils import (
+        split_partitions_using_pivots_for_sort,
+    )
+
+    df = pandas.DataFrame(
+        {
+            "numeric_col": range(9),
+            "non_numeric_col": list("abcdefghi"),
+        }
+    )
+    index_name = "custom_name"
+    df.index.name = index_name
+    df.columns.name = index_name
+
+    # Pivots that contain empty bins
+    pivots = [2, 2, 5, 7]
+    splits = split_partitions_using_pivots_for_sort(
+        df, df, column="numeric_col", pivots=pivots, ascending=ascending
+    )
+
+    for part in splits:
+        assert part.index.name == index_name
+        assert part.columns.name == index_name
+
+
 @pytest.mark.parametrize("has_cols_metadata", [True, False])
 @pytest.mark.parametrize("has_dtypes_metadata", [True, False])
 def test_merge_preserves_metadata(has_cols_metadata, has_dtypes_metadata):