Repository: modin-project/modin

Issue:
Merge partial dtype caches on `concat(axis=0)`
we could have merged 'known_dtypes':
```python
import modin.pandas as pd
import numpy as np
from modin.core.dataframe.pandas.metadata import ModinDtypes, DtypesDescriptor

df1 = pd.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]})
df2 = pd.DataFrame({"a": [3.0, 4.0, 5.4], "b": [True, True, False]})
df2._query_compiler._modin_frame.set_dtypes_cache(
    ModinDtypes(
        DtypesDescriptor({"a": np.dtype(float)}, cols_with_unknown_dtypes=["b"])
    )
)

res = pd.concat([df1, df2])

# known dtypes: {};
# cols with unknown dtypes: ['a', 'b'];
print(res._query_compiler._modin_frame._dtypes)
# Expected:
# known_dtypes: {"a": float}
# cols_with_unknown_dtypes: ["b"]
```

Discussion:

Patch:
diff --git a/modin/core/dataframe/base/dataframe/utils.py b/modin/core/dataframe/base/dataframe/utils.py
index 7a1d5d98..75c1c999 100644
--- 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/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
index b9a57271..81092b7e 100644
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -2804,6 +2804,7 @@ class PandasDataframe(ClassLogger):
         new_index=None,
         new_columns=None,
         keep_remaining=False,
