Repository: modin-project/modin

Issue:
`.merge()` may produce unbalanced partitioning
'Merge' implementation always preserves partitioning of the left frame:
[`qc.merge -> mf.apply_full_axis()`](https://github.com/modin-project/modin/blob/b2f7680fa2ad387ee0eb7de0633d0c86572e5f2b/modin/core/storage_formats/pandas/query_compiler.py#L463) -> [`keep_partitioning=True`](https://github.com/modin-project/modin/blob/b2f7680fa2ad387ee0eb7de0633d0c86572e5f2b/modin/core/dataframe/pandas/dataframe/dataframe.py#L2189)

This could potentially lead to unbalanced partitioning in case a user merges a bigger frame into a smaller one:
```python
import modin.pandas as pd
import numpy as np

NROWS = 10_000

small_frame = pd.DataFrame({"col0": np.arange(NROWS), "a": np.arange(NROWS)})
big_frame = pd.DataFrame({f"col{i}": np.arange(NROWS) for i in range(300)})

print(f"Left frame partitioning: {small_frame._query_compiler._modin_frame._partitions.shape}")
print(f"Right frame partitioning: {big_frame._query_compiler._modin_frame._partitions.shape}")

result = small_frame.merge(big_frame, on="col0")
print(f"\nMerged frame partitioning: {result._query_compiler._modin_frame._partitions.shape}")
print(f"Merger frame shape: {result.shape}")
```
Output:
```
Left frame partitioning: (112, 1)
Right frame partitioning: (112, 10)

Merged frame partitioning: (112, 1)
Merger frame shape: (10000, 301)
```
In the example above the `.merge` clearly produce suboptimal partitioning by bounding 300 columns into a single column partition. As `.merge` is able to significantly change the shape of the 'left' frame it's proposed to explicitly pass `keep_partitioning=False` to `.apply_full_axis` so partitioning would automatically adapt to the new number of columns.
https://github.com/modin-project/modin/blob/b2f7680fa2ad387ee0eb7de0633d0c86572e5f2b/modin/core/storage_formats/pandas/query_compiler.py#L462-L464


Discussion:

Resolving patch:
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -460,7 +460,13 @@ def map_func(left, right=right, kwargs=kwargs):
                 return pandas.merge(left, right, **kwargs)
 
             new_self = self.__constructor__(
-                self._modin_frame.apply_full_axis(1, map_func)
+                self._modin_frame.apply_full_axis(
+                    axis=1,
+                    func=map_func,
+                    # We're going to explicitly change the shape across the 1-axis,
+                    # so we want for partitioning to adapt as well
+                    keep_partitioning=False,
+                )
             )
             is_reset_index = True
             if left_on and right_on:
@@ -505,7 +511,13 @@ def map_func(left, right=right, kwargs=kwargs):
                 return pandas.DataFrame.join(left, right, **kwargs)
 
             new_self = self.__constructor__(
-                self._modin_frame.apply_full_axis(1, map_func)
+                self._modin_frame.apply_full_axis(
+                    axis=1,
+                    func=map_func,
+                    # We're going to explicitly change the shape across the 1-axis,
+                    # so we want for partitioning to adapt as well
+                    keep_partitioning=False,
+                )
             )
             return new_self.sort_rows_by_column_values(on) if sort else new_self
         else:


Test patch: