Repository: modin-project/modin (python)

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)
```

Pull request:
PERF-#6753: Preserve dtypes cache on '.__setitem__()'
<!--
Thank you for your contribution!
Please review the contributing docs: https://modin.readthedocs.io/en/latest/development/contributing.html
if you have questions about contributing.
-->

## What do these changes do?

Detect the value's dtype and insert in into partial dtypes for the result of `.setitem()`

- [x] first commit message and PR title follow format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
  > **_NOTE:_**  If you edit the PR title to match this format, you need to add another commit (even if it's empty) or amend your last commit for the CI job that checks the PR title to pick up the new PR title.
- [ ] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [ ] passes `black --check modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [ ] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [ ] Resolves #6753 <!-- issue must be created for each patch -->
- [ ] tests added and passing
- [ ] module layout described at `docs/development/architecture.rst` is up-to-date <!-- if you have added, renamed or removed files or directories please update the documentation accordingly -->


Patch:
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
index 88c97843..e32c1693 100644
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -2915,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=}"
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..680825fa 100644
--- 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
index a34c8d4b..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 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 (
@@ -2943,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):
@@ -2953,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(
@@ -2961,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)
@@ -3135,13 +3157,7 @@ class PandasQueryCompiler(BaseQueryCompiler):
             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
index 243a496c..f315aa0b 100644
--- 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",