Repository: modin-project/modin

Issue:
BUG: applying function to virtual partitions with NPartitions = 1 produces dataframe with 3 dimensional partitions
### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: macOS Monterey
- **Modin version** (`modin.__version__`): 980f126dd941da534da10d6a5e919cf10410cae0
- **Python version**: 3.10.4
- **Code we can use to reproduce**:

```python
import modin.pandas as pd
from modin.config import NPartitions

NPartitions.put(1)
df = pd.DataFrame([1])
df = pd.concat([df] * 100)
result = df.applymap(lambda x: x)
print(result._query_compiler._modin_frame._partitions.shape)
```

### Describe the problem

`df._query_compiler._modin_frame._partitions[0,0].apply(lambda x: x)` returns a list of 100 partitions instead of a single partition. Probably there is something wrong with how the virtual partition class is using the parent axis partition class's `apply`.


Discussion:

Resolving patch:
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
@@ -1422,11 +1422,14 @@ def rebalance_partitions(cls, partitions):
                         for obj in partitions[stop]
                     ]
                     partition_size = ideal_partition_size
+                # The new virtual partitions are not `full_axis`, even if they
+                # happen to span all rows in the dataframe, because they are
+                # meant to be the final partitions of the dataframe. They've
+                # already been split up correctly along axis 0, but using the
+                # default full_axis=True would cause partition.apply() to split
+                # its result along axis 0.
                 new_partitions.append(
-                    cls.column_partitions(
-                        (partitions[start : stop + 1]),
-                        full_axis=partition_size == total_rows,
-                    )
+                    cls.column_partitions(partitions[start : stop + 1], full_axis=False)
                 )
                 start = stop + 1
             return np.array(new_partitions)


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
@@ -59,6 +59,14 @@ def put(x):
     virtual_row_partition_class = PandasOnDaskDataframeRowPartition
 
 
+@pytest.fixture
+def set_num_partitions(request):
+    old_num_partitions = NPartitions.get()
+    NPartitions.put(request.param)
+    yield
+    NPartitions.put(old_num_partitions)
+
+
 def test_aligning_blocks():
     # Test problem when modin frames have the same number of rows, but different
     # blocks (partition.list_of_blocks). See #2322 for details
@@ -163,7 +171,13 @@ def func_to_apply(partition, row_internal_indices, col_internal_indices, item):
         "large_df_plus_small_dfs",
     ],
 )
-def test_rebalance_partitions(test_type):
+@pytest.mark.parametrize(
+    "set_num_partitions",
+    [1, 4],
+    indirect=True,
+)
+def test_rebalance_partitions(test_type, set_num_partitions):
+    num_partitions = NPartitions.get()
     if test_type == "many_small_dfs":
         small_dfs = [
             pd.DataFrame(
@@ -204,8 +218,8 @@ def test_rebalance_partitions(test_type):
         col_length = 103
     large_modin_frame = large_df._query_compiler._modin_frame
     assert large_modin_frame._partitions.shape == (
-        NPartitions.get(),
-        NPartitions.get(),
+        num_partitions,
+        num_partitions,
     ), "Partitions were not rebalanced after concat."
     assert all(
         isinstance(ptn, large_modin_frame._partition_mgr_cls._column_partitions_class)
@@ -218,16 +232,16 @@ def col_apply_func(col):
         assert len(col) == col_length, "Partial axis partition detected."
         return col + 1
 
-    large_df = large_df.apply(col_apply_func)
-    new_large_modin_frame = large_df._query_compiler._modin_frame
-    assert new_large_modin_frame._partitions.shape == (
-        NPartitions.get(),
-        NPartitions.get(),
+    large_apply_result = large_df.apply(col_apply_func)
+    large_apply_result_frame = large_apply_result._query_compiler._modin_frame
+    assert large_apply_result_frame._partitions.shape == (
+        num_partitions,
+        num_partitions,
     ), "Partitions list shape is incorrect."
     assert all(
-        isinstance(ptn, new_large_modin_frame._partition_mgr_cls._partition_class)
-        for ptn in new_large_modin_frame._partitions.flatten()
-    ), "Partitions are not block partitioned after apply."
+        isinstance(ptn, large_apply_result_frame._partition_mgr_cls._partition_class)
+        for ptn in large_apply_result_frame._partitions.flatten()
+    ), "Partitions are not block partitioned after column-wise apply."
     large_df = pd.DataFrame(
         query_compiler=large_df._query_compiler.__constructor__(large_modin_frame)
     )
@@ -238,16 +252,27 @@ def row_apply_func(row):
         assert len(row) == 1000, "Partial axis partition detected."
         return row + 1
 
-    large_df = large_df.apply(row_apply_func, axis=1)
-    new_large_modin_frame = large_df._query_compiler._modin_frame
-    assert new_large_modin_frame._partitions.shape == (
-        4,
-        4,
+    large_apply_result = large_df.apply(row_apply_func, axis=1)
+    large_apply_result_frame = large_apply_result._query_compiler._modin_frame
+    assert large_apply_result_frame._partitions.shape == (
+        num_partitions,
+        num_partitions,
+    ), "Partitions list shape is incorrect."
+    assert all(
+        isinstance(ptn, large_apply_result_frame._partition_mgr_cls._partition_class)
+        for ptn in large_apply_result_frame._partitions.flatten()
+    ), "Partitions are not block partitioned after row-wise apply."
+
+    large_apply_result = large_df.applymap(lambda x: x)
+    large_apply_result_frame = large_apply_result._query_compiler._modin_frame
+    assert large_apply_result_frame._partitions.shape == (
+        num_partitions,
+        num_partitions,
     ), "Partitions list shape is incorrect."
     assert all(
-        isinstance(ptn, new_large_modin_frame._partition_mgr_cls._partition_class)
-        for ptn in new_large_modin_frame._partitions.flatten()
-    ), "Partitions are not block partitioned after apply."
+        isinstance(ptn, large_apply_result_frame._partition_mgr_cls._partition_class)
+        for ptn in large_apply_result_frame._partitions.flatten()
+    ), "Partitions are not block partitioned after element-wise apply."
 
 
 @pytest.mark.skipif(