+        new_dtypes=None,
     ):
         """
         Apply a function across an entire axis for a subset of the data.
@@ -2826,6 +2827,10 @@ class PandasDataframe(ClassLogger):
             advance, and if not provided it must be computed.
         keep_remaining : boolean, default: False
             Whether or not to drop the data that is not computed over.
+        new_dtypes : ModinDtypes or pandas.Series, optional
+            The data types of the result. This is an optimization
+            because there are functions that always result in a particular data
+            type, and allows us to avoid (re)computing it.
 
         Returns
         -------
@@ -2854,7 +2859,9 @@ class PandasDataframe(ClassLogger):
             new_index = self.index if axis == 1 else None
         if new_columns is None:
             new_columns = self.columns if axis == 0 else None
-        return self.__constructor__(new_partitions, new_index, new_columns, None, None)
+        return self.__constructor__(
+            new_partitions, new_index, new_columns, None, None, dtypes=new_dtypes
+        )
 
     @lazy_metadata_decorator(apply_axis="both")
     def apply_select_indices(
@@ -2908,7 +2915,7 @@ class PandasDataframe(ClassLogger):
             new_index = self.index if axis == 1 else None
         if new_columns is None:
             new_columns = self.columns if axis == 0 else None
-        if new_columns is not None and new_dtypes is not None:
+        if new_columns is not None and isinstance(new_dtypes, pandas.Series):
             assert new_dtypes.index.equals(
                 new_columns
             ), f"{new_dtypes=} doesn't have the same columns as in {new_columns=}"
@@ -3688,16 +3695,7 @@ class PandasDataframe(ClassLogger):
                 new_index = self.index.append([other.index for other in others])
             new_columns = joined_index
             frames = [self] + others
-            if all(frame.has_materialized_dtypes for frame in frames):
-                all_dtypes = [frame.dtypes for frame in frames]
-                if not all(dtypes.empty for dtypes in all_dtypes):
-                    new_dtypes = pandas.concat(all_dtypes, axis=1)
-                    # 'nan' value will be placed in a row if a column doesn't exist in all frames;
-                    # this value is np.float64 type so we need an explicit conversion
-                    new_dtypes.fillna(np.dtype("float64"), inplace=True)
-                    new_dtypes = new_dtypes.apply(
-                        lambda row: find_common_type(row.values), axis=1
-                    )
+            new_dtypes = ModinDtypes.concat([frame._dtypes for frame in frames], axis=1)
             # If we have already cached the length of each row in at least one
             # of the row's partitions, we can build new_lengths for the new
             # frame. Typically, if we know the length for any partition in a
diff --git a/modin/core/dataframe/pandas/metadata/__init__.py b/modin/core/dataframe/pandas/metadata/__init__.py
index 4caf2383..2ace6d8e 100644
--- a/modin/core/dataframe/pandas/metadata/__init__.py
+++ b/modin/core/dataframe/pandas/metadata/__init__.py
@@ -13,7 +13,18 @@
 
 """Utilities and classes to handle work with metadata."""
 
-from .dtypes import DtypesDescriptor, LazyProxyCategoricalDtype, ModinDtypes
+from .dtypes import (
+    DtypesDescriptor,
+    LazyProxyCategoricalDtype,
+    ModinDtypes,
+    extract_dtype,
+)
 from .index import ModinIndex
 
-__all__ = ["ModinDtypes", "ModinIndex", "LazyProxyCategoricalDtype", "DtypesDescriptor"]
+__all__ = [
+    "ModinDtypes",
+    "ModinIndex",
+    "LazyProxyCategoricalDtype",
+    "DtypesDescriptor",
+    "extract_dtype",
+]
diff --git a/modin/core/dataframe/pandas/metadata/dtypes.py b/modin/core/dataframe/pandas/metadata/dtypes.py
index ffd81eb7..35674441 100644
--- a/modin/core/dataframe/pandas/metadata/dtypes.py
+++ b/modin/core/dataframe/pandas/metadata/dtypes.py
@@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, Callable, Optional, Union
 import numpy as np
 import pandas
 from pandas._typing import IndexLabel
+from pandas.core.dtypes.cast import find_common_type
 
 if TYPE_CHECKING:
     from modin.core.dataframe.pandas.dataframe.dataframe import PandasDataframe
@@ -446,20 +447,120 @@ class DtypesDescriptor:
         return known_dtypes
 
     @classmethod
-    def concat(
+    def _merge_dtypes(
         cls, values: list[Union["DtypesDescriptor", pandas.Series, None]]
-    ) -> "DtypesDescriptor":  # noqa: GL08
+    ) -> "DtypesDescriptor":
+        """
+        Union columns described by ``values`` and compute common dtypes for them.
+
+        Parameters
+        ----------
+        values : list of DtypesDescriptors, pandas.Series or Nones
+
+        Returns
+        -------
+        DtypesDescriptor
+        """
+        known_dtypes = {}
+        cols_with_unknown_dtypes = []
+        know_all_names = True
+        dtypes_are_unknown = False
+
+        # index - joined column names, columns - dtypes taken from 'values'
+        #        0     1     2      3
+        #  col1  int   bool  float  int
+        #  col2  int   int   int    int
+        #  colN  bool  bool  bool   int
+        dtypes_matrix = pandas.DataFrame()
+
+        for i, val in enumerate(values):
+            if isinstance(val, cls):
+                know_all_names &= val._know_all_names
+                dtypes = val._known_dtypes.copy()
+                dtypes.update({col: "unknown" for col in val._cols_with_unknown_dtypes})
+                if val._remaining_dtype is not None:
+                    # we can't process remaining dtypes, so just discarding them
+                    know_all_names = False
+
+                # setting a custom name to the Series to prevent duplicated names
+                # in the 'dtypes_matrix'
+                series = pandas.Series(dtypes, name=i)
+                dtypes_matrix = pandas.concat([dtypes_matrix, series], axis=1)
+                dtypes_matrix.fillna(
+                    value={
+                        # If we encountered a 'NaN' while 'val' describes all the columns, then
+                        # it means, that the missing columns for this instance will be filled with NaNs (floats),
+                        # otherwise, it may indicate missing columns that this 'val' has no info about,
+                        # meaning that we shouldn't try computing a new dtype for this column,
+                        # so marking it as 'unknown'
+                        i: np.dtype(float)
+                        if val._know_all_names and val._remaining_dtype is None
+                        else "unknown"
+                    },
+                    inplace=True,
+                )
+            elif isinstance(val, pandas.Series):
+                dtypes_matrix = pandas.concat([dtypes_matrix, val], axis=1)
+            elif val is None:
+                # one of the 'dtypes' is None, meaning that we wouldn't been infer a valid result dtype,
+                # however, we're continuing our loop so we would at least know the columns we're missing
+                # dtypes for
+                dtypes_are_unknown = True
+                know_all_names = False
+            else:
+                raise NotImplementedError(type(val))
+
+        if dtypes_are_unknown:
+            return DtypesDescriptor(
+                cols_with_unknown_dtypes=dtypes_matrix.index,
+                know_all_names=know_all_names,
+            )
+
+        def combine_dtypes(row):
+            if (row == "unknown").any():
+                return "unknown"
+            row = row.fillna(np.dtype("float"))
+            return find_common_type(list(row.values))
+
+        dtypes = dtypes_matrix.apply(combine_dtypes, axis=1)
+
+        for col, dtype in dtypes.items():
+            if dtype == "unknown":
+                cols_with_unknown_dtypes.append(col)
+            else:
+                known_dtypes[col] = dtype
+
+        return DtypesDescriptor(
+            known_dtypes,
+            cols_with_unknown_dtypes,
+            remaining_dtype=None,
+            know_all_names=know_all_names,
+        )
+
+    @classmethod
+    def concat(
+        cls, values: list[Union["DtypesDescriptor", pandas.Series, None]], axis: int = 0
+    ) -> "DtypesDescriptor":
         """
         Concatenate dtypes descriptors into a single descriptor.
 
         Parameters
         ----------
         values : list of DtypesDescriptors and pandas.Series
