# 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 itertools
import operator
from copy import copy as copy_obj
from numbers import Integral
from typing import Type, Sequence

import numpy as np
import pandas as pd
from pandas._libs import lib
from pandas.api.indexers import check_array_indexer
from pandas.api.types import pandas_dtype, is_scalar, is_array_like
from pandas.api.extensions import ExtensionArray, ExtensionDtype, register_extension_dtype
from pandas.core import ops
from pandas.core.algorithms import take
from pandas.compat import set_function_name
try:
    from pandas.arrays import StringArray as StringArrayBase
except ImportError:  # for pandas < 1.0
    StringArrayBase = ExtensionArray

try:
    import pyarrow as pa
    pa_null = pa.NULL
except ImportError:  # pragma: no cover
    pa = None
    pa_null = None


@register_extension_dtype
class ArrowStringDtype(ExtensionDtype):
    """
    Extension dtype for arrow string data.

    .. warning::

       ArrowStringDtype is considered experimental. The implementation and
       parts of the API may change without warning.

       In particular, ArrowStringDtype.na_value may change to no longer be
       ``numpy.nan``.

    Attributes
    ----------
    None

    Methods
    -------
    None

    Examples
    --------
    >>> import mars.dataframe as md
    >>> md.ArrowStringDtype()
    ArrowStringDtype
    """

    type = str
    kind = "U"
    name = "arrow_string"
    na_value = pa_null

    @classmethod
    def construct_from_string(cls, string):
        if string == cls.name:
            return cls()
        else:
            raise TypeError(f"Cannot construct a '{cls}' from '{string}'")

    @classmethod
    def construct_array_type(cls) -> "Type[ArrowStringArray]":
        return ArrowStringArray

    def __from_arrow__(self, array):
        return ArrowStringArray(array)


