Repository: modin-project/modin

Issue:
BUG: `.broadcast_apply_full_axis` precomputes new partition shapes incorrectly in certain cases
### Modin version checks

- [X] I have checked that this issue has not already been reported.

- [X] I have confirmed this bug exists on the latest released version of Modin.

- [X] I have confirmed this bug exists on the main branch of Modin. (In order to do this you can follow [this guide](https://modin.readthedocs.io/en/stable/getting_started/installation.html#installing-from-the-github-master-branch).)


### Reproducible Example

```python
import pandas
from modin.test.storage_formats.pandas.test_internals import (
    construct_modin_df_by_scheme,
)

import modin.config as cfg

# Has to be equal to the number of column splits
cfg.NPartitions.put(2)

md_df = construct_modin_df_by_scheme(
    pandas_df=pandas.DataFrame(
        {"a": [1, 1, 2, 2], "b": [3, 4, 5, 6], "c": [1, 2, 3, 4], "d": [4, 5, 6, 7]}
    ),
    partitioning_scheme={"row_lengths": [4], "column_widths": [2, 2]},
)

mf = md_df._query_compiler._modin_frame

print(mf._partitions.shape)  # (1, 2)
print(mf.column_widths)  # [2, 2]

res = mf.apply_full_axis(
    axis=1,
    func=lambda df: df,
    keep_partitioning=False,
    new_index=[0, 1, 2, 3],
    new_columns=["a", "b", "c", "d"],
)
print(res._partitions.shape)  # (1, 1)
print(res.column_widths)  # [4]
print(
    res._partitions[0][0].to_pandas().columns
)  # pandas.Index(["a", "b"]); the other two columns are missing!!!
```


### Issue Description

This logic of precomputing new partition lengths seems to be not always right:
https://github.com/modin-project/modin/blob/0df5fbf7a66663110878633467e4488ad37b0337/modin/core/dataframe/pandas/dataframe/dataframe.py#L2719-L2723

Here's a logic inside of `AxisPartition` that may preserve old partitioning even if we explicitly passed `keep_partitioning=False`:
https://github.com/modin-project/modin/blob/0df5fbf7a66663110878633467e4488ad37b0337/modin/core/dataframe/pandas/partitioning/axis_partition.py#L196-L203
If we hit the case then we result in an inconsistency of `new_column_widths` and the actual partitioning. In the reproducer above, the result of `.apply_full_axis` has `column_widths == [4]` however actual widths were unchanged (it's still `[2, 2]`). This causes the second column partition to be filtered out as judging by the `columns_widths == [4]` it's considered to be empty.

There's a parameter for `AxisPartition.apply` called `maintain_partitioning` that controls this behavior, if it's `False` then the old partitioning is never preserved, the only problem is that it's always `True` besides of `keep_partitioning` parameter. 
https://github.com/modin-project/modin/blob/0df5fbf7a66663110878633467e4488ad37b0337/modin/core/dataframe/pandas/partitioning/axis_partition.py#L53

It seems that all we have to do here is just forward `keep_partitioning` parameter to `AxisPartition.apply` as `maintain_partitioning=keep_partitioning`.


### Expected Behavior

...

### Error Logs

<details>

```python-traceback

Replace this line with the error backtrace (if applicable).

```

</details>


### Installed Versions

<details>

Replace this line with the output of pd.show_versions()

</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
@@ -1103,14 +1103,19 @@ def _reorder_labels(self, row_positions=None, col_positions=None):
         """
         new_dtypes = self._dtypes
         if row_positions is not None:
+            # We want to preserve the frame's partitioning so passing in ``keep_partitioning=True``
+            # in order to use the cached `row_lengths` values for the new frame.
+            # If the frame's is re-partitioned using the "standard" partitioning,
+            # then knowing that, we can compute new row lengths.
             ordered_rows = self._partition_mgr_cls.map_axis_partitions(
-                0, self._partitions, lambda df: df.iloc[row_positions]
+                0,
+                self._partitions,
+                lambda df: df.iloc[row_positions],
+                keep_partitioning=True,
             )
             row_idx = self.index[row_positions]
 
-            if self._partitions.shape[0] != ordered_rows.shape[0] or len(
-                row_idx
-            ) != len(self.index):
+            if 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.
@@ -1126,16 +1131,21 @@ def _reorder_labels(self, row_positions=None, col_positions=None):
             row_idx = self.index
             new_lengths = self._row_lengths_cache
         if col_positions is not None:
+            # We want to preserve the frame's partitioning so passing in ``keep_partitioning=True``
+            # in order to use the cached `column_widths` values for the new frame.
+            # If the frame's is re-partitioned using the "standard" partitioning,
+            # then knowing that, we can compute new column widths.
             ordered_cols = self._partition_mgr_cls.map_axis_partitions(
-                1, ordered_rows, lambda df: df.iloc[:, col_positions]
+                1,
+                ordered_rows,
+                lambda df: df.iloc[:, col_positions],
+                keep_partitioning=True,
             )
             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):
+            if 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.
diff --git a/modin/core/dataframe/pandas/partitioning/partition_manager.py b/modin/core/dataframe/pandas/partitioning/partition_manager.py
--- a/modin/core/dataframe/pandas/partitioning/partition_manager.py
+++ b/modin/core/dataframe/pandas/partitioning/partition_manager.py
@@ -450,6 +450,7 @@ def broadcast_axis_partitions(
         kw = {
             "num_splits": num_splits,
             "other_axis_partition": right_partitions,
+            "maintain_partitioning": keep_partitioning,
         }
         if lengths:
             kw["lengths"] = lengths


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
@@ -605,3 +605,42 @@ def test_reorder_labels_dtypes():
         row_positions=None, col_positions=np.arange(len(md_df.columns) - 1, -1, -1)
     )
     df_equals(result.dtypes, result.to_pandas().dtypes)
+
+
+@pytest.mark.parametrize("set_num_partitions", [2], indirect=True)
+def test_repartitioning(set_num_partitions):
+    """
+    This test verifies that 'keep_partitioning=False' doesn't actually preserve partitioning.
+
+    For more details see: https://github.com/modin-project/modin/issues/5621
+    """
+    assert NPartitions.get() == 2
+
+    pandas_df = pandas.DataFrame(
+        {"a": [1, 1, 2, 2], "b": [3, 4, 5, 6], "c": [1, 2, 3, 4], "d": [4, 5, 6, 7]}
+    )
+
+    modin_df = construct_modin_df_by_scheme(
+        pandas_df=pandas.DataFrame(
+            {"a": [1, 1, 2, 2], "b": [3, 4, 5, 6], "c": [1, 2, 3, 4], "d": [4, 5, 6, 7]}
+        ),
+        partitioning_scheme={"row_lengths": [4], "column_widths": [2, 2]},
+    )
+
+    modin_frame = modin_df._query_compiler._modin_frame
+
+    assert modin_frame._partitions.shape == (1, 2)
+    assert modin_frame.column_widths == [2, 2]
+
+    res = modin_frame.apply_full_axis(
+        axis=1,
+        func=lambda df: df,
+        keep_partitioning=False,
+        new_index=[0, 1, 2, 3],
+        new_columns=["a", "b", "c", "d"],
+    )
+
+    assert res._partitions.shape == (1, 1)
+    assert res.column_widths == [4]
+    df_equals(res._partitions[0, 0].to_pandas(), pandas_df)
+    df_equals(res.to_pandas(), pandas_df)