Repository: mars-project/mars

Issue:
Support DataFrame.astype
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**

Support DataFrame.astype to cast a DataFrame object to a specified dtype, `CategoricalDtype` may do some special handling.



Discussion:

Resolving patch:
diff --git a/mars/dataframe/base/__init__.py b/mars/dataframe/base/__init__.py
--- a/mars/dataframe/base/__init__.py
+++ b/mars/dataframe/base/__init__.py
@@ -28,6 +28,7 @@
 from .shift import shift, tshift
 from .diff import df_diff, series_diff
 from .value_counts import value_counts
+from .astype import astype
 
 
 def _install():
@@ -55,6 +56,7 @@ def _install():
         setattr(t, 'shift', shift)
         setattr(t, 'tshift', tshift)
         setattr(t, 'diff', df_diff)
+        setattr(t, 'astype', astype)
 
     for t in SERIES_TYPE:
         setattr(t, 'to_gpu', to_gpu)
@@ -78,6 +80,7 @@ def _install():
         setattr(t, 'tshift', tshift)
         setattr(t, 'diff', series_diff)
         setattr(t, 'value_counts', value_counts)
+        setattr(t, 'astype', astype)
 
     for t in INDEX_TYPE:
         setattr(t, 'rechunk', rechunk)
diff --git a/mars/dataframe/base/astype.py b/mars/dataframe/base/astype.py
new file mode 100644
--- /dev/null
+++ b/mars/dataframe/base/astype.py
@@ -0,0 +1,316 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pandas as pd
+from pandas.api.types import CategoricalDtype
+
+from ... import opcodes as OperandDef
+from ...serialize import AnyField, StringField, ListField
+from ...utils import recursive_tile
+from ...tensor.base import sort
+from ..utils import build_empty_df, build_empty_series
+from ..core import SERIES_TYPE
+from ..operands import DataFrameOperand, DataFrameOperandMixin, ObjectType
+
+
+class DataFrameAstype(DataFrameOperand, DataFrameOperandMixin):
+    _op_type_ = OperandDef.ASTYPE
+
+    _dtype_values = AnyField('dtype_values')
+    _errors = StringField('errors')
+    _category_cols = ListField('category_cols')
+
+    def __init__(self, dtype_values=None, copy=None, errors=None,
+                 category_cols=None, object_type=None, **kw):
+        super().__init__(_dtype_values=dtype_values,
+                         _errors=errors, _category_cols=category_cols,
+                         _object_type=object_type, **kw)
+
+    @property
+    def dtype_values(self):
+        return self._dtype_values
+
+    @property
+    def errors(self):
+        return self._errors
+
+    @property
+    def category_cols(self):
+        return self._category_cols
+
+    @classmethod
+    def _tile_one_chunk(cls, op):
+        c = op.inputs[0].chunks[0]
+        chunk_op = op.copy().reset_key()
+        chunk_params = op.outputs[0].params.copy()
+        chunk_params['index'] = c.index
+        out_chunks = [chunk_op.new_chunk([c], **chunk_params)]
+
+        new_op = op.copy()
+        return new_op.new_tileables(op.inputs, nsplits=op.inputs[0].nsplits,
+                                    chunks=out_chunks, **op.outputs[0].params.copy())
+
+    @classmethod
+    def _tile_series(cls, op):
+        in_series = op.inputs[0]
+        out = op.outputs[0]
+
+        unique_chunk = None
+        if op.dtype_values == 'category' and isinstance(op.dtype_values, str):
+            unique_chunk = recursive_tile(sort(in_series.unique())).chunks[0]
+
+        chunks = []
+        for c in in_series.chunks:
+            chunk_op = op.copy().reset_key()
+            params = c.params.copy()
+            params['dtype'] = out.dtype
+            if unique_chunk is not None:
+                chunk_op._category_cols = [in_series.name]
+                new_chunk = chunk_op.new_chunk([c, unique_chunk], **params)
+            else:
+                new_chunk = chunk_op.new_chunk([c], **params)
+            chunks.append(new_chunk)
+
+        new_op = op.copy()
+        return new_op.new_seriess(op.inputs, nsplits=in_series.nsplits,
+                                  chunks=chunks, **out.params.copy())
+
+    @classmethod
+    def _tile_dataframe(cls, op):
+        in_df = op.inputs[0]
+        out = op.outputs[0]
+        cum_nsplits = np.cumsum((0,) + in_df.nsplits[1])
+        out_chunks = []
+
+        if op.dtype_values == 'category':
+            # all columns need unique values
+            for c in in_df.chunks:
+                chunk_op = op.copy().reset_key()
+                params = c.params.copy()
+                dtypes = out.dtypes[cum_nsplits[c.index[1]]: cum_nsplits[c.index[1] + 1]]
+                params['dtypes'] = dtypes
+                chunk_op._category_cols = list(c.columns_value.to_pandas())
+                unique_chunks = []
+                for col in c.columns_value.to_pandas():
+                    unique_chunks.append(recursive_tile(sort(in_df[col].unique())).chunks[0])
+                new_chunk = chunk_op.new_chunk([c] + unique_chunks, **params)
+                out_chunks.append(new_chunk)
+        elif isinstance(op.dtype_values, dict) and 'category' in op.dtype_values.values():
+            # some columns' types are category
+            category_cols = [c for c, v in op.dtype_values.items()
+                             if isinstance(v, str) and v == 'category']
+            unique_chunks = dict((col, recursive_tile(sort(in_df[col].unique())).chunks[0])
+                                 for col in category_cols)
+            for c in in_df.chunks:
+                chunk_op = op.copy().reset_key()
+                params = c.params.copy()
+                dtypes = out.dtypes[cum_nsplits[c.index[1]]: cum_nsplits[c.index[1] + 1]]
+                params['dtypes'] = dtypes
+                chunk_category_cols = []
+                chunk_unique_chunks = []
+                for col in c.columns_value.to_pandas():
+                    if col in category_cols:
+                        chunk_category_cols.append(col)
+                        chunk_unique_chunks.append(unique_chunks[col])
+                chunk_op._category_cols = chunk_category_cols
+                new_chunk = chunk_op.new_chunk([c] + chunk_unique_chunks, **params)
+                out_chunks.append(new_chunk)
+        else:
+            for c in in_df.chunks:
+                chunk_op = op.copy().reset_key()
+                params = c.params.copy()
+                dtypes = out.dtypes[cum_nsplits[c.index[1]]: cum_nsplits[c.index[1] + 1]]
+                params['dtypes'] = dtypes
+                new_chunk = chunk_op.new_chunk([c], **params)
+                out_chunks.append(new_chunk)
+
+        new_op = op.copy()
+        return new_op.new_dataframes(op.inputs, nsplits=in_df.nsplits,
+                                     chunks=out_chunks, **out.params.copy())
+
+    @classmethod
+    def tile(cls, op):
+        if len(op.inputs[0].chunks) == 1:
+            return cls._tile_one_chunk(op)
+        elif isinstance(op.inputs[0], SERIES_TYPE):
+            return cls._tile_series(op)
+        else:
+            return cls._tile_dataframe(op)
+
+    @classmethod
+    def execute(cls, ctx, op):
+        in_data = ctx[op.inputs[0].key]
+        if not isinstance(op.dtype_values, dict):
+            if op.category_cols is not None:
+                uniques = [ctx[c.key] for c in op.inputs[1:]]
+                dtype = dict((col, CategoricalDtype(unique_values)) for
+                             col, unique_values in zip(op.category_cols, uniques))
+                ctx[op.outputs[0].key] = in_data.astype(dtype, errors=op.errors)
+
+            else:
+                ctx[op.outputs[0].key] = in_data.astype(op.dtype_values, errors=op.errors)
+        else:
+            selected_dtype = dict((k, v) for k, v in op.dtype_values.items()
+                                  if k in in_data.columns)
+            if op.category_cols is not None:
+                uniques = [ctx[c.key] for c in op.inputs[1:]]
+                for col, unique_values in zip(op.category_cols, uniques):
+                    selected_dtype[col] = CategoricalDtype(unique_values)
+            ctx[op.outputs[0].key] = in_data.astype(selected_dtype, errors=op.errors)
+
+    def __call__(self, df):
+        if isinstance(df, SERIES_TYPE):
+            self._object_type = ObjectType.series
+            empty_series = build_empty_series(df.dtype)
+            new_series = empty_series.astype(self.dtype_values, errors=self.errors)
+            if new_series.dtype != df.dtype:
+                dtype = CategoricalDtype() if isinstance(
+                    new_series.dtype, CategoricalDtype) else new_series.dtype
+            else:  # pragma: no cover
+                dtype = df.dtype
+            return self.new_series([df], shape=df.shape, dtype=dtype,
+                                   name=df.name, index_value=df.index_value)
+        else:
+            self._object_type = ObjectType.dataframe
+            empty_df = build_empty_df(df.dtypes)
+            new_df = empty_df.astype(self.dtype_values, errors=self.errors)
+            dtypes = []
+            for dt, new_dt in zip(df.dtypes, new_df.dtypes):
+                if new_dt != dt and isinstance(new_dt, CategoricalDtype):
+                    dtypes.append(CategoricalDtype())
+                else:
+                    dtypes.append(new_dt)
+            dtypes = pd.Series(dtypes, index=new_df.dtypes.index)
+            return self.new_dataframe([df], shape=df.shape, dtypes=dtypes,
+                                      index_value=df.index_value,
+                                      columns_value=df.columns_value)
+
+
+def astype(df, dtype, copy=True, errors='raise'):
+    """
+    Cast a pandas object to a specified dtype ``dtype``.
+
+    Parameters
+    ----------
+    dtype : data type, or dict of column name -> data type
+        Use a numpy.dtype or Python type to cast entire pandas object to
+        the same type. Alternatively, use {col: dtype, ...}, where col is a
+        column label and dtype is a numpy.dtype or Python type to cast one
+        or more of the DataFrame's columns to column-specific types.
+    copy : bool, default True
+        Return a copy when ``copy=True`` (be very careful setting
+        ``copy=False`` as changes to values then may propagate to other
+        pandas objects).
+    errors : {'raise', 'ignore'}, default 'raise'
+        Control raising of exceptions on invalid data for provided dtype.
+
+        - ``raise`` : allow exceptions to be raised
+        - ``ignore`` : suppress exceptions. On error return original object.
+
+    Returns
+    -------
+    casted : same type as caller
+
+    See Also
+    --------
+    to_datetime : Convert argument to datetime.
+    to_timedelta : Convert argument to timedelta.
+    to_numeric : Convert argument to a numeric type.
+    numpy.ndarray.astype : Cast a numpy array to a specified type.
+
+    Examples
+    --------
+    Create a DataFrame:
+
+    >>> import mars.dataframe as md
+    >>> df = md.DataFrame(pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}))
+    >>> df.dtypes
+    col1    int64
+    col2    int64
+    dtype: object
+
+    Cast all columns to int32:
+
+    >>> df.astype('int32').dtypes
+    col1    int32
+    col2    int32
+    dtype: object
+
+    Cast col1 to int32 using a dictionary:
+
+    >>> df.astype({'col1': 'int32'}).dtypes
+    col1    int32
+    col2    int64
+    dtype: object
+
+    Create a series:
+
+    >>> ser = md.Series(pd.Series([1, 2], dtype='int32'))
+    >>> ser.execute()
+    0    1
+    1    2
+    dtype: int32
+    >>> ser.astype('int64').execute()
+    0    1
+    1    2
+    dtype: int64
+
+    Convert to categorical type:
+
+    >>> ser.astype('category').execute()
+    0    1
+    1    2
+    dtype: category
+    Categories (2, int64): [1, 2]
+
+    Convert to ordered categorical type with custom ordering:
+
+    >>> cat_dtype = pd.api.types.CategoricalDtype(
+    ...     categories=[2, 1], ordered=True)
+    >>> ser.astype(cat_dtype).execute()
+    0    1
+    1    2
+    dtype: category
+    Categories (2, int64): [2 < 1]
+
+    Note that using ``copy=False`` and changing data on a new
+    pandas object may propagate changes:
+
+    >>> s1 = md.Series(pd.Series([1, 2]))
+    >>> s2 = s1.astype('int64', copy=False)
+    >>> s1.execute()  # note that s1[0] has changed too
+    0     1
+    1     2
+    dtype: int64
+    """
+    if isinstance(dtype, dict):
+        keys = list(dtype.keys())
+        if isinstance(df, SERIES_TYPE):
+            if len(keys) != 1 or keys[0] != df.name:
+                raise KeyError('Only the Series name can be used for the key in Series dtype mappings.')
+            else:
+                dtype = list(dtype.values())[0]
+        else:
+            for k in keys:
+                columns = df.columns_value.to_pandas()
+                if k not in columns:
+                    raise KeyError('Only a column name can be used for the key in a dtype mappings argument.')
+    op = DataFrameAstype(dtype_values=dtype, errors=errors)
+    r = op(df)
+    if not copy:
+        df.data = r.data
+        return df
+    else:
+        return r
diff --git a/mars/web/apihandlers.py b/mars/web/apihandlers.py
--- a/mars/web/apihandlers.py
+++ b/mars/web/apihandlers.py
@@ -101,7 +101,7 @@ def post(self, session_id):
         try:
             graph = self.get_argument('graph')
             target = self.get_argument('target').split(',')