class ArrowStringArray(StringArrayBase):
    def __init__(self, values, copy=False):
        if isinstance(values, (pd.Index, pd.Series)):
            # for pandas Index and Series,
            # convert to PandasArray
            values = values.array

        if isinstance(values, type(self)):
            arrow_array = values._arrow_array
        elif isinstance(values, ExtensionArray):
            # if come from pandas object like index,
            # convert to pandas StringArray first,
            # validation will be done in construct
            arrow_array = pa.chunked_array([pa.array(values, from_pandas=True)])
        elif isinstance(values, pa.ChunkedArray):
            arrow_array = values
        elif isinstance(values, pa.StringArray):
            arrow_array = pa.chunked_array([values])
        else:
            arrow_array = pa.chunked_array([pa.array(values, type=pa.string())])

        if copy:
            arrow_array = copy_obj(arrow_array)

        self._arrow_array = arrow_array
        self._dtype = ArrowStringDtype()

        # for test purpose
        self._force_use_pandas = False

    @classmethod
    def from_scalars(cls, values):
        arrow_array = pa.chunked_array(
            [pa.array(np.asarray(values)).cast(pa.string())])
        return cls(arrow_array)

    @classmethod
    def _from_sequence(cls, scalars, dtype=None, copy=False):
        if not hasattr(scalars, 'dtype'):
            scalars = np.asarray(scalars)
        if isinstance(scalars, cls):
            if copy:
                scalars = scalars.copy()
            return scalars
        arrow_array = pa.chunked_array([pa.array(scalars).cast(pa.string())])
        return cls(arrow_array, copy=copy)

    @classmethod
    def _from_sequence_of_strings(cls, strings, dtype=None, copy=False):
        return cls._from_sequence(strings, dtype=dtype, copy=copy)

    @classmethod
    def _from_factorized(cls, values, original):
        return cls(values)

    def __repr__(self):
        return f"{type(self).__name__}({repr(self._arrow_array)})"

    @property
    def dtype(self):
        return self._dtype

    @property
    def nbytes(self) -> int:
        return sum(x.size
                   for chunk in self._arrow_array.chunks
                   for x in chunk.buffers()
                   if x is not None)

    @staticmethod
    def _can_process_slice_via_arrow(slc):
        if not isinstance(slc, slice):
            return False
        if slc.step is not None and slc.step != 1:
            return False
        if slc.start is not None and \
                not isinstance(slc.start, Integral):  # pragma: no cover
            return False
        if slc.stop is not None and \
                not isinstance(slc.stop, Integral):  # pragma: no cover
            return False
        return True

    def _values_for_factorize(self):
        arr = self.to_numpy()
        mask = self.isna()
        arr[mask] = -1
        return arr, -1

    def _values_for_argsort(self):
        return self.to_numpy()

    @staticmethod
    def _process_pos(pos, length, is_start):
        if pos is None:
            return 0 if is_start else length
        return pos + length if pos < 0 else pos

    def __getitem__(self, item):
        has_take = hasattr(self._arrow_array, 'take')
        if not self._force_use_pandas and has_take:
            if pd.api.types.is_scalar(item):
                item = item + len(self) if item < 0 else item
                return self._arrow_array.take([item]).to_pandas()[0]
            elif self._can_process_slice_via_arrow(item):
                length = len(self)
                start, stop = item.start, item.stop
                start = self._process_pos(start, length, True)
                stop = self._process_pos(stop, length, False)
                return ArrowStringArray(
                    self._arrow_array.slice(offset=start,
                                            length=stop - start))
            elif hasattr(item, 'dtype') and np.issubdtype(item.dtype, np.bool_):
                return ArrowStringArray(self._arrow_array.filter(
                    pa.array(item, from_pandas=True)))
            elif hasattr(item, 'dtype'):
                length = len(self)
                item = np.where(item < 0, item + length, item)
                return ArrowStringArray(self._arrow_array.take(item))

        array = np.asarray(self._arrow_array.to_pandas())
        return ArrowStringArray(array[item])

    def __setitem__(self, key, value):
        if isinstance(value, (pd.Index, pd.Series)):
            value = value.to_numpy()
        if isinstance(value, type(self)):
            value = value.to_numpy()

        key = check_array_indexer(self, key)
        scalar_key = is_scalar(key)
        scalar_value = is_scalar(value)
        if scalar_key and not scalar_value:
            raise ValueError("setting an array element with a sequence.")

        # validate new items
        if scalar_value:
            if pd.isna(value):
                value = None
            elif not isinstance(value, str):
                raise ValueError(
                    f"Cannot set non-string value '{value}' into a StringArray."
                )
        else:
            if not is_array_like(value):
                value = np.asarray(value, dtype=object)
            if len(value) and not lib.is_string_array(value, skipna=True):
                raise ValueError("Must provide strings.")

        string_array = np.asarray(self._arrow_array.to_pandas())
        string_array[key] = value
        self._arrow_array = pa.chunked_array([pa.array(string_array)])

    def __len__(self):
        return len(self._arrow_array)

    def __array__(self, dtype=None):
        return self.to_numpy(dtype=dtype)

    def to_numpy(self, dtype=None, copy=False, na_value=lib.no_default):
        array = np.asarray(self._arrow_array.to_pandas())
        if copy or na_value is not lib.no_default:
            array = array.copy()
        if na_value is not lib.no_default:
            array[self.isna()] = na_value
        return array

    def fillna(self, value=None, method=None, limit=None):
        chunks = []
        for chunk_array in self._arrow_array.chunks:
            array = chunk_array.to_pandas()
            result_array = array.fillna(value=value, method=method,
                                        limit=limit)
            chunks.append(pa.array(result_array, from_pandas=True))
        return ArrowStringArray(pa.chunked_array(chunks))

    def astype(self, dtype, copy=True):
        dtype = pandas_dtype(dtype)
        if isinstance(dtype, ArrowStringDtype):
            if copy:
                return self.copy()
            return self

        # try to slice 1 record to get the result dtype
        test_array = self._arrow_array.slice(0, 1).to_pandas()
        test_result_array = test_array.astype(dtype).array

        result_array = \
            type(test_result_array)(
                np.full(self.shape, test_result_array.dtype.na_value,
                        dtype=np.asarray(test_result_array).dtype))

        start = 0
        # use chunks to do astype
        for chunk_array in self._arrow_array.chunks:
            result_array[start: start + len(chunk_array)] = \
                chunk_array.to_pandas().astype(dtype).array
            start += len(chunk_array)
        return result_array

    def isna(self):
        if not self._force_use_pandas and hasattr(self._arrow_array, 'is_null'):
            return self._arrow_array.is_null().to_pandas().to_numpy()
        else:
            return pd.isna(self._arrow_array.to_pandas()).to_numpy()

    def take(self, indices, allow_fill=False, fill_value=None):
        if allow_fill is False:
            return ArrowStringArray(self[indices])

        string_array = self._arrow_array.to_pandas().to_numpy()

        replace = False
        if allow_fill and fill_value is None:
            fill_value = self.dtype.na_value
            replace = True

        result = take(string_array, indices, fill_value=fill_value,
                      allow_fill=allow_fill)
        if replace:
            # pyarrow cannot recognize pa.NULL
            result[result == self.dtype.na_value] = None
        return ArrowStringArray(result)

    def copy(self):
        return type(self)(copy_obj(self._arrow_array))

    @classmethod
    def _concat_same_type(
            cls, to_concat: Sequence["ArrowStringArray"]) -> "ArrowStringArray":
        chunks = list(itertools.chain.from_iterable(
            x._arrow_array.chunks for x in to_concat))
        if len(chunks) == 0:
            chunks = [pa.array([], type=pa.string())]
        return cls(pa.chunked_array(chunks))

    def value_counts(self, dropna=False):
        string_array = self._arrow_array.to_pandas()
        return ArrowStringArray(string_array.value_counts(dropna=dropna))

    def any(self, axis=0, out=None):
        return self._arrow_array.to_pandas().any(axis=axis, out=out)

    def all(self, axis=0, out=None):
        return self._arrow_array.to_pandas().all(axis=axis, out=out)

    # Overrride parent because we have different return types.
    @classmethod
    def _create_arithmetic_method(cls, op):
        # Note: this handles both arithmetic and comparison methods.
        def method(self, other):
            is_arithmetic = \
                True if op.__name__ in ops.ARITHMETIC_BINOPS else False

            is_other_array = False
            if not is_scalar(other):
                is_other_array = True
                other = np.asarray(other)

            self_is_na = self.isna()
            other_is_na = pd.isna(other)
            mask = self_is_na | other_is_na

            chunks = []
            mask_chunks = []
            start = 0
            for chunk_array in self._arrow_array.chunks:
                chunk_array = np.asarray(chunk_array.to_pandas())
                end = start + len(chunk_array)
                chunk_mask = mask[start: end]
                chunk_valid = ~chunk_mask

                if is_arithmetic:
                    result = np.empty(chunk_array.shape, dtype=object)
                else:
                    result = np.zeros(chunk_array.shape, dtype=bool)

                chunk_other = other
                if is_other_array:
                    chunk_other = other[start: end]
                    chunk_other = chunk_other[chunk_valid]

                # calculate only for both not None
                result[chunk_valid] = op(chunk_array[chunk_valid],
                                         chunk_other)

                if is_arithmetic:
                    chunks.append(pa.array(result, type=pa.string(),
                                           from_pandas=True))
                else:
                    chunks.append(result)
                    mask_chunks.append(chunk_mask)

            if is_arithmetic:
                return ArrowStringArray(pa.chunked_array(chunks))
            else:
                return pd.arrays.BooleanArray(np.concatenate(chunks),
                                              np.concatenate(mask_chunks))

        return set_function_name(method, "__{}__".format(op.__name__), cls)

    @classmethod
    def _add_arithmetic_ops(cls):
        cls.__add__ = cls._create_arithmetic_method(operator.add)
        cls.__radd__ = cls._create_arithmetic_method(ops.radd)

        cls.__mul__ = cls._create_arithmetic_method(operator.mul)
        cls.__rmul__ = cls._create_arithmetic_method(ops.rmul)

    @classmethod
    def _add_comparison_ops(cls):
        cls.__eq__ = cls._create_comparison_method(operator.eq)
        cls.__ne__ = cls._create_comparison_method(operator.ne)
        cls.__lt__ = cls._create_comparison_method(operator.lt)
        cls.__gt__ = cls._create_comparison_method(operator.gt)
        cls.__le__ = cls._create_comparison_method(operator.le)
        cls.__ge__ = cls._create_comparison_method(operator.ge)

    _create_comparison_method = _create_arithmetic_method

    def __mars_tokenize__(self):
        return [memoryview(x) for chunk in self._arrow_array.chunks
                for x in chunk.buffers()
                if x is not None]


