Repository: modin-project/modin

Issue:
Preserve columns for merge on index in simple cases
At the moment, we're not preserving columns/dtypes cache if merging on an index level as it appears to be quite complex to mimic pandas with a limited amount of metadata available:
https://github.com/modin-project/modin/blob/bee2c28a3cededa4c5c4b61e9e59c77401ae39a8/modin/core/dataframe/base/dataframe/utils.py#L99-L104

However, there are simple cases (and quite popular cases) when merging on an index that could be supported by simply adding a few lines of code. For example, when merging on a single index level.




Discussion:

Resolving patch:
diff --git a/modin/core/dataframe/base/dataframe/utils.py b/modin/core/dataframe/base/dataframe/utils.py
--- a/modin/core/dataframe/base/dataframe/utils.py
+++ b/modin/core/dataframe/base/dataframe/utils.py
@@ -96,6 +96,19 @@ def join_columns(
         Sequence[IndexLabel], [right_on] if is_scalar(right_on) else right_on
     )
 
+    # handling a simple case of merging on one column and when the column is located in an index
+    if len(left_on) == 1 and len(right_on) == 1 and left_on[0] == right_on[0]:
+        if left_on[0] not in left and right_on[0] not in right:
+            # in this case the 'on' column will stay in the index, so we can simply
+            # drop the 'left/right_on' values and proceed as normal
+            left_on = []
+            right_on = []
+        # in other cases, we can simply add the index name to columns and proceed as normal
+        elif left_on[0] not in left:
+            left = left.insert(loc=0, item=left_on[0])  # type: ignore
+        elif right_on[0] not in right:
+            right = right.insert(loc=0, item=right_on[0])  # type: ignore
+
     if any(col not in left for col in left_on) or any(
         col not in right for col in right_on
     ):
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
@@ -573,13 +573,36 @@ def map_func(
                     # is really complicated in this case, so we're not computing resulted columns for now.
                     pass
                 else:
-                    if self._modin_frame.has_materialized_dtypes:
-                        new_dtypes = []
-                        for old_col in left_renamer.keys():
-                            new_dtypes.append(self.dtypes[old_col])
-                        for old_col in right_renamer.keys():
-                            new_dtypes.append(right_pandas.dtypes[old_col])
-                        new_dtypes = pandas.Series(new_dtypes, index=new_columns)
+                    # renamers may contain columns from 'index', so trying to merge index and column dtypes here
+                    right_index_dtypes = (
+                        right_pandas.index.dtypes
+                        if isinstance(right_pandas.index, pandas.MultiIndex)
+                        else pandas.Series(
+                            [right_pandas.index.dtype], index=[right_pandas.index.name]
+                        )
+                    )
+                    right_dtypes = pandas.concat(
+                        [right_pandas.dtypes, right_index_dtypes]
+                    )[right_renamer.keys()].rename(right_renamer)
+
+                    left_index_dtypes = None
+                    if self._modin_frame.has_materialized_index:
+                        left_index_dtypes = (
+                            self.index.dtypes
+                            if isinstance(self.index, pandas.MultiIndex)
+                            else pandas.Series(
+                                [self.index.dtype], index=[self.index.name]
+                            )
+                        )
+
+                    left_dtypes = (
+                        ModinDtypes.concat(
+                            [self._modin_frame._dtypes, left_index_dtypes]
+                        )
+                        .lazy_get(left_renamer.keys())
+                        .set_index(list(left_renamer.values()))
+                    )
+                    new_dtypes = ModinDtypes.concat([left_dtypes, right_dtypes])
 
             new_self = self.__constructor__(
                 self._modin_frame.apply_full_axis(


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
@@ -471,6 +471,36 @@ def setup_cache():
         )
 
 
+@pytest.mark.parametrize(
+    "left_index", [[], ["key"], ["key", "b"], ["key", "b", "c"], ["b"], ["b", "c"]]
+)
+@pytest.mark.parametrize(
+    "right_index", [[], ["key"], ["key", "e"], ["key", "e", "f"], ["e"], ["e", "f"]]
+)
+def test_merge_on_single_index(left_index, right_index):
+    """
+    Test ``.merge()`` method when merging on a single column, that is located in an index level of one of the frames.
+    """
+    modin_df1, pandas_df1 = create_test_dfs(
+        {"b": [3, 4, 4, 5], "key": [1, 1, 2, 2], "c": [2, 3, 2, 2], "d": [2, 1, 3, 1]}
+    )
+    if len(left_index):
+        modin_df1 = modin_df1.set_index(left_index)
+        pandas_df1 = pandas_df1.set_index(left_index)
+
+    modin_df2, pandas_df2 = create_test_dfs(
+        {"e": [3, 4, 4, 5], "f": [2, 3, 2, 2], "key": [1, 1, 2, 2], "h": [2, 1, 3, 1]}
+    )
+    if len(right_index):
+        modin_df2 = modin_df2.set_index(right_index)
+        pandas_df2 = pandas_df2.set_index(right_index)
+    eval_general(
+        (modin_df1, modin_df2),
+        (pandas_df1, pandas_df2),
+        lambda dfs: dfs[0].merge(dfs[1], on="key"),
+    )
+
+
 @pytest.mark.parametrize("axis", [0, 1])
 @pytest.mark.parametrize(
     "ascending", bool_arg_values, ids=arg_keys("ascending", bool_arg_keys)