-            names = self.get_argument('names').split(',')
+            names = self.get_argument('names', default='').split(',')
             compose = bool(int(self.get_argument('compose', '1')))
         except web.MissingArgumentError as ex:
             self.write(json.dumps(dict(msg=str(ex))))


Test patch:
diff --git a/mars/dataframe/base/tests/test_base.py b/mars/dataframe/base/tests/test_base.py
--- a/mars/dataframe/base/tests/test_base.py
+++ b/mars/dataframe/base/tests/test_base.py
@@ -23,7 +23,7 @@
     CATEGORICAL_TYPE, CATEGORICAL_CHUNK_TYPE
 from mars.dataframe.datasource.dataframe import from_pandas as from_pandas_df
 from mars.dataframe.datasource.series import from_pandas as from_pandas_series
-from mars.dataframe.base import to_gpu, to_cpu, df_reset_index, series_reset_index, cut
+from mars.dataframe.base import to_gpu, to_cpu, df_reset_index, series_reset_index, cut, astype
 from mars.dataframe.operands import ObjectType
 from mars.operands import OperandStage
 from mars.tensor.core import TENSOR_TYPE
@@ -820,3 +820,14 @@ def testCut(self):
         self.assertIsInstance(r, TENSOR_TYPE)
         e = pd.cut([0, 1, 1, 2], bins=4, labels=False)
         self.assertEqual(r.dtype, e.dtype)