ArrowStringArray._add_arithmetic_ops()
ArrowStringArray._add_comparison_ops()

mars/dataframe/arrays.py
# 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 itertools
import operator
from copy import copy as copy_obj
from numbers import Integral
from typing import Type, Sequence

import numpy as np
import pandas as pd
from pandas._libs import lib
from pandas.api.indexers import check_array_indexer
from pandas.api.types import pandas_dtype, is_scalar, is_array_like
from pandas.api.extensions import ExtensionArray, ExtensionDtype, register_extension_dtype
from pandas.core import ops
from pandas.core.algorithms import take
from pandas.compat import set_function_name
try:
    from pandas.arrays import StringArray as StringArrayBase
except ImportError:  # for pandas < 1.0
    StringArrayBase = ExtensionArray

try:
    import pyarrow as pa
    pa_null = pa.NULL
except ImportError:  # pragma: no cover
    pa = None
    pa_null = None


@register_extension_dtype
class ArrowStringDtype(ExtensionDtype):
    """
    Extension dtype for arrow string data.

    .. warning::

       ArrowStringDtype is considered experimental. The implementation and
       parts of the API may change without warning.

       In particular, ArrowStringDtype.na_value may change to no longer be
       ``numpy.nan``.

    Attributes
    ----------
    None

    Methods
    -------
    None

    Examples
    --------
    >>> import mars.dataframe as md
    >>> md.ArrowStringDtype()
    ArrowStringDtype
    """

    type = str
    kind = "U"
    name = "arrow_string"
    na_value = pa_null

    @classmethod
    def construct_from_string(cls, string):
        if string == cls.name:
            return cls()
        else:
            raise TypeError(f"Cannot construct a '{cls}' from '{string}'")

    @classmethod
    def construct_array_type(cls) -> "Type[ArrowStringArray]":
        return ArrowStringArray

    def __from_arrow__(self, array):
        return ArrowStringArray(array)


