Repository: mars-project/mars

Issue:
[BUG] `astype` cannot accept ExtensionDtype
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

`astype` cannot accept ExtensionDtype.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version
2. The version of Mars you use
3. Versions of crucial packages, such as numpy, scipy and protobuf
4. Full stack of the error.
5. Minimized code to reproduce the error.

```
In [1]: import mars.dataframe as md                                                                                                                                                                         

In [2]: df = md.DataFrame({'a': [1, 2, 3]})                                                                                                                                                                 

In [3]: df                                                                                                                                                                                                  
Out[3]: DataFrame <op=DataFrameDataSource, key=5446e9fe2c3fcbd0484646129f2c04f9>

In [4]: df['a'] = df['a'].astype('string')                                                                                                                                                                  
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/Workspace/mars/mars/_utils.pyx in mars._utils.Tokenizer.tokenize()
    101         try:
--> 102             handler = self._handlers[object_type]
    103             return handler(obj)

KeyError: <class 'pandas.core.arrays.string_.StringDtype'>

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-4-205311623643> in <module>
----> 1 df['a'] = df['a'].astype('string')

~/Workspace/mars/mars/dataframe/base/astype.py in astype(df, dtype, copy, errors)
    305                     raise KeyError('Only a column name can be used for the key in a dtype mappings argument.')
    306     op = DataFrameAstype(dtype_values=dtype, errors=errors)
--> 307     r = op(df)
    308     if not copy:
    309         df.data = r.data

~/Workspace/mars/mars/dataframe/base/astype.py in __call__(self, df)
    179                 dtype = df.dtype
    180             return self.new_series([df], shape=df.shape, dtype=dtype,
--> 181                                    name=df.name, index_value=df.index_value)
    182         else:
    183             empty_df = build_empty_df(df.dtypes)

~/Workspace/mars/mars/dataframe/operands.py in new_series(self, inputs, shape, dtype, index_value, name, **kw)
     59 
     60         return self.new_seriess(inputs, shape=shape, dtype=dtype,
---> 61                                 index_value=index_value, name=name, **kw)[0]
     62 
     63     def new_indexes(self, inputs, shape=None, dtype=None, index_value=None, name=None,

~/Workspace/mars/mars/dataframe/operands.py in new_seriess(self, inputs, shape, dtype, index_value, name, chunks, nsplits, output_limit, kws, **kw)
     52         return self.new_tileables(inputs, shape=shape, dtype=dtype, index_value=index_value,
     53                                   name=name, chunks=chunks, nsplits=nsplits,
---> 54                                   output_limit=output_limit, kws=kws, **kw)
     55 
     56     def new_series(self, inputs, shape=None, dtype=None, index_value=None, name=None, **kw):

~/Workspace/mars/mars/operands.py in new_tileables(self, inputs, kws, **kw)
    424         """
    425 
--> 426         tileables = self._new_tileables(inputs, kws=kws, **kw)
    427         if is_eager_mode():
    428             ExecutableTuple(tileables).execute(fetch=False)

~/Workspace/mars/mars/operands.py in _new_tileables(self, inputs, kws, **kw)
    398             if kws:
    399                 create_tensor_kw.update(kws[j])
--> 400             tileable = self._create_tileable(j, **create_tensor_kw)
    401             tileables.append(tileable)
    402 

~/Workspace/mars/mars/operands.py in _create_tileable(self, output_idx, **kw)
    378 
    379         kw = self._fill_nan_shape(kw)
--> 380         data = tileable_data_type(**kw)
    381         return tileable_type(data)
    382 

~/Workspace/mars/mars/dataframe/core.py in __init__(self, op, shape, nsplits, dtype, name, index_value, chunks, **kw)
    679                  name=None, index_value=None, chunks=None, **kw):
    680         super().__init__(_op=op, _shape=shape, _nsplits=nsplits, _dtype=dtype, _name=name,
--> 681                          _index_value=index_value, _chunks=chunks, **kw)
    682         self._accessors = dict()
    683 

~/Workspace/mars/mars/core.py in __init__(self, *args, **kwargs)
    415             kwargs['_nsplits'] = tuple(tuple(s) for s in kwargs['_nsplits'])
    416 
--> 417         super().__init__(*args, **kwargs)
    418 
    419         if hasattr(self, '_chunks') and self._chunks:

~/Workspace/mars/mars/core.py in __init__(self, *args, **kwargs)
    194         extras = AttributeDict((k, kwargs.pop(k)) for k in set(kwargs) - set(self.__slots__))
    195         kwargs['_extra_params'] = kwargs.pop('_extra_params', extras)
--> 196         super().__init__(*args, **kwargs)
    197 
    198     @property

~/Workspace/mars/mars/core.py in __init__(self, *args, **kwargs)
     44 
     45         if self._init_update_key_ and (not hasattr(self, '_key') or not self._key):
---> 46             self._update_key()
     47         if not hasattr(self, '_id') or not self._id:
     48             self._id = str(id(self))

~/Workspace/mars/mars/core.py in _update_key(self)
     74 
     75     def _update_key(self):
---> 76         self._obj_set('_key', tokenize(type(self).__name__, *self._values_))
     77         return self
     78 

~/Workspace/mars/mars/_utils.pyx in mars._utils.tokenize()
     82 
     83 def tokenize(*args, **kwargs):
---> 84     return mmh_hash_bytes(build_canonical_bytes(args, kwargs)).hex()
     85 
     86 

~/Workspace/mars/mars/_utils.pyx in mars._utils.build_canonical_bytes()
     78     if kwargs:
     79         args = args + (kwargs,)
---> 80     return str([tokenize_handler.tokenize(arg) for arg in args]).encode('utf-8')
     81 
     82 

~/Workspace/mars/mars/_utils.pyx in mars._utils.Tokenizer.tokenize()
    111                     self._handlers[object_type] = self._handlers[clz]
    112                     return self._handlers[clz](obj)
--> 113             raise TypeError('Cannot generate token for %s, type: %s' % (obj, object_type))
    114 
    115 

TypeError: Cannot generate token for string, type: <class 'pandas.core.arrays.string_.StringDtype'>
```