+
+    def testAstype(self):
+        s = from_pandas_series(pd.Series([1, 2, 1, 2], name='a'), chunk_size=2)
+        with self.assertRaises(KeyError):
+            astype(s, {'b': 'str'})
+
+        df = from_pandas_df(pd.DataFrame({'a': [1, 2, 1, 2],
+                                          'b': ['a', 'b', 'a', 'b']}), chunk_size=2)
+
+        with self.assertRaises(KeyError):
+            astype(df, {'c': 'str', 'a': 'str'})
diff --git a/mars/dataframe/base/tests/test_base_execution.py b/mars/dataframe/base/tests/test_base_execution.py
--- a/mars/dataframe/base/tests/test_base_execution.py
+++ b/mars/dataframe/base/tests/test_base_execution.py
@@ -1182,3 +1182,115 @@ def testValueCountsExecution(self):
             r = series.value_counts(bins=5, normalize=True)
             pd.testing.assert_series_equal(executor.execute_dataframes([r])[0],
                                            s.value_counts(bins=5, normalize=True))
+
+    def testAsType(self):
+        rs = np.random.RandomState(0)
+        raw = pd.DataFrame(rs.randint(1000, size=(20, 8)),
+                           columns=['c' + str(i + 1) for i in range(8)])
+        # single chunk
+        df = from_pandas_df(raw)
+        r = df.astype('int32')
+
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.astype('int32')
+        pd.testing.assert_frame_equal(expected, result)
+
+        # multiply chunks
+        df = from_pandas_df(raw, chunk_size=6)
+        r = df.astype('int32')
+
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.astype('int32')
+        pd.testing.assert_frame_equal(expected, result)
+
+        # dict type
+        df = from_pandas_df(raw, chunk_size=5)
+        r = df.astype({'c1': 'int32', 'c2': 'float', 'c8': 'str'})
+
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.astype({'c1': 'int32', 'c2': 'float', 'c8': 'str'})
+        pd.testing.assert_frame_equal(expected, result)
+
+        # test series
+        s = pd.Series(rs.randint(5, size=20))
+        series = from_pandas_series(s)
+        r = series.astype('int32')
+
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = s.astype('int32')
+        pd.testing.assert_series_equal(result, expected)
+
+        # multiply chunks
+        series = from_pandas_series(s, chunk_size=6)
+        r = series.astype('str')
+
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = s.astype('str')
+        pd.testing.assert_series_equal(result, expected)
+
+        # test category
+        raw = pd.DataFrame(rs.randint(3, size=(20, 8)),
+                           columns=['c' + str(i + 1) for i in range(8)])
+
+        df = from_pandas_df(raw)
+        r = df.astype('category')
+
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.astype('category')
+        pd.testing.assert_frame_equal(expected, result)
+
+        df = from_pandas_df(raw)
+        r = df.astype({'c1': 'category', 'c8': 'int32', 'c4': 'str'})
+
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.astype({'c1': 'category', 'c8': 'int32', 'c4': 'str'})
+        pd.testing.assert_frame_equal(expected, result)
+
+        df = from_pandas_df(raw, chunk_size=5)
+        r = df.astype('category')
+
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.astype('category')
+        pd.testing.assert_frame_equal(expected, result)
+
+        df = from_pandas_df(raw, chunk_size=3)
+        r = df.astype({'c1': 'category', 'c8': 'int32', 'c4': 'str'})
+
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.astype({'c1': 'category', 'c8': 'int32', 'c4': 'str'})
+        pd.testing.assert_frame_equal(expected, result)
+
+        df = from_pandas_df(raw, chunk_size=6)
+        r = df.astype({'c1': 'category', 'c5': 'float', 'c2': 'int32',
+                       'c7': pd.CategoricalDtype([1, 3, 4, 2]),
+                       'c4': pd.CategoricalDtype([1, 3, 2])})
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.astype({'c1': 'category', 'c5': 'float', 'c2': 'int32',
+                               'c7': pd.CategoricalDtype([1, 3, 4, 2]),
+                               'c4': pd.CategoricalDtype([1, 3, 2])})
+        pd.testing.assert_frame_equal(expected, result)
+
+        df = from_pandas_df(raw, chunk_size=8)
+        r = df.astype({'c2': 'category'})
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.astype({'c2': 'category'})
+        pd.testing.assert_frame_equal(expected, result)
+
+        # test series category
+        raw = pd.Series(np.random.choice(['a', 'b', 'c'], size=(10,)))
+        series = from_pandas_series(raw, chunk_size=4)
+        result = self.executor.execute_dataframe(series.astype('category'), concat=True)[0]
+        expected = raw.astype('category')
+        pd.testing.assert_series_equal(expected, result)
+
+        series = from_pandas_series(raw, chunk_size=3)
+        result = self.executor.execute_dataframe(
+            series.astype(pd.CategoricalDtype(['a', 'c', 'b']), copy=False), concat=True)[0]
+        expected = raw.astype(pd.CategoricalDtype(['a', 'c', 'b']),  copy=False)
+        pd.testing.assert_series_equal(expected, result)
+
+        series = from_pandas_series(raw, chunk_size=6)
+        result = self.executor.execute_dataframe(
+            series.astype(pd.CategoricalDtype(['a', 'c', 'b', 'd'])), concat=True)[0]
+        expected = raw.astype(pd.CategoricalDtype(['a', 'c', 'b', 'd']))
+        pd.testing.assert_series_equal(expected, result)