Repository: modin-project/modin
Pull request: FIX-#4314: Allow passing a series of dtypes to astype.
Signed-off-by: mvashishtha <mahesh@ponder.io>

<!--
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] commit message follows format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
- [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 #4314
- [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 -->
- [x] added (Issue Number: PR title (PR Number)) and github username to release notes for next major release <!-- e.g. DOCS-#4077: Add release notes template to docs folder (#4078) -->


Problem statement:
Can't use astype() to assign a series of dtypes to a frame.
### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: MacOS Monterey
- **Modin version** (`modin.__version__`): latest source
- **Python version**: 3.9.10
- **Code we can use to reproduce**:

```python
import modin.pandas as pd
df = pd.DataFrame([[0]])
df.astype(df.dtypes)
```

<!--
You can obtain the Modin version with

python -c "import modin; print(modin.__version__)"
-->

### Describe the problem
If I try to pass a series of dtypes to `astype`, I get `ValueError: The truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().` However, I can do that in pandas.

### Source code / logs
<details>
<summary>Show stack trace</summary>

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [10], in <module>
      1 import modin.pandas as pd
      3 df = pd.DataFrame([[0]])
----> 4 df.astype(df.dtypes)

File ~/modin/modin/pandas/base.py:921, in BasePandasDataset.astype(self, dtype, copy, errors)
    918     for column in self._query_compiler.columns:
    919         col_dtypes[column] = dtype
--> 921 new_query_compiler = self._query_compiler.astype(col_dtypes)
    922 return self._create_or_update_from_compiler(new_query_compiler, not copy)

File ~/modin/modin/core/storage_formats/pandas/query_compiler.py:1508, in PandasQueryCompiler.astype(self, col_dtypes, **kwargs)
   1507 def astype(self, col_dtypes, **kwargs):
-> 1508     return self.__constructor__(self._modin_frame.astype(col_dtypes))

File ~/modin/modin/core/dataframe/pandas/dataframe/dataframe.py:114, in lazy_metadata_decorator.<locals>.decorator.<locals>.run_f_on_minimally_updated_metadata(self, *args, **kwargs)
    112     elif apply_axis == "rows":
    113         obj._propagate_index_objs(axis=0)
--> 114 result = f(self, *args, **kwargs)
    115 if apply_axis is None and not transpose:
    116     result._deferred_index = self._deferred_index

File ~/modin/modin/core/dataframe/pandas/dataframe/dataframe.py:1003, in PandasDataframe.astype(self, col_dtypes)
   1000 except TypeError:
   1001     new_dtype = dtype
-> 1003 if dtype != np.int32 and new_dtype == np.int32:
   1004     new_dtypes[column] = np.dtype("int64")
   1005 elif dtype != np.float32 and new_dtype == np.float32:

File /usr/local/lib/python3.9/site-packages/pandas/core/generic.py:1535, in NDFrame.__nonzero__(self)
   1533 @final
   1534 def __nonzero__(self):
-> 1535     raise ValueError(
   1536         f"The truth value of a {type(self).__name__} is ambiguous. "
   1537         "Use a.empty, a.bool(), a.item(), a.any() or a.all()."
   1538     )

ValueError: The truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().

```
</details>

Hints from issue discussion:

Commit message: Delete a redundant import.

Signed-off-by: mvashishtha <mahesh@ponder.io>

Patch to review:
diff --git a/docs/release_notes/release_notes-0.14.0.rst b/docs/release_notes/release_notes-0.14.0.rst
index 99c9b6a6190..53145249af3 100644
--- a/docs/release_notes/release_notes-0.14.0.rst
+++ b/docs/release_notes/release_notes-0.14.0.rst
@@ -23,6 +23,7 @@ Key Features and Updates
   * FIX-#4185: Fix rolling across column partitions (#4262)
   * FIX-#4303: Fix the syntax error in reading from postgres (#4304)
   * FIX-#4308: Add proper error handling in df.set_index (#4309)
+  * FIX-#4314: Allow passing a series of dtypes to astype (#4318)
 * Performance enhancements
   * FIX-#4138, FIX-#4009: remove redundant sorting in the internal '.mask()' flow (#4140)
   * FIX-#4183: Stop shallow copies from creating global shared state. (#4184)
diff --git a/modin/pandas/base.py b/modin/pandas/base.py
index c4cdd6dab01..2feb8e3d7e1 100644
--- a/modin/pandas/base.py
+++ b/modin/pandas/base.py
@@ -51,6 +51,7 @@
 from .utils import is_full_grab_slice
 from modin.utils import try_cast_to_pandas, _inherit_docstrings
 from modin.error_message import ErrorMessage
+import modin.pandas as pd
 from modin.pandas.utils import is_scalar
 from modin.config import IsExperimental
 
@@ -487,8 +488,6 @@ def _default_to_pandas(self, op, *args, **kwargs):
             return Series(result)
         # inplace
         elif result is None:
-            import modin.pandas as pd
-
             return self._create_or_update_from_compiler(
                 getattr(pd, type(pandas_obj).__name__)(pandas_obj)._query_compiler,
                 inplace=True,
@@ -903,20 +902,31 @@ def asof(self, where, subset=None):
         return result
 
     def astype(self, dtype, copy=True, errors="raise"):
-        col_dtypes = {}
+        # dtype can be a series, a dict, or a scalar. If it's series or scalar,
+        # convert it to a dict before passing it to the query compiler.
+        if isinstance(dtype, (pd.Series, pandas.Series)):
+            if not dtype.index.is_unique:
+                raise ValueError(
+                    "The new Series of types must have a unique index, i.e. "
+                    + "it must be one-to-one mapping from column names to "
+                    + " their new dtypes."
+                )
+            dtype = {column: dtype for column, dtype in dtype.items()}
+        # If we got a series or dict originally, dtype is a dict now. Its keys
+        # must be column names.
         if isinstance(dtype, dict):
             if (
                 not set(dtype.keys()).issubset(set(self._query_compiler.columns))
                 and errors == "raise"
             ):
                 raise KeyError(
-                    "Only a column name can be used for the key in"
+                    "Only a column name can be used for the key in "
                     + "a dtype mappings argument."
                 )
             col_dtypes = dtype
         else:
-            for column in self._query_compiler.columns:
-                col_dtypes[column] = dtype
+            # Assume that the dtype is a scalar.
+            col_dtypes = {column: dtype for column in self._query_compiler.columns}
 
         new_query_compiler = self._query_compiler.astype(col_dtypes)
         return self._create_or_update_from_compiler(new_query_compiler, not copy)
diff --git a/modin/pandas/test/dataframe/test_map_metadata.py b/modin/pandas/test/dataframe/test_map_metadata.py
index 53358805d3e..5290ffc7457 100644
--- a/modin/pandas/test/dataframe/test_map_metadata.py
+++ b/modin/pandas/test/dataframe/test_map_metadata.py
@@ -483,6 +483,32 @@ def test_astype():
     with pytest.raises(KeyError):
         modin_df.astype({"not_exists": np.uint8})
 
+    # The dtypes series must have a unique index.
+    with pytest.raises(ValueError):
+        modin_df.astype(pd.Series([str, str], index=["col1", "col1"]))
+
+
+@pytest.mark.parametrize("dtypes_are_dict", [True, False])
+def test_astype_dict_or_series_multiple_column_partitions(dtypes_are_dict):
+    # Test astype with a dtypes dict that is complex in that:
+    # - It applies to columns spanning multiple column partitions
+    # - Within a partition frame df:
+    #   - dtypes.index is not a subset of df.columns
+    #   - df.columns is not a subset of dtypes.index
+
+    pandas_df = pandas.DataFrame(test_data["int_data"])
+    modin_df = pd.DataFrame(pandas_df)
+    if dtypes_are_dict:
+        new_dtypes = {}
+    else:
+        new_dtypes = pandas.Series()
+    for i, column in enumerate(pandas_df.columns):
+        if i % 3 == 1:
+            new_dtypes[column] = "string"
+        elif i % 3 == 2:
+            new_dtypes[column] = float
+    df_equals(modin_df.astype(new_dtypes), pandas_df.astype(new_dtypes))
+
 
 def test_astype_category():
     modin_df = pd.DataFrame(
diff --git a/modin/pandas/test/test_series.py b/modin/pandas/test/test_series.py
index 03bd1f69ea5..7f6013d7921 100644
--- a/modin/pandas/test/test_series.py
+++ b/modin/pandas/test/test_series.py
@@ -1016,6 +1016,8 @@ def test_asof_large(where):
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
 def test_astype(data):
     modin_series, pandas_series = create_test_series(data)
+    series_name = "test_series"
+    modin_series.name = pandas_series.name = series_name
     try:
         pandas_result = pandas_series.astype(str)
     except Exception as e:
@@ -1040,6 +1042,17 @@ def test_astype(data):
     else:
         df_equals(modin_series.astype(np.float64), pandas_result)
 
+    df_equals(
+        modin_series.astype({series_name: str}),
+        pandas_series.astype({series_name: str}),
+    )
+
+    with pytest.raises(KeyError):
+        modin_series.astype({"wrong_name": str})
+
+    # TODO(https://github.com/modin-project/modin/issues/4317): Test passing a
+    # dict to astype() for a series with no name.
+
 
 def test_astype_categorical():
     modin_df = pd.Series(["A", "A", "B", "B", "A"])


Code review comments:
diff_hunk: @@ -483,6 +483,32 @@ def test_astype():
     with pytest.raises(KeyError):
         modin_df.astype({"not_exists": np.uint8})
 
+    # The dtypes series must have a unique index.
+    with pytest.raises(ValueError):
+        modin_df.astype(pd.Series([str, str], index=["col1", "col1"]))
+
+
+@pytest.mark.parametrize("dtypes_are_dict", [True, False])
+def test_astype_dict_or_series_multiple_column_partitions(dtypes_are_dict):
+    # Test astype with a dtypes dict that is complex in that:
+    # - It applies to columns spanning multiple column partitions
+    # - Within a partition frame df:
+    #   - dtypes.index is not a subset of df.columns
+    #   - df.columns is not a subset of dtypes.index
+
+    pandas_df = pandas.DataFrame(test_data["int_data"])
+    modin_df = pd.DataFrame(pandas_df)
path: modin/pandas/test/dataframe/test_map_metadata.py
text: @user1:
```suggestion
    modin_df, pandas_df = create_test_dfs(test_data["int_data"])
```

diff_hunk: @@ -1040,6 +1042,17 @@ def test_astype(data):
     else:
         df_equals(modin_series.astype(np.float64), pandas_result)
 
+    df_equals(
+        modin_series.astype({series_name: str}),
+        pandas_series.astype({series_name: str}),
+    )
+
+    with pytest.raises(KeyError):
+        modin_series.astype({"wrong_name": str})
path: modin/pandas/test/test_series.py
text: @user1:
Use eval_general where possible.

@author:
Done.

diff_hunk: @@ -483,6 +483,32 @@ def test_astype():
     with pytest.raises(KeyError):
         modin_df.astype({"not_exists": np.uint8})
 
+    # The dtypes series must have a unique index.
+    with pytest.raises(ValueError):
+        modin_df.astype(pd.Series([str, str], index=["col1", "col1"]))
+
+
+@pytest.mark.parametrize("dtypes_are_dict", [True, False])
+def test_astype_dict_or_series_multiple_column_partitions(dtypes_are_dict):
+    # Test astype with a dtypes dict that is complex in that:
+    # - It applies to columns spanning multiple column partitions
+    # - Within a partition frame df:
+    #   - dtypes.index is not a subset of df.columns
+    #   - df.columns is not a subset of dtypes.index
+
+    pandas_df = pandas.DataFrame(test_data["int_data"])
+    modin_df = pd.DataFrame(pandas_df)
+    if dtypes_are_dict:
+        new_dtypes = {}
+    else:
+        new_dtypes = pandas.Series()
+    for i, column in enumerate(pandas_df.columns):
+        if i % 3 == 1:
+            new_dtypes[column] = "string"
+        elif i % 3 == 2:
+            new_dtypes[column] = float
+    df_equals(modin_df.astype(new_dtypes), pandas_df.astype(new_dtypes))
path: modin/pandas/test/dataframe/test_map_metadata.py
text: @user1:
```suggestion
    eval_general(modin_df, pandas_df, lambda df: df.astype(new_dtypes))
```