class ArrowStringArray(StringArrayBase):
    def __init__(self, values, copy=False):
        if isinstance(values, (pd.Index, pd.Series)):
            # for pandas Index and Series,
            # convert to PandasArray
            values = values.array

        if isinstance(values, type(self)):
            arrow_array = values._arrow_array
        elif isinstance(values, ExtensionArray):
            # if come from pandas object like index,
            # convert to pandas StringArray first,
            # validation will be done in construct
            arrow_array = pa.chunked_array([pa.array(values, from_pandas=True)])
        elif isinstance(values, pa.ChunkedArray):
            arrow_array = values
        elif isinstance(values, pa.StringArray):
            arrow_array = pa.chunked_array([values])
        else:
            arrow_array = pa.chunked_array([pa.array(values, type=pa.string())])

        if copy:
            arrow_array = copy_obj(arrow_array)

        self._arrow_array = arrow_array
        self._dtype = ArrowStringDtype()

        # for test purpose
        self._force_use_pandas = False

    @classmethod
    def from_scalars(cls, values):
        arrow_array = pa.chunked_array(
            [pa.array(np.asarray(values)).cast(pa.string())])
        return cls(arrow_array)

    @classmethod
    def _from_sequence(cls, scalars, dtype=None, copy=False):
        if not hasattr(scalars, 'dtype'):
            scalars = np.asarray(scalars)
        if isinstance(scalars, cls):
            if copy:
                scalars = scalars.copy()
            return scalars
        arrow_array = pa.chunked_array([pa.array(scalars).cast(pa.string())])
        return cls(arrow_array, copy=copy)

    @classmethod
    def _from_sequence_of_strings(cls, strings, dtype=None, copy=False):
        return cls._from_sequence(strings, dtype=dtype, copy=copy)

    @classmethod
    def _from_factorized(cls, values, original):
        return cls(values)

    def __repr__(self):
        return f"{type(self).__name__}({repr(self._arrow_array)})"

    @property
    def dtype(self):
        return self._dtype

    @property
    def nbytes(self) -> int:
        return sum(x.size
                   for chunk in self._arrow_array.chunks
                   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):
            return False
        if slc.step is not None and slc.step != 1:
            return False
        if slc.start is not None and \
                not isinstance(slc.start, Integral):  # pragma: no cover
            return False
        if slc.stop is not None and \
                not isinstance(slc.stop, Integral):  # pragma: no cover
            return False
        return True

    def _values_for_factorize(self):
        arr = self.to_numpy()
        mask = self.isna()
        arr[mask] = -1
        return arr, -1

    def _values_for_argsort(self):
        return self.to_numpy()

    @staticmethod
    def _process_pos(pos, length, is_start):
        if pos is None:
            return 0 if is_start else length
        return pos + length if pos < 0 else pos

    def __getitem__(self, item):
        has_take = hasattr(self._arrow_array, 'take')
        if not self._force_use_pandas and has_take:
            if pd.api.types.is_scalar(item):
                item = item + len(self) if item < 0 else item
                return self._arrow_array.take([item]).to_pandas()[0]
            elif self._can_process_slice_via_arrow(item):
                length = len(self)
                start, stop = item.start, item.stop
                start = self._process_pos(start, length, True)
                stop = self._process_pos(stop, length, False)
                return ArrowStringArray(
                    self._arrow_array.slice(offset=start,
                                            length=stop - start))
            elif hasattr(item, 'dtype') and np.issubdtype(item.dtype, np.bool_):
                return ArrowStringArray(self._arrow_array.filter(
                    pa.array(item, from_pandas=True)))
            elif hasattr(item, 'dtype'):
                length = len(self)
                item = np.where(item < 0, item + length, item)
                return ArrowStringArray(self._arrow_array.take(item))

        array = np.asarray(self._arrow_array.to_pandas())
        return ArrowStringArray(array[item])

    def __setitem__(self, key, value):
        if isinstance(value, (pd.Index, pd.Series)):
            value = value.to_numpy()
        if isinstance(value, type(self)):
            value = value.to_numpy()

        key = check_array_indexer(self, key)
        scalar_key = is_scalar(key)
        scalar_value = is_scalar(value)
        if scalar_key and not scalar_value:
            raise ValueError("setting an array element with a sequence.")

        # validate new items
        if scalar_value:
            if pd.isna(value):
                value = None
            elif not isinstance(value, str):
                raise ValueError(
                    f"Cannot set non-string value '{value}' into a StringArray."
                )
        else:
            if not is_array_like(value):
                value = np.asarray(value, dtype=object)
            if len(value) and not lib.is_string_array(value, skipna=True):
                raise ValueError("Must provide strings.")

        string_array = np.asarray(self._arrow_array.to_pandas())
        string_array[key] = value
        self._arrow_array = pa.chunked_array([pa.array(string_array)])

    def __len__(self):
        return len(self._arrow_array)

    def __array__(self, dtype=None):
        return self.to_numpy(dtype=dtype)

    def to_numpy(self, dtype=None, copy=False, na_value=lib.no_default):
        array = np.asarray(self._arrow_array.to_pandas())
        if copy or na_value is not lib.no_default:
            array = array.copy()
        if na_value is not lib.no_default:
            array[self.isna()] = na_value
        return array

    def fillna(self, value=None, method=None, limit=None):
        chunks = []
        for chunk_array in self._arrow_array.chunks:
            array = chunk_array.to_pandas()
            result_array = array.fillna(value=value, method=method,
                                        limit=limit)
            chunks.append(pa.array(result_array, from_pandas=True))
        return ArrowStringArray(pa.chunked_array(chunks))

    def astype(self, dtype, copy=True):
        dtype = pandas_dtype(dtype)
        if isinstance(dtype, ArrowStringDtype):
            if copy:
                return self.copy()
            return self

        # try to slice 1 record to get the result dtype
        test_array = self._arrow_array.slice(0, 1).to_pandas()
        test_result_array = test_array.astype(dtype).array

        result_array = \
            type(test_result_array)(
                np.full(self.shape, test_result_array.dtype.na_value,
                        dtype=np.asarray(test_result_array).dtype))

        start = 0
        # use chunks to do astype
        for chunk_array in self._arrow_array.chunks:
            result_array[start: start + len(chunk_array)] = \
                chunk_array.to_pandas().astype(dtype).array
            start += len(chunk_array)
        return result_array

    def isna(self):
        if not self._force_use_pandas and hasattr(self._arrow_array, 'is_null'):
            return self._arrow_array.is_null().to_pandas().to_numpy()
        else:
            return pd.isna(self._arrow_array.to_pandas()).to_numpy()

    def take(self, indices, allow_fill=False, fill_value=None):
        if allow_fill is False:
            return ArrowStringArray(self[indices])

        string_array = self._arrow_array.to_pandas().to_numpy()

        replace = False
        if allow_fill and fill_value is None:
            fill_value = self.dtype.na_value
            replace = True

        result = take(string_array, indices, fill_value=fill_value,
                      allow_fill=allow_fill)
        if replace:
            # pyarrow cannot recognize pa.NULL
            result[result == self.dtype.na_value] = None
        return ArrowStringArray(result)

    def copy(self):
        return type(self)(copy_obj(self._arrow_array))

    @classmethod
    def _concat_same_type(
            cls, to_concat: Sequence["ArrowStringArray"]) -> "ArrowStringArray":
        chunks = list(itertools.chain.from_iterable(
            x._arrow_array.chunks for x in to_concat))
        if len(chunks) == 0:
            chunks = [pa.array([], type=pa.string())]
        return cls(pa.chunked_array(chunks))

    def value_counts(self, dropna=False):
        string_array = self._arrow_array.to_pandas()
        return ArrowStringArray(string_array.value_counts(dropna=dropna))

    def any(self, axis=0, out=None):
        return self._arrow_array.to_pandas().any(axis=axis, out=out)

    def all(self, axis=0, out=None):
        return self._arrow_array.to_pandas().all(axis=axis, out=out)

    # Overrride parent because we have different return types.
    @classmethod
    def _create_arithmetic_method(cls, op):
        # Note: this handles both arithmetic and comparison methods.
        def method(self, other):
            is_arithmetic = \
                True if op.__name__ in ops.ARITHMETIC_BINOPS else False

            is_other_array = False
            if not is_scalar(other):
                is_other_array = True
                other = np.asarray(other)

            self_is_na = self.isna()
            other_is_na = pd.isna(other)
            mask = self_is_na | other_is_na

            chunks = []
            mask_chunks = []
            start = 0
            for chunk_array in self._arrow_array.chunks:
                chunk_array = np.asarray(chunk_array.to_pandas())
                end = start + len(chunk_array)
                chunk_mask = mask[start: end]
                chunk_valid = ~chunk_mask

                if is_arithmetic:
                    result = np.empty(chunk_array.shape, dtype=object)
                else:
                    result = np.zeros(chunk_array.shape, dtype=bool)

                chunk_other = other
                if is_other_array:
                    chunk_other = other[start: end]
                    chunk_other = chunk_other[chunk_valid]

                # calculate only for both not None
                result[chunk_valid] = op(chunk_array[chunk_valid],
                                         chunk_other)

                if is_arithmetic:
                    chunks.append(pa.array(result, type=pa.string(),
                                           from_pandas=True))
                else:
                    chunks.append(result)
                    mask_chunks.append(chunk_mask)

            if is_arithmetic:
                return ArrowStringArray(pa.chunked_array(chunks))
            else:
                return pd.arrays.BooleanArray(np.concatenate(chunks),
                                              np.concatenate(mask_chunks))

        return set_function_name(method, "__{}__".format(op.__name__), cls)

    @classmethod
    def _add_arithmetic_ops(cls):
        cls.__add__ = cls._create_arithmetic_method(operator.add)
        cls.__radd__ = cls._create_arithmetic_method(ops.radd)

        cls.__mul__ = cls._create_arithmetic_method(operator.mul)
        cls.__rmul__ = cls._create_arithmetic_method(ops.rmul)

    @classmethod
    def _add_comparison_ops(cls):
        cls.__eq__ = cls._create_comparison_method(operator.eq)
        cls.__ne__ = cls._create_comparison_method(operator.ne)
        cls.__lt__ = cls._create_comparison_method(operator.lt)
        cls.__gt__ = cls._create_comparison_method(operator.gt)
        cls.__le__ = cls._create_comparison_method(operator.le)
        cls.__ge__ = cls._create_comparison_method(operator.ge)

    _create_comparison_method = _create_arithmetic_method

    def __mars_tokenize__(self):
        return [memoryview(x) for chunk in self._arrow_array.chunks
                for x in chunk.buffers()
                if x is not None]


