Repository: modin-project/modin

Issue:
Preserve dtypes cache on `df[existing_col] = scalar`
It currently loses all dtypes on `.__setitem__()`
```python
import modin.pandas as pd

df = pd.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]})
df["b"] = 10

# known dtypes: {};
# cols with unknown dtypes: ['a', 'b'];
print(df._query_compiler._modin_frame._dtypes)
```


Discussion:

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
@@ -2915,7 +2915,7 @@ def apply_select_indices(
             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=}"
diff --git a/modin/core/dataframe/pandas/metadata/__init__.py b/modin/core/dataframe/pandas/metadata/__init__.py
--- 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
--- a/modin/core/dataframe/pandas/metadata/dtypes.py
+++ b/modin/core/dataframe/pandas/metadata/dtypes.py
@@ -529,14 +529,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}'")
 
@@ -985,3 +994,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
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -58,7 +58,11 @@
     SeriesGroupByDefault,
 )
 from modin.core.dataframe.base.dataframe.utils import join_columns
-from modin.core.dataframe.pandas.metadata import DtypesDescriptor, 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 (
@@ -2920,6 +2924,22 @@ def setitem_builder(df, internal_indices=[]):  # pragma: no cover
             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 +2950,7 @@ def setitem_builder(df, internal_indices=[]):  # pragma: no cover
                 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 +2959,7 @@ def setitem_builder(df, internal_indices=[]):  # pragma: no cover
                 [key],
                 new_index=self.index,
                 new_columns=self.columns,
+                new_dtypes=new_dtypes,
                 keep_remaining=True,
             )
         return self.__constructor__(new_modin_frame)
@@ -3112,13 +3134,7 @@ def insert(df, internal_indices=[]):  # pragma: no cover
             df.insert(internal_idx, column, value)
             return df
 
-        if hasattr(value, "dtype"):
-            value_dtype = value.dtype
-        elif is_scalar(value):
-            value_dtype = np.dtype(type(value))
-        else:
-            value_dtype = np.array(value).dtype
-
+        value_dtype = extract_dtype(value)
         new_columns = self.columns.insert(loc, column)
         new_dtypes = ModinDtypes.concat(
             [


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
@@ -1965,6 +1965,64 @@ 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",