+        axis : int, default: 0
+            If ``axis == 0``: concatenate column names. This implements the logic of
+            how dtypes are combined on ``pd.concat([df1, df2], axis=1)``.
+            If ``axis == 1``: perform a union join for the column names described by
+            `values` and then find common dtypes for the columns appeared to be in
+            an intersection. This implements the logic of how dtypes are combined on
+            ``pd.concat([df1, df2], axis=0).dtypes``.
 
         Returns
         -------
         DtypesDescriptor
         """
+        if axis == 1:
+            return cls._merge_dtypes(values)
         known_dtypes = {}
         cols_with_unknown_dtypes = []
         schema_is_known = True
@@ -529,14 +630,23 @@ class ModinDtypes:
 
     Parameters
     ----------
-    value : pandas.Series or callable
+    value : pandas.Series, callable, DtypesDescriptor or ModinDtypes, optional
     """
 
-    def __init__(self, value: Union[Callable, pandas.Series, DtypesDescriptor]):
+    def __init__(
+        self,
+        value: Optional[
+            Union[Callable, pandas.Series, DtypesDescriptor, "ModinDtypes"]
+        ],
+    ):
         if callable(value) or isinstance(value, pandas.Series):
             self._value = value
         elif isinstance(value, DtypesDescriptor):
             self._value = value.to_series() if value.is_materialized else value
+        elif isinstance(value, type(self)):
+            self._value = value.copy()._value
+        elif isinstance(value, None):
+            self._value = DtypesDescriptor()
         else:
             raise ValueError(f"ModinDtypes doesn't work with '{value}'")
 
@@ -627,13 +737,20 @@ class ModinDtypes:
         return ModinDtypes(self._value.iloc[ids] if numeric_index else self._value[ids])
 
     @classmethod
