Repository: modin-project/modin

Issue:
[PERF] Precompute row/column lengths in `._reorder_labels`
Currently, the resulting frame of [`._reorder_labels()`](https://github.com/modin-project/modin/blob/3cc33a2be050a2472fd084ccc3b9bbb33986f770/modin/core/dataframe/pandas/dataframe/dataframe.py#L1054) contain neither row nor column lengths cache:
https://github.com/modin-project/modin/blob/3cc33a2be050a2472fd084ccc3b9bbb33986f770/modin/core/dataframe/pandas/dataframe/dataframe.py#L1088
There's no reason though why we can't precompute or at least propagate this cache from the source frame.

The reordering happens via [`.map_axis_partitions(..., lengths=None, keep_partitioning=False)`](https://github.com/modin-project/modin/blob/3cc33a2be050a2472fd084ccc3b9bbb33986f770/modin/core/dataframe/pandas/partitioning/partition_manager.py#L530) method. The given set of parameters results in either:
1. Repartitioning the axis using the default partitioning scheme if the [current partitioning doesn't satisfy the desired* one](https://github.com/modin-project/modin/blob/3cc33a2be050a2472fd084ccc3b9bbb33986f770/modin/core/dataframe/pandas/partitioning/axis_partition.py#L193-L194) or if [the reordered result has changed its shape](https://github.com/modin-project/modin/blob/3cc33a2be050a2472fd084ccc3b9bbb33986f770/modin/core/dataframe/pandas/partitioning/axis_partition.py#L197-L199).
2. [Preserving the partitioning](https://github.com/modin-project/modin/blob/3cc33a2be050a2472fd084ccc3b9bbb33986f770/modin/core/dataframe/pandas/partitioning/axis_partition.py#L197) if the current partitioning is already desirable* and the result's shape hasn't been changed.

For the first case we can replicate the logic of [`split_result_of_axis_func_pandas`](https://github.com/modin-project/modin/blob/3cc33a2be050a2472fd084ccc3b9bbb33986f770/modin/core/storage_formats/pandas/utils.py#L68) to compute new row/column lengths.

For the second case, we can just preserve the existing cache as the partitioning scheme stays the same.

<details><summary>* desired partitioning</summary>

That means that the number of splits along the axis equals the `NPartitions.get()`.

</details>



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
@@ -28,6 +28,7 @@
 from typing import List, Hashable, Optional, Callable, Union, Dict, TYPE_CHECKING
 
 from modin.core.storage_formats.pandas.query_compiler import PandasQueryCompiler
+from modin.core.storage_formats.pandas.utils import get_length_list
 from modin.error_message import ErrorMessage
 from modin.core.storage_formats.pandas.parsers import (
     find_common_type_cat as find_common_type,
@@ -1083,23 +1084,59 @@ def _reorder_labels(self, row_positions=None, col_positions=None):
         PandasDataframe
             A new PandasDataframe with reordered columns and/or rows.
         """
+        new_lengths, new_widths = None, None
+        new_dtypes = self._dtypes
         if row_positions is not None:
             ordered_rows = self._partition_mgr_cls.map_axis_partitions(
                 0, self._partitions, lambda df: df.iloc[row_positions]
             )
             row_idx = self.index[row_positions]
+
+            if self._partitions.shape[0] != ordered_rows.shape[0] or len(
+                row_idx
+            ) != len(self.index):
+                # The frame was re-partitioned along the 0 axis during reordering using
+                # the "standard" partitioning. Knowing the standard partitioning scheme
+                # we are able to compute new row lengths.
+                new_lengths = get_length_list(
+                    axis_len=len(row_idx), num_splits=ordered_rows.shape[0]
+                )
+            else:
+                # If the frame's partitioning was preserved then
+                # we can use previous row lengths cache
+                new_lengths = self._row_lengths_cache
         else:
             ordered_rows = self._partitions
             row_idx = self.index
+            new_lengths = self._row_lengths_cache
         if col_positions is not None:
             ordered_cols = self._partition_mgr_cls.map_axis_partitions(
                 1, ordered_rows, lambda df: df.iloc[:, col_positions]
             )
             col_idx = self.columns[col_positions]
+            if new_dtypes is not None:
+                new_dtypes = self._dtypes.iloc[col_positions]
+
+            if self._partitions.shape[1] != ordered_cols.shape[1] or len(
+                col_idx
+            ) != len(self.columns):
+                # The frame was re-partitioned along the 1 axis during reordering using
+                # the "standard" partitioning. Knowing the standard partitioning scheme
+                # we are able to compute new column widths.
+                new_widths = get_length_list(
+                    axis_len=len(col_idx), num_splits=ordered_cols.shape[1]
+                )
+            else:
+                # If the frame's partitioning was preserved then
+                # we can use previous column widths cache
+                new_widths = self._column_widths_cache
         else:
             ordered_cols = ordered_rows
             col_idx = self.columns
-        return self.__constructor__(ordered_cols, row_idx, col_idx)
+            new_widths = self._column_widths_cache
+        return self.__constructor__(
+            ordered_cols, row_idx, col_idx, new_lengths, new_widths, new_dtypes
+        )
 
     @lazy_metadata_decorator(apply_axis=None)
     def copy(self):
diff --git a/modin/core/storage_formats/pandas/utils.py b/modin/core/storage_formats/pandas/utils.py
--- a/modin/core/storage_formats/pandas/utils.py
+++ b/modin/core/storage_formats/pandas/utils.py
@@ -90,12 +90,11 @@ def split_result_of_axis_func_pandas(axis, num_splits, result, length_list=None)
     if num_splits == 1:
         return [result]
 
-    if length_list is not None:
-        length_list.insert(0, 0)
-    else:
-        chunksize = compute_chunksize(result.shape[axis], num_splits)
-        length_list = np.full(num_splits + 1, chunksize)
-        length_list[0] = 0
+    if length_list is None:
+        length_list = get_length_list(result.shape[axis], num_splits)
+    # Inserting the first "zero" to properly compute cumsum indexing slices
+    length_list = np.insert(length_list, obj=0, values=[0])
+
     sums = np.cumsum(length_list)
     # We do this to restore block partitioning
     if axis == 0 or isinstance(result, pandas.Series):
@@ -104,6 +103,31 @@ def split_result_of_axis_func_pandas(axis, num_splits, result, length_list=None)
         return [result.iloc[:, sums[i] : sums[i + 1]] for i in range(len(sums) - 1)]
 
 
+def get_length_list(axis_len: int, num_splits: int) -> list:
+    """
+    Compute partitions lengths along the axis with the specified number of splits.
+
+    Parameters
+    ----------
+    axis_len : int
+        Element count in an axis.
+    num_splits : int
+        Number of splits along the axis.
+
+    Returns
+    -------
+    list of ints
+        List of integer lengths of partitions.
+    """
+    chunksize = compute_chunksize(axis_len, num_splits)
+    return [
+        chunksize
+        if (i + 1) * chunksize <= axis_len
+        else max(0, axis_len - i * chunksize)
+        for i in range(num_splits)
+    ]
+
+
 def length_fn_pandas(df):
     """
     Compute number of rows of passed `pandas.DataFrame`.


Test patch:
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
@@ -18,6 +18,8 @@
     df_equals,
 )
 from modin.config import NPartitions, Engine
+from modin.distributed.dataframe.pandas import from_partitions
+from modin.core.storage_formats.pandas.utils import split_result_of_axis_func_pandas
 
 import numpy as np
 import pandas
@@ -60,6 +62,57 @@ def put(x):
     virtual_row_partition_class = PandasOnDaskDataframeRowPartition
 
 
+def construct_modin_df_by_scheme(pandas_df, partitioning_scheme):
+    """
+    Build ``modin.pandas.DataFrame`` from ``pandas.DataFrame`` according the `partitioning_scheme`.
+
+    Parameters
+    ----------
+    pandas_df : pandas.DataFrame
+    partitioning_scheme : dict[{"row_lengths", "column_widths"}] -> list of ints
+
+    Returns
+    -------
+    modin.pandas.DataFrame
+    """
+    row_partitions = split_result_of_axis_func_pandas(
+        axis=0,
+        num_splits=len(partitioning_scheme["row_lengths"]),
+        result=pandas_df,
+        length_list=partitioning_scheme["row_lengths"],
+    )
+    partitions = [
+        split_result_of_axis_func_pandas(
+            axis=1,
+            num_splits=len(partitioning_scheme["column_widths"]),
+            result=row_part,
+            length_list=partitioning_scheme["column_widths"],
+        )
+        for row_part in row_partitions
+    ]
+
+    md_df = from_partitions(
+        [[put(part) for part in row_parts] for row_parts in partitions], axis=None
+    )
+    return md_df
+
+
+def validate_partitions_cache(df):
+    """Assert that the ``PandasDataframe`` shape caches correspond to the actual partition's shapes."""
+    row_lengths = df._row_lengths_cache
+    column_widths = df._column_widths_cache
+
+    assert row_lengths is not None
+    assert column_widths is not None
+    assert df._partitions.shape[0] == len(row_lengths)
+    assert df._partitions.shape[1] == len(column_widths)
+
+    for i in range(df._partitions.shape[0]):
+        for j in range(df._partitions.shape[1]):
+            assert df._partitions[i, j].length() == row_lengths[i]
+            assert df._partitions[i, j].width() == column_widths[j]
+
+
 @pytest.fixture
 def set_num_partitions(request):
     old_num_partitions = NPartitions.get()
@@ -433,3 +486,102 @@ def test_virtual_partition_dup_object_ref():
     ), "Test setup did not contain duplicate objects"
     # The below call to wait() should not crash
     partition.wait()
+
+
+__test_reorder_labels_cache_axis_positions = [
+    pytest.param(lambda index: None, id="no_reordering"),
+    pytest.param(lambda index: np.arange(len(index) - 1, -1, -1), id="reordering_only"),
+    pytest.param(
+        lambda index: [0, 1, 2, len(index) - 3, len(index) - 2, len(index) - 1],
+        id="projection_only",
+    ),
+    pytest.param(
+        lambda index: np.repeat(np.arange(len(index)), repeats=3), id="size_grow"
+    ),
+]
+
+
+@pytest.mark.parametrize("row_positions", __test_reorder_labels_cache_axis_positions)
+@pytest.mark.parametrize("col_positions", __test_reorder_labels_cache_axis_positions)
+@pytest.mark.parametrize(
+    "partitioning_scheme",
+    [
+        pytest.param(
+            lambda df: {
+                "row_lengths": [df.shape[0]],
+                "column_widths": [df.shape[1]],
+            },
+            id="single_partition",
+        ),
+        pytest.param(
+            lambda df: {
+                "row_lengths": [32, max(0, df.shape[0] - 32)],
+                "column_widths": [32, max(0, df.shape[1] - 32)],
+            },
+            id="two_unbalanced_partitions",
+        ),
+        pytest.param(
+            lambda df: {
+                "row_lengths": [df.shape[0] // NPartitions.get()] * NPartitions.get(),
+                "column_widths": [df.shape[1] // NPartitions.get()] * NPartitions.get(),
+            },
+            id="perfect_partitioning",
+        ),
+        pytest.param(
+            lambda df: {
+                "row_lengths": [2**i for i in range(NPartitions.get())],
+                "column_widths": [2**i for i in range(NPartitions.get())],
+            },
+            id="unbalanced_partitioning_equals_npartition",
+        ),
+        pytest.param(
+            lambda df: {
+                "row_lengths": [2] * (df.shape[0] // 2),
+                "column_widths": [2] * (df.shape[1] // 2),
+            },
+            id="unbalanced_partitioning",
+        ),
+    ],
+)
+def test_reorder_labels_cache(
+    row_positions,
+    col_positions,
+    partitioning_scheme,
+):
+    pandas_df = pandas.DataFrame(test_data_values[0])
+
+    md_df = construct_modin_df_by_scheme(pandas_df, partitioning_scheme(pandas_df))
+    md_df = md_df._query_compiler._modin_frame
+
+    result = md_df._reorder_labels(
+        row_positions(md_df.index), col_positions(md_df.columns)
+    )
+    validate_partitions_cache(result)
+
+
+def test_reorder_labels_dtypes():
+    pandas_df = pandas.DataFrame(
+        {
+            "a": [1, 2, 3, 4],
+            "b": [1.0, 2.4, 3.4, 4.5],
+            "c": ["a", "b", "c", "d"],
+            "d": pd.to_datetime([1, 2, 3, 4], unit="D"),
+        }
+    )
+
+    md_df = construct_modin_df_by_scheme(
+        pandas_df,
+        partitioning_scheme={
+            "row_lengths": [len(pandas_df)],
+            "column_widths": [
+                len(pandas_df) // 2,
+                len(pandas_df) // 2 + len(pandas_df) % 2,
+            ],
+        },
+    )
+    md_df = md_df._query_compiler._modin_frame
+
+    result = md_df._reorder_labels(
+        row_positions=None, col_positions=np.arange(len(md_df.columns) - 1, -1, -1)
+    )
+    df_equals(result.dtypes, result.to_pandas().dtypes)