Repository: modin-project/modin
Pull request: FIX-#6394: preserve dtypes for __setitem__ op when using not hashable key
<!--
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?

<!-- Please give a short brief about these changes. -->

- [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.
- [x] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] passes `black --check modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [x] Resolves #6394 <!-- issue must be created for each patch -->
- [x] tests added and passing
- [x] 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 -->


Problem statement:
`df[existing_cols] = df` should compute result dtypes where possible
Example:
```python
import modin.pandas as pd

df = pd.DataFrame([[1,2,3,4], [5,6,7,8]])
print(df._query_compiler._modin_frame.has_materialized_dtypes)  # True

df2 = pd.DataFrame([[9,9], [5,5]])
print(df2._query_compiler._modin_frame.has_materialized_dtypes)  # True

df[[1,2]] = df2
print(df._query_compiler._modin_frame.has_materialized_dtypes)  # False, but it should be True
```

Hints from issue discussion:

Commit message: add test

Signed-off-by: Anatoly Myachev <anatoly.myachev@intel.com>

Patch to review:
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
index a5edd031b51..b5f4597593b 100644
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -2804,6 +2804,7 @@ def apply_select_indices(
         col_labels=None,
         new_index=None,
         new_columns=None,
+        new_dtypes=None,
         keep_remaining=False,
         item_to_distribute=no_default,
     ):
@@ -2825,11 +2826,11 @@ def apply_select_indices(
             The column labels to apply over. Must be provided
             with `row_labels` to apply over both axes.
         new_index : list-like, optional
-            The index of the result. We may know this in advance,
-            and if not provided it must be computed.
+            The index of the result, we may know this in advance.
         new_columns : list-like, optional
-            The columns of the result. We may know this in
-            advance, and if not provided it must be computed.
+            The columns of the result, we may know this in advance.
+        new_dtypes : pandas.Series, optional
+            The dtypes of the result, we may know this in advance.
         keep_remaining : boolean, default: False
             Whether or not to drop the data that is not computed over.
         item_to_distribute : np.ndarray or scalar, default: no_default
@@ -2845,6 +2846,12 @@ 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 not new_dtypes.index.equals(new_columns):
+                # sanity check
+                raise ValueError(
+                    f"{new_dtypes=} doesn't have the same columns as in {new_columns=}"
+                )
         if axis is not None:
             assert apply_indices is not None
             # Convert indices to numeric indices
@@ -2872,7 +2879,12 @@ def apply_select_indices(
                 axis ^ 1: [self.row_lengths, self.column_widths][axis ^ 1],
             }
             return self.__constructor__(
-                new_partitions, new_index, new_columns, lengths_objs[0], lengths_objs[1]
+                new_partitions,
+                new_index,
+                new_columns,
+                lengths_objs[0],
+                lengths_objs[1],
+                new_dtypes,
             )
         else:
             # We are applying over both axes here, so make sure we have all the right
@@ -2899,6 +2911,7 @@ def apply_select_indices(
                 new_columns,
                 self._row_lengths_cache,
                 self._column_widths_cache,
+                new_dtypes,
             )
 
     @lazy_metadata_decorator(apply_axis="both")
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
index f3491bb7045..ffd4188e9de 100644
--- a/modin/core/storage_formats/base/query_compiler.py
+++ b/modin/core/storage_formats/base/query_compiler.py
@@ -4326,7 +4326,9 @@ def setitem(df, axis, key, value):
 
         return DataFrameDefault.register(setitem)(self, axis=axis, key=key, value=value)
 
-    def write_items(self, row_numeric_index, col_numeric_index, broadcasted_items):
+    def write_items(
+        self, row_numeric_index, col_numeric_index, item, need_columns_reindex=True
+    ):
         """
         Update QueryCompiler elements at the specified positions by passed values.
 
@@ -4338,15 +4340,21 @@ def write_items(self, row_numeric_index, col_numeric_index, broadcasted_items):
             Row positions to write value.
         col_numeric_index : list of ints
             Column positions to write value.
-        broadcasted_items : 2D-array
-            Values to write. Have to be same size as defined by `row_numeric_index`
-            and `col_numeric_index`.
+        item : Any
+            Values to write. If not a scalar will be broadcasted according to
+            `row_numeric_index` and `col_numeric_index`.
+        need_columns_reindex : bool, default: True
+            In the case of assigning columns to a dataframe (broadcasting is
+            part of the flow), reindexing is not needed.
 
         Returns
         -------
         BaseQueryCompiler
             New QueryCompiler with updated values.
         """
+        # We have to keep this import away from the module level to avoid circular import
+        from modin.pandas.utils import broadcast_item, is_scalar
+
         if not isinstance(row_numeric_index, slice):
             row_numeric_index = list(row_numeric_index)
         if not isinstance(col_numeric_index, slice):
@@ -4358,8 +4366,19 @@ def write_items(df, broadcasted_items):
             df.iloc[row_numeric_index, col_numeric_index] = broadcasted_items
             return df
 
+        if not is_scalar(item):
+            broadcasted_item, _ = broadcast_item(
+                self,
+                row_numeric_index,
+                col_numeric_index,
+                item,
+                need_columns_reindex=need_columns_reindex,
+            )
+        else:
+            broadcasted_item, _ = item, pandas.Series([np.array(item).dtype])
+
         return DataFrameDefault.register(write_items)(
-            self, broadcasted_items=broadcasted_items
+            self, broadcasted_items=broadcasted_item
         )
 
     # END Abstract methods for QueryCompiler
diff --git a/modin/core/storage_formats/cudf/query_compiler.py b/modin/core/storage_formats/cudf/query_compiler.py
index de6e0fcc789..89d48cd6237 100644
--- a/modin/core/storage_formats/cudf/query_compiler.py
+++ b/modin/core/storage_formats/cudf/query_compiler.py
@@ -31,7 +31,12 @@ def transpose(self, *args, **kwargs):
         # Switch the index and columns and transpose the data within the blocks.
         return self.__constructor__(self._modin_frame.transpose())
 
-    def write_items(self, row_numeric_index, col_numeric_index, broadcasted_items):
+    def write_items(
+        self, row_numeric_index, col_numeric_index, item, need_columns_reindex=True
+    ):
+        # We have to keep this import away from the module level to avoid circular import
+        from modin.pandas.utils import broadcast_item, is_scalar
+
         def iloc_mut(partition, row_internal_indices, col_internal_indices, item):
             partition = partition.copy()
             unique_items = np.unique(item)
@@ -54,6 +59,17 @@ def iloc_mut(partition, row_internal_indices, col_internal_indices, item):
                     partition.iloc[i, j] = it
             return partition
 
+        if not is_scalar(item):
+            broadcasted_item, _ = broadcast_item(
+                self,
+                row_numeric_index,
+                col_numeric_index,
+                item,
+                need_columns_reindex=need_columns_reindex,
+            )
+        else:
+            broadcasted_item, _ = item, pandas.Series([np.array(item).dtype])
+
         new_modin_frame = self._modin_frame.apply_select_indices(
             axis=None,
             func=iloc_mut,
@@ -62,6 +78,6 @@ def iloc_mut(partition, row_internal_indices, col_internal_indices, item):
             new_index=self.index,
             new_columns=self.columns,
             keep_remaining=True,
-            item_to_distribute=broadcasted_items,
+            item_to_distribute=broadcasted_item,
         )
         return self.__constructor__(new_modin_frame)
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
index cbe41184d00..59101b80bd3 100644
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -4217,7 +4217,12 @@ def take_2d_positional(self, index=None, columns=None):
             )
         )
 
-    def write_items(self, row_numeric_index, col_numeric_index, broadcasted_items):
+    def write_items(
+        self, row_numeric_index, col_numeric_index, item, need_columns_reindex=True
+    ):
+        # We have to keep this import away from the module level to avoid circular import
+        from modin.pandas.utils import broadcast_item, is_scalar
+
         def iloc_mut(partition, row_internal_indices, col_internal_indices, item):
             """
             Write `value` in a specified location in a single partition.
@@ -4253,6 +4258,30 @@ def iloc_mut(partition, row_internal_indices, col_internal_indices, item):
                 partition.iloc[row_internal_indices, col_internal_indices] = item.copy()
             return partition
 
+        if not is_scalar(item):
+            broadcasted_item, broadcasted_dtypes = broadcast_item(
+                self,
+                row_numeric_index,
+                col_numeric_index,
+                item,
+                need_columns_reindex=need_columns_reindex,
+            )
+        else:
+            broadcasted_item, broadcasted_dtypes = item, pandas.Series(
+                [np.array(item).dtype] * len(col_numeric_index)
+            )
+
+        new_dtypes = None
+        if (
+            # compute dtypes only if assigning entire columns
+            isinstance(row_numeric_index, slice)
+            and row_numeric_index == slice(None)
+            and self._modin_frame.has_materialized_dtypes
+            and broadcasted_dtypes is not None
+        ):
+            new_dtypes = self.dtypes.copy()
+            new_dtypes.iloc[col_numeric_index] = broadcasted_dtypes.values
+
         new_modin_frame = self._modin_frame.apply_select_indices(
             axis=None,
             func=iloc_mut,
@@ -4260,8 +4289,9 @@ def iloc_mut(partition, row_internal_indices, col_internal_indices, item):
             col_labels=col_numeric_index,
             new_index=self.index,
             new_columns=self.columns,
+            new_dtypes=new_dtypes,
             keep_remaining=True,
-            item_to_distribute=broadcasted_items,
+            item_to_distribute=broadcasted_item,
         )
         return self.__constructor__(new_modin_frame)
 
diff --git a/modin/numpy/indexing.py b/modin/numpy/indexing.py
index 272f277eede..aabe28efa8e 100644
--- a/modin/numpy/indexing.py
+++ b/modin/numpy/indexing.py
@@ -472,22 +472,6 @@ def get_axis(axis):
             axis = None
         return axis
 
-    def _write_items(self, row_lookup, col_lookup, item):
-        """
-        Perform remote write and replace blocks.
-
-        Parameters
-        ----------
-        row_lookup : slice or scalar
-            The global row index to write item to.
-        col_lookup : slice or scalar
-            The global col index to write item to.
-        item : numpy.ndarray
-            The new item value that needs to be assigned to `self`.
-        """
-        new_qc = self.arr._query_compiler.write_items(row_lookup, col_lookup, item)
-        self.arr._update_inplace(new_qc)
-
     def _setitem_positional(self, row_lookup, col_lookup, item, axis=None):
         """
         Assign `item` value to located dataset.
@@ -512,16 +496,9 @@ def _setitem_positional(self, row_lookup, col_lookup, item, axis=None):
             row_lookup = range(len(self.arr._query_compiler.index))[row_lookup]
         if isinstance(col_lookup, slice):
             col_lookup = range(len(self.arr._query_compiler.columns))[col_lookup]
-        if axis is None:
-            if not is_scalar(item):
-                item = broadcast_item(self.arr, row_lookup, col_lookup, item)
-            self.arr._query_compiler = self.arr._query_compiler.write_items(
-                row_lookup, col_lookup, item
-            )
-        else:
-            if not is_scalar(item):
-                item = broadcast_item(self.arr, row_lookup, col_lookup, item)
-            self._write_items(row_lookup, col_lookup, item)
+
+        new_qc = self.arr._query_compiler.write_items(row_lookup, col_lookup, item)
+        self.arr._update_inplace(new_qc)
 
     def __setitem__(self, key, item):
         """
diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py
index a00c137fa37..a4a37bbf7b3 100644
--- a/modin/pandas/dataframe.py
+++ b/modin/pandas/dataframe.py
@@ -58,7 +58,6 @@
 from .utils import (
     from_pandas,
     from_non_pandas,
-    broadcast_item,
     cast_function_modin2pandas,
     SET_DATAFRAME_ATTRIBUTE_WARNING,
 )
@@ -2530,16 +2529,12 @@ def __setitem__(self, key, value):
                         value = np.array(value)
                     if len(key) != value.shape[-1]:
                         raise ValueError("Columns must be same length as key")
-                item = broadcast_item(
-                    self,
+                new_qc = self._query_compiler.write_items(
                     slice(None),
-                    key,
+                    self.columns.get_indexer_for(key),
                     value,
                     need_columns_reindex=False,
                 )
-                new_qc = self._query_compiler.write_items(
-                    slice(None), self.columns.get_indexer_for(key), item
-                )
                 self._update_inplace(new_qc)
                 # self.loc[:, key] = value
                 return
diff --git a/modin/pandas/indexing.py b/modin/pandas/indexing.py
index b0e4f37c569..5b2e6d518b7 100644
--- a/modin/pandas/indexing.py
+++ b/modin/pandas/indexing.py
@@ -40,7 +40,7 @@
 
 from .dataframe import DataFrame
 from .series import Series
-from .utils import is_scalar, broadcast_item
+from .utils import is_scalar
 
 
 def is_slice(x):
@@ -435,27 +435,12 @@ def _setitem_positional(self, row_lookup, col_lookup, item, axis=None):
             assert len(row_lookup) == 1
             new_qc = self.qc.setitem(1, self.qc.index[row_lookup[0]], item)
             self.df._create_or_update_from_compiler(new_qc, inplace=True)
+            self.qc = self.df._query_compiler
         # Assignment to both axes.
         else:
-            if not is_scalar(item):
-                item = broadcast_item(self.df, row_lookup, col_lookup, item)
-            self._write_items(row_lookup, col_lookup, item)
-
-    def _write_items(self, row_lookup, col_lookup, item):
-        """
-        Perform remote write and replace blocks.
-
-        Parameters
-        ----------
-        row_lookup : slice or scalar
-            The global row index to write item to.
-        col_lookup : slice or scalar
-            The global col index to write item to.
-        item : numpy.ndarray
-            The new item value that needs to be assigned to `self`.
-        """
-        new_qc = self.qc.write_items(row_lookup, col_lookup, item)
-        self.df._create_or_update_from_compiler(new_qc, inplace=True)
+            new_qc = self.qc.write_items(row_lookup, col_lookup, item)
+            self.df._create_or_update_from_compiler(new_qc, inplace=True)
+            self.qc = self.df._query_compiler
 
     def _determine_setitem_axis(self, row_lookup, col_lookup, row_scalar, col_scalar):
         """
diff --git a/modin/pandas/utils.py b/modin/pandas/utils.py
index abef69318be..6ff02fe7d33 100644
--- a/modin/pandas/utils.py
+++ b/modin/pandas/utils.py
@@ -302,7 +302,7 @@ def broadcast_item(
 
     Parameters
     ----------
-    obj : DataFrame or Series
+    obj : DataFrame or Series or query compiler
         The object containing the necessary information about the axes.
     row_lookup : slice or scalar
         The global row index to locate inside of `item`.
@@ -316,8 +316,9 @@ def broadcast_item(
 
     Returns
     -------
-    np.ndarray
-        `item` after it was broadcasted to `to_shape`.
+    (np.ndarray, Optional[Series])
+        * np.ndarray - `item` after it was broadcasted to `to_shape`.
+        * Optional[Series] - item's dtypes, if previously contained.
 
     Raises
     ------
@@ -345,6 +346,7 @@ def broadcast_item(
     )
     to_shape = new_row_len, new_col_len
 
+    dtypes = None
     if isinstance(item, (pandas.Series, pandas.DataFrame, Series, DataFrame)):
         # convert indices in lookups to names, as pandas reindex expects them to be so
         axes_to_reindex = {}
@@ -358,12 +360,26 @@ def broadcast_item(
         # New value for columns/index make that reindex add NaN values
         if axes_to_reindex:
             item = item.reindex(**axes_to_reindex)
+
+        if (
+            hasattr(item, "_query_compiler")
+            and item._query_compiler._modin_frame.has_materialized_dtypes
+            or hasattr(item, "dtypes")
+        ):
+            dtypes = item.dtypes
+            if not isinstance(dtypes, pandas.Series):
+                dtypes = pandas.Series([dtypes])
+
     try:
+        # Cast to numpy drop information about heterogeneous types (cast to common)
+        # TODO: we shouldn't do that, maybe there should be the if branch
         item = np.array(item)
+        if dtypes is None:
+            dtypes = pandas.Series([item.dtype] * len(col_lookup))
         if np.prod(to_shape) == np.prod(item.shape):
-            return item.reshape(to_shape)
+            return item.reshape(to_shape), dtypes
         else:
-            return np.broadcast_to(item, to_shape)
+            return np.broadcast_to(item, to_shape), dtypes
     except ValueError:
         from_shape = np.array(item).shape
         raise ValueError(
diff --git a/modin/test/storage_formats/pandas/test_internals.py b/modin/test/storage_formats/pandas/test_internals.py
index 5a5f191ea93..0c9c6a7128f 100644
--- a/modin/test/storage_formats/pandas/test_internals.py
+++ b/modin/test/storage_formats/pandas/test_internals.py
@@ -1113,6 +1113,17 @@ def test_setitem_bool_preserve_dtypes():
     assert df._query_compiler._modin_frame.has_materialized_dtypes
 
 
+def test_setitem_unhashable_preserve_dtypes():
+    df = pd.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]])
+    assert df._query_compiler._modin_frame.has_materialized_dtypes
+
+    df2 = pd.DataFrame([[9, 9], [5, 5]])
+    assert df2._query_compiler._modin_frame.has_materialized_dtypes
+
+    df[[1, 2]] = df2
+    assert df._query_compiler._modin_frame.has_materialized_dtypes
+
+
 @pytest.mark.parametrize(
     "modify_config", [{ExperimentalGroupbyImpl: True}], indirect=True
 )


Code review comments:
diff_hunk: @@ -4358,8 +4366,19 @@ def write_items(df, broadcasted_items):
             df.iloc[row_numeric_index, col_numeric_index] = broadcasted_items
             return df
 
+        if not is_scalar(item):
+            broadcasted_item, _ = broadcast_item(
+                self,
+                row_numeric_index,
+                col_numeric_index,
+                item,
+                need_columns_reindex=need_columns_reindex,
+            )
+        else:
+            broadcasted_item, _ = item, pandas.Series([np.array(item).dtype])
path: modin/core/storage_formats/base/query_compiler.py
text: @user1:
huh, what's the point of assigning `pandas.Series()` to the `_` variable here?..

@author:
fixed

diff_hunk: @@ -4253,15 +4258,40 @@ def iloc_mut(partition, row_internal_indices, col_internal_indices, item):
                 partition.iloc[row_internal_indices, col_internal_indices] = item.copy()
             return partition
 
+        if not is_scalar(item):
+            broadcasted_item, broadcasted_dtypes = broadcast_item(
+                self,
+                row_numeric_index,
+                col_numeric_index,
+                item,
+                need_columns_reindex=need_columns_reindex,
+            )
+        else:
+            broadcasted_item, broadcasted_dtypes = item, pandas.Series(
+                [np.array(item).dtype] * len(col_numeric_index)
+            )
+
+        new_dtypes = None
+        if (
+            # compute dtypes only if assigning entire columns
+            isinstance(row_numeric_index, slice)
+            and row_numeric_index == slice(None)
+            and self._modin_frame.has_materialized_dtypes
+            and broadcasted_dtypes is not None
path: modin/core/storage_formats/pandas/query_compiler.py
text: @user1:
I would suggest move this check to be the very first in this condition, as it's the cheapest to compute (super-tiny optimization which is probably useless here but could be a good habit of writing new code that way - from cheapest to heavier, to determine the most probable result with as small cost as possible)

@author:
agree, done

diff_hunk: @@ -2845,6 +2846,12 @@ 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 not new_dtypes.index.equals(new_columns):
+                # sanity check
+                raise ValueError(
+                    f"{new_dtypes=} doesn't have the same columns as in {new_columns=}"
+                )
path: modin/core/dataframe/pandas/dataframe/dataframe.py
text: @author:
```suggestion
            assert new_dtypes.index.equals(new_columns), f"{new_dtypes=} doesn't have the same columns as in {new_columns=}"