-    def concat(cls, values: list) -> "ModinDtypes":
+    def concat(cls, values: list, axis: int = 0) -> "ModinDtypes":
         """
-        Concatenate dtypes..
+        Concatenate dtypes.
 
         Parameters
         ----------
         values : list of DtypesDescriptors, pandas.Series, ModinDtypes and Nones
+        axis : int, default: 0
+            If ``axis == 0``: concatenate column names. This implements the logic of
+            how dtypes are combined on ``pd.concat([df1, df2], axis=1)``.
+            If ``axis == 1``: perform a union join for the column names described by
+            `values` and then find common dtypes for the columns appeared to be in
+            an intersection. This implements the logic of how dtypes are combined on
+            ``pd.concat([df1, df2], axis=0).dtypes``.
 
         Returns
         -------
@@ -649,12 +766,16 @@ class ModinDtypes:
                 raise NotImplementedError(type(val))
 
         try:
-            desc = DtypesDescriptor.concat(preprocessed_vals)
+            desc = DtypesDescriptor.concat(preprocessed_vals, axis=axis)
         except NotImplementedError as e:
-            # 'DtypesDescriptor' doesn't support duplicated labels, however, if all values are pandas Serieses,
+            # 'DtypesDescriptor' doesn't support duplicated labels, however, if all values are pandas Series,
             # we still can perform concatenation using pure pandas
-            if "duplicated" not in e.args[0].lower() or not all(
-                isinstance(val, pandas.Series) for val in values
+            if (
+                # 'pd.concat(axis=1)' fails on duplicated labels anyway, so doing this logic
+                # only in case 'axis=0'
+                axis == 0
+                and "duplicated" not in e.args[0].lower()
+                or not all(isinstance(val, pandas.Series) for val in values)
             ):
                 raise e
             desc = pandas.concat(values)
@@ -985,3 +1106,27 @@ def get_categories_dtype(
         if isinstance(cdt, LazyProxyCategoricalDtype)
         else cdt.categories.dtype
     )
+
+
+def extract_dtype(value):
+    """
+    Extract dtype(s) from the passed `value`.
+
+    Parameters
+    ----------
+    value : object
+
+    Returns
+    -------
+    numpy.dtype or pandas.Series of numpy.dtypes
+    """
+    from modin.pandas.utils import is_scalar
+
+    if hasattr(value, "dtype"):
+        return value.dtype
+    elif hasattr(value, "dtypes"):
+        return value.dtypes
+    elif is_scalar(value):
+        return np.dtype(type(value))
+    else:
+        return np.array(value).dtype
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
index 0565a7d1..14020e03 100644
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -58,7 +58,11 @@ from modin.core.dataframe.algebra.default2pandas.groupby import (
     SeriesGroupByDefault,
 )
 from modin.core.dataframe.base.dataframe.utils import join_columns
-from modin.core.dataframe.pandas.metadata import ModinDtypes
+from modin.core.dataframe.pandas.metadata import (
+    DtypesDescriptor,
+    ModinDtypes,
+    extract_dtype,
+)
 from modin.core.storage_formats.base.query_compiler import BaseQueryCompiler
 from modin.error_message import ErrorMessage
 from modin.utils import (
@@ -573,13 +577,36 @@ class PandasQueryCompiler(BaseQueryCompiler):
                     # 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(
@@ -2920,6 +2947,22 @@ class PandasQueryCompiler(BaseQueryCompiler):
             idx = self.get_axis(axis ^ 1).get_indexer_for([key])[0]
             return self.insert_item(axis ^ 1, idx, value, how, replace=True)
 
+        if axis == 0:
+            value_dtype = extract_dtype(value)
+
+            old_columns = self.columns.difference(pandas.Index([key]))
+            old_dtypes = ModinDtypes(self._modin_frame._dtypes).lazy_get(old_columns)
+            new_dtypes = ModinDtypes.concat(
+                [
+                    old_dtypes,
+                    DtypesDescriptor({key: value_dtype}, cols_with_unknown_dtypes=[]),
+                ]
+                # get dtypes in a proper order
+            ).lazy_get(self.columns)
+        else:
+            # TODO: apply 'find_common_dtype' to the value's dtype and old column dtypes
+            new_dtypes = None
+
         # TODO: rework by passing list-like values to `apply_select_indices`
         # as an item to distribute
         if is_list_like(value):
@@ -2930,6 +2973,7 @@ class PandasQueryCompiler(BaseQueryCompiler):
                 new_index=self.index,
                 new_columns=self.columns,
                 keep_remaining=True,
+                new_dtypes=new_dtypes,
             )
         else:
             new_modin_frame = self._modin_frame.apply_select_indices(
@@ -2938,6 +2982,7 @@ class PandasQueryCompiler(BaseQueryCompiler):
                 [key],
                 new_index=self.index,
                 new_columns=self.columns,
+                new_dtypes=new_dtypes,
                 keep_remaining=True,
             )
         return self.__constructor__(new_modin_frame)
@@ -3112,6 +3157,17 @@ class PandasQueryCompiler(BaseQueryCompiler):
             df.insert(internal_idx, column, value)
             return df
 
+        value_dtype = extract_dtype(value)
+        new_columns = self.columns.insert(loc, column)
+        new_dtypes = ModinDtypes.concat(
+            [
+                self._modin_frame._dtypes,
+                DtypesDescriptor({column: value_dtype}, cols_with_unknown_dtypes=[]),
+            ]
+        ).lazy_get(
+            new_columns
+        )  # get dtypes in a proper order
+
         # TODO: rework by passing list-like values to `apply_select_indices`
         # as an item to distribute
         new_modin_frame = self._modin_frame.apply_full_axis_select_indices(
@@ -3120,7 +3176,8 @@ class PandasQueryCompiler(BaseQueryCompiler):
             numeric_indices=[loc],
             keep_remaining=True,
             new_index=self.index,
-            new_columns=self.columns.insert(loc, column),
+            new_columns=new_columns,
+            new_dtypes=new_dtypes,
         )
         return self.__constructor__(new_modin_frame)
 


Test patch:
diff --git a/modin/pandas/test/dataframe/test_join_sort.py b/modin/pandas/test/dataframe/test_join_sort.py
index 19f0f3d6..0d4ab89f 100644
--- a/modin/pandas/test/dataframe/test_join_sort.py
+++ b/modin/pandas/test/dataframe/test_join_sort.py
@@ -471,6 +471,36 @@ def test_merge_on_index(has_index_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)
diff --git a/modin/test/storage_formats/pandas/test_internals.py b/modin/test/storage_formats/pandas/test_internals.py
index 5279614b..c205bf1b 100644
--- a/modin/test/storage_formats/pandas/test_internals.py
+++ b/modin/test/storage_formats/pandas/test_internals.py
@@ -1729,7 +1729,7 @@ class TestModinDtypes:
         )
         assert res.equals(exp)
 
-    def test_concat(self):
+    def test_concat_axis_0(self):
         res = DtypesDescriptor.concat(
             [
                 DtypesDescriptor(self.schema[["a", "b"]]),
@@ -1815,6 +1815,147 @@ class TestModinDtypes:
         )
         assert res.equals(exp)
 
+    @pytest.mark.parametrize(
+        "initial_dtypes, result_cols_with_known_dtypes, result_cols_with_unknown_dtypes",
+        [
+            [
+                # initial dtypes (cols_with_known_dtypes, cols_with_unknown_dtypes, remaining_dtype):
+                # dtypes for all columns are known
+                [
+                    (["a", "b", "c", "d"], [], None),
+                    (["a", "b", "e", "d"], [], None),
+                    (["a", "b"], [], None),
+                ],
+                # result_cols_with_known_dtypes:
+                # all dtypes were known in the beginning, expecting the same
+                # for the result
+                ["a", "b", "c", "d", "e"],
+                # result_cols_with_unknown_dtypes
+                [],
+            ],
+            [
+                # initial dtypes (cols_with_known_dtypes, cols_with_unknown_dtypes, remaining_dtype)
+                [
+                    (["a", "b"], ["c", "d"], None),
+                    (["a", "b", "d"], ["e"], None),
+                    (["a", "b"], [], None),
+                ],
+                # result_cols_with_known_dtypes:
+                # across all dataframes, dtypes were only known for 'a' and 'b' columns
+                ["a", "b"],
+                # result_cols_with_unknown_dtypes
+                ["c", "d", "e"],
+            ],
+            [
+                # initial dtypes (cols_with_known_dtypes, cols_with_unknown_dtypes, remaining_dtype):
+                # the 'e' column in the second frame is missing here, emulating 'know_all_names=False' case
+                [
+                    (["a", "b"], ["c", "d"], None),
+                    (["a", "b", "d"], [], None),
+                    (["a", "b"], [], None),
+                ],
+                # result_cols_with_known_dtypes
+                ["a", "b"],
+                # result_cols_with_unknown_dtypes:
+                # the missing 'e' column will be deducted from the resulted frame after '.concat()'
+                ["c", "d", "e"],
+            ],
+            [
+                # initial dtypes (cols_with_known_dtypes, cols_with_unknown_dtypes, remaining_dtype)
+                # the 'c' column in the first frame is described using 'remaining_dtype'
+                [
+                    (["a", "b", "d"], [], np.dtype(bool)),
+                    (["a", "b", "e", "d"], [], None),
+                    (["a", "b"], [], None),
+                ],
+                # result_cols_with_known_dtypes:
+                # remaining dtypes are not supported by 'concat(axis=0)', so dtype for the 'c'
+                # column is missing here
+                ["a", "b", "e", "d"],
+                # result_cols_with_unknown_dtypes:
+                ["c"],
+            ],
+        ],
+    )
+    def test_concat_axis_1(
+        self,
+        initial_dtypes,
+        result_cols_with_known_dtypes,
+        result_cols_with_unknown_dtypes,
+    ):
+        """
+        Test that ``DtypesDescriptor.concat(axis=1)`` works as expected.
+
+        Parameters
+        ----------
+        initial_dtypes : list of tuples: (cols_with_known_dtypes, cols_with_unknown_dtypes, remaining_dtype)
+            Describe how to build ``DtypesDescriptor`` for each of the three dataframes.
+        result_cols_with_known_dtypes : list of labels
+            Column names for which dtypes has to be determined after ``.concat()``.
+        result_cols_with_unknown_dtypes : list of labels
+            Column names for which dtypes has to be unknown after ``.concat()``.
+        """
+        md_df1, pd_df1 = create_test_dfs(
+            {
+                "a": [1, 2, 3],
+                "b": [3.5, 4.5, 5.5],
+                "c": [True, False, True],
+                "d": ["a", "b", "c"],
+            }
+        )
+        md_df2, pd_df2 = create_test_dfs(
+            {
+                "a": [1.5, 2.5, 3.5],
+                "b": [3.5, 4.5, 5.5],
+                "e": [True, False, True],
+                "d": ["a", "b", "c"],
+            }
+        )
+        md_df3, pd_df3 = create_test_dfs({"a": [1, 2, 3], "b": [3.5, 4.5, 5.5]})
+
+        for md_df, (known_cols, unknown_cols, remaining_dtype) in zip(
+            [md_df1, md_df2, md_df3], initial_dtypes
+        ):
+            known_dtypes = {col: md_df.dtypes[col] for col in known_cols}
+            know_all_names = (
+                len(known_cols) + len(unknown_cols) == len(md_df.columns)
+                or remaining_dtype is not None
+            )
+            # setting columns cache to 'None', in order to prevent completing 'dtypes' with the materialized columns
+            md_df._query_compiler._modin_frame.set_columns_cache(None)
+            md_df._query_compiler._modin_frame.set_dtypes_cache(
+                ModinDtypes(
+                    DtypesDescriptor(
+                        known_dtypes,
+                        unknown_cols,
+                        remaining_dtype,
+                        know_all_names=know_all_names,
+                    )
+                )
+            )
+        md_dtypes = pd.concat(
+            [md_df1, md_df2, md_df3]
+        )._query_compiler._modin_frame._dtypes
+        pd_dtypes = pandas.concat([pd_df1, pd_df2, pd_df3]).dtypes
+        if len(result_cols_with_known_dtypes) == len(pd_dtypes):
+            md_dtypes = (
+                md_dtypes if isinstance(md_dtypes, pandas.Series) else md_dtypes._value
+            )
+            assert isinstance(md_dtypes, pandas.Series)
+            assert md_dtypes.equals(pd_dtypes)
+        else:
+            assert set(md_dtypes._value._known_dtypes.keys()) == set(
+                result_cols_with_known_dtypes
+            )
+            # reindexing to ensure proper order
+            md_known_dtypes = pandas.Series(md_dtypes._value._known_dtypes).reindex(
+                result_cols_with_known_dtypes
+            )
+            assert md_known_dtypes.equals(pd_dtypes[result_cols_with_known_dtypes])
+            assert set(md_dtypes._value._cols_with_unknown_dtypes) == set(
+                result_cols_with_unknown_dtypes
+            )
+
     def test_ModinDtypes_duplicated_concat(self):
         # test that 'ModinDtypes' is able to perform dtypes concatenation on duplicated labels
         # if all of them are Serieses
@@ -1965,6 +2106,119 @@ class TestZeroComputationDtypes:
     Test cases that shouldn't trigger dtypes computation during their execution.
     """
 