Discussion:

Resolving patch:
diff --git a/mars/dataframe/arrays.py b/mars/dataframe/arrays.py
--- a/mars/dataframe/arrays.py
+++ b/mars/dataframe/arrays.py
@@ -157,6 +157,9 @@ def nbytes(self) -> int:
                    for x in chunk.buffers()
                    if x is not None)
 
+    def memory_usage(self, deep=True) -> int:
+        return self.nbytes
+
     @staticmethod
     def _can_process_slice_via_arrow(slc):
         if not isinstance(slc, slice):


Test patch:
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
@@ -1144,6 +1144,14 @@ def testAsType(self):
         expected = raw.astype({'c1': 'int32', 'c2': 'float', 'c8': 'str'})
         pd.testing.assert_frame_equal(expected, result)
 
+        # test arrow_string dtype
+        df = from_pandas_df(raw, chunk_size=8)
+        r = df.astype({'c1': 'arrow_string'})
+
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.astype({'c1': 'arrow_string'})
+        pd.testing.assert_frame_equal(expected, result)
+
         # test series
         s = pd.Series(rs.randint(5, size=20))
         series = from_pandas_series(s)
@@ -1153,6 +1161,13 @@ def testAsType(self):
         expected = s.astype('int32')
         pd.testing.assert_series_equal(result, expected)
 
+        series = from_pandas_series(s, chunk_size=6)
+        r = series.astype('arrow_string')
+
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = s.astype('arrow_string')
+        pd.testing.assert_series_equal(result, expected)
+
         # multiply chunks
         series = from_pandas_series(s, chunk_size=6)
         r = series.astype('str')
diff --git a/mars/dataframe/tests/test_arrays.py b/mars/dataframe/tests/test_arrays.py
--- a/mars/dataframe/tests/test_arrays.py
+++ b/mars/dataframe/tests/test_arrays.py
@@ -164,6 +164,10 @@ def testArrowStringArrayFunctions(self):
         self.assertLess(arrow_array.nbytes,
                         pd.Series(string_array).memory_usage(deep=True))
 
+        # test memory_usage
+        self.assertEqual(arrow_array.memory_usage(deep=True),
+                         arrow_array.nbytes)
+
         # test isna
         np.testing.assert_array_equal(has_na_arrow_array.isna(),
                                       has_na_string_array.isna())