ArrowStringArray._add_arithmetic_ops()
ArrowStringArray._add_comparison_ops()

mars/dataframe/arrays.py
diff --git a/mars/_utils.pyx b/mars/_utils.pyx
index bbdb03644d..92e1e89869 100644
--- a/mars/_utils.pyx
+++ b/mars/_utils.pyx
@@ -199,6 +199,10 @@ cdef list tokenize_pandas_categorical(ob):
     return iterative_tokenize(l)
 
 
+cdef list tokenize_pd_extension_dtype(ob):
+    return iterative_tokenize([ob.name])
+
+
 cdef list tokenize_categories_dtype(ob):
     return iterative_tokenize([ob.categories, ob.ordered])
 
@@ -275,6 +279,8 @@ tokenize_handler.register(pd.arrays.DatetimeArray, tokenize_pandas_time_arrays)
 tokenize_handler.register(pd.arrays.TimedeltaArray, tokenize_pandas_time_arrays)
 tokenize_handler.register(pd.arrays.PeriodArray, tokenize_pandas_time_arrays)
 tokenize_handler.register(pd.arrays.IntervalArray, tokenize_pandas_interval_arrays)
+tokenize_handler.register(pd.api.extensions.ExtensionDtype, tokenize_pd_extension_dtype)
+
 if SATypeEngine is not None:
     tokenize_handler.register(SATypeEngine, tokenize_sqlalchemy_data_type)
 if SASelectable is not None:
diff --git a/mars/dataframe/arrays.py b/mars/dataframe/arrays.py
index 8cf8890bbe..2a5aff03ff 100644
--- 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):
diff --git a/mars/dataframe/base/tests/test_base_execution.py b/mars/dataframe/base/tests/test_base_execution.py
index 28905ed15c..62eadba248 100644
--- 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
index dbaab54014..9d0f5c1639 100644
--- 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())

[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'>
```