+    @pytest.mark.parametrize("self_dtype", ["materialized", "partial", "unknown"])
+    @pytest.mark.parametrize(
+        "value, value_dtype",
+        [
+            [3.5, np.dtype(float)],
+            [[3.5, 2.4], np.dtype(float)],
+            [np.array([3.5, 2.4]), np.dtype(float)],
+            [pd.Series([3.5, 2.4]), np.dtype(float)],
+        ],
+    )
+    def test_preserve_dtypes_setitem(self, self_dtype, value, value_dtype):
+        """
+        Test that ``df[single_existing_column] = value`` preserves dtypes cache.
+        """
+        with mock.patch.object(PandasDataframe, "_compute_dtypes") as patch:
+            df = pd.DataFrame({"a": [1, 2], "b": [3, 4], "c": [3, 4]})
+            if self_dtype == "materialized":
+                assert df._query_compiler._modin_frame.has_materialized_dtypes
+            elif self_dtype == "partial":
+                df._query_compiler._modin_frame.set_dtypes_cache(
+                    ModinDtypes(
+                        DtypesDescriptor(
+                            {"a": np.dtype(int)}, cols_with_unknown_dtypes=["b", "c"]
+                        )
+                    )
+                )
+            elif self_dtype == "unknown":
+                df._query_compiler._modin_frame.set_dtypes_cache(None)
+            else:
+                raise NotImplementedError(self_dtype)
+
+            df["b"] = value
+
+            if self_dtype == "materialized":
+                result_dtype = pandas.Series(
+                    [np.dtype(int), value_dtype, np.dtype(int)], index=["a", "b", "c"]
+                )
+                assert df._query_compiler._modin_frame.has_materialized_dtypes
+                assert df.dtypes.equals(result_dtype)
+            elif self_dtype == "partial":
+                result_dtype = DtypesDescriptor(
+                    {"a": np.dtype(int), "b": value_dtype},
+                    cols_with_unknown_dtypes=["c"],
+                    columns_order={0: "a", 1: "b", 2: "c"},
+                )
+                df._query_compiler._modin_frame._dtypes._value.equals(result_dtype)
+            elif self_dtype == "unknown":
+                result_dtype = DtypesDescriptor(
+                    {"b": value_dtype},
+                    cols_with_unknown_dtypes=["a", "b"],
+                    columns_order={0: "a", 1: "b", 2: "c"},
+                )
+                df._query_compiler._modin_frame._dtypes._value.equals(result_dtype)
+            else:
+                raise NotImplementedError(self_dtype)
+
+        patch.assert_not_called()
+
+    @pytest.mark.parametrize("self_dtype", ["materialized", "partial", "unknown"])
+    @pytest.mark.parametrize(
+        "value, value_dtype",
+        [
+            [3.5, np.dtype(float)],
+            [[3.5, 2.4], np.dtype(float)],
+            [np.array([3.5, 2.4]), np.dtype(float)],
+            [pd.Series([3.5, 2.4]), np.dtype(float)],
+        ],
+    )
+    def test_preserve_dtypes_insert(self, self_dtype, value, value_dtype):
+        with mock.patch.object(PandasDataframe, "_compute_dtypes") as patch:
+            df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
+            if self_dtype == "materialized":
+                assert df._query_compiler._modin_frame.has_materialized_dtypes
+            elif self_dtype == "partial":
+                df._query_compiler._modin_frame.set_dtypes_cache(
+                    ModinDtypes(
+                        DtypesDescriptor(
+                            {"a": np.dtype(int)}, cols_with_unknown_dtypes=["b"]
+                        )
+                    )
+                )
+            elif self_dtype == "unknown":
+                df._query_compiler._modin_frame.set_dtypes_cache(None)
+            else:
+                raise NotImplementedError(self_dtype)
+
+            df.insert(loc=0, column="c", value=value)
+
+            if self_dtype == "materialized":
+                result_dtype = pandas.Series(
+                    [value_dtype, np.dtype(int), np.dtype(int)], index=["c", "a", "b"]
+                )
+                assert df._query_compiler._modin_frame.has_materialized_dtypes
+                assert df.dtypes.equals(result_dtype)
+            elif self_dtype == "partial":
+                result_dtype = DtypesDescriptor(
+                    {"a": np.dtype(int), "c": value_dtype},
+                    cols_with_unknown_dtypes=["b"],
+                    columns_order={0: "c", 1: "a", 2: "b"},
+                )
+                df._query_compiler._modin_frame._dtypes._value.equals(result_dtype)
+            elif self_dtype == "unknown":
+                result_dtype = DtypesDescriptor(
+                    {"c": value_dtype},
+                    cols_with_unknown_dtypes=["a", "b"],
+                    columns_order={0: "c", 1: "a", 2: "b"},
+                )
+                df._query_compiler._modin_frame._dtypes._value.equals(result_dtype)
+            else:
+                raise NotImplementedError(self_dtype)
+
+        patch.assert_not_called()
+
     def test_get_dummies_case(self):
         with mock.patch.object(PandasDataframe, "_compute_dtypes") as patch:
             df = pd.DataFrame(