```

diff_hunk: @@ -4358,8 +4366,19 @@ def write_items(df, broadcasted_items):
             df.iloc[row_numeric_index, col_numeric_index] = broadcasted_items
             return df
 
+        if not is_scalar(item):
+            broadcasted_item, _ = broadcast_item(
+                self,
+                row_numeric_index,
+                col_numeric_index,
+                item,
+                need_columns_reindex=need_columns_reindex,
+            )
+        else:
+            broadcasted_item, _ = item, pandas.Series([np.array(item).dtype])
path: modin/core/storage_formats/base/query_compiler.py
text: @author:
```suggestion
            broadcasted_item = item
```

diff_hunk: @@ -54,6 +59,17 @@ def iloc_mut(partition, row_internal_indices, col_internal_indices, item):
                     partition.iloc[i, j] = it
             return partition
 
+        if not is_scalar(item):
+            broadcasted_item, _ = broadcast_item(
+                self,
+                row_numeric_index,
+                col_numeric_index,
+                item,
+                need_columns_reindex=need_columns_reindex,
+            )
+        else:
+            broadcasted_item, _ = item, pandas.Series([np.array(item).dtype])
path: modin/core/storage_formats/cudf/query_compiler.py
text: @author:
```suggestion
            broadcasted_item = item
```

diff_hunk: @@ -2825,11 +2826,11 @@ def apply_select_indices(
             The column labels to apply over. Must be provided
             with `row_labels` to apply over both axes.
         new_index : list-like, optional
-            The index of the result. We may know this in advance,
-            and if not provided it must be computed.
+            The index of the result, we may know this in advance.
         new_columns : list-like, optional
-            The columns of the result. We may know this in
-            advance, and if not provided it must be computed.
+            The columns of the result, we may know this in advance.
+        new_dtypes : pandas.Series, optional
+            The dtypes of the result, we may know this in advance.
path: modin/core/dataframe/pandas/dataframe/dataframe.py
text: @user1:
```suggestion
            The columns of the result, if known in advance.
        new_dtypes : pandas.Series, optional
            The dtypes of the result, if known in advance.
```

diff_hunk: @@ -54,6 +59,17 @@ def iloc_mut(partition, row_internal_indices, col_internal_indices, item):
                     partition.iloc[i, j] = it
             return partition
 
+        if not is_scalar(item):
+            broadcasted_item, _ = broadcast_item(
+                self,
+                row_numeric_index,
+                col_numeric_index,
+                item,
+                need_columns_reindex=need_columns_reindex,
+            )
+        else:
+            broadcasted_item, _ = item, pandas.Series([np.array(item).dtype])
path: modin/core/storage_formats/cudf/query_compiler.py
text: @user1:
same question as in base QC - what's the point?

@author:
fixed

diff_hunk: @@ -2845,6 +2846,12 @@ 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 not new_dtypes.index.equals(new_columns):
+                # sanity check
path: modin/core/dataframe/pandas/dataframe/dataframe.py
text: @user1:
if that's a sanity check, maybe it should be
```python
assert new_dtypes.index.equals(new_columns), "error message"
```
instead of `if` then `raise ValueError`?