ENJH: create NumpyBoolArray and ArrawBoolArray subclasses for MaskArray

Author: jreback

pandas/core/arrays/_mask.py

@@ -14,16 +14,13 @@
     is_bool_dtype, is_float, is_float_dtype, is_integer, is_integer_dtype,
     is_list_like, is_object_dtype, is_scalar)
 from pandas.core.dtypes.dtypes import register_extension_dtype
-from pandas.core.dtypes.generic import ABCIndexClass, ABCSeries
+from pandas.core.dtypes.generic import ABCIndexClass, ABCMaskArray, ABCSeries
 from pandas.core.dtypes.missing import isna, notna
 
 from pandas.core import nanops
 from pandas.core.arrays import ExtensionArray, ExtensionOpsMixin
-<<<<<<< HEAD
+from pandas.core.arrays.mask import get_mask_array_type
 from pandas.core.tools.numeric import to_numeric
-=======
-from pandas.core.arrays._mask import NAMask
->>>>>>> 384287e71... Revert unnecessary changes from master
 
 
 class _IntegerDtype(ExtensionDtype):
@@ -291,8 +288,8 @@ def __init__(self, values, mask, copy=False):
                 and is_integer_dtype(values.dtype)):
             raise TypeError("values should be integer numpy array. Use "
                             "the 'integer_array' function instead")
-        if not (isinstance(mask, NAMask) or (
-                isinstance(mask, np.ndarray) and is_bool_dtype(mask.dtype))):
+        if not (isinstance(mask, (np.ndarray, ABCMaskArray)) and
+                is_bool_dtype(mask.dtype)):
             raise TypeError("mask should be boolean numpy array. Use "
                             "the 'integer_array' function instead")
 
@@ -301,7 +298,7 @@ def __init__(self, values, mask, copy=False):
             mask = mask.copy()
 
         self._data = values
-        self._mask = NAMask(mask)
+        self._mask = get_mask_array_type()._from_sequence(mask, copy=False)
 
     @classmethod
     def _from_sequence(cls, scalars, dtype=None, copy=False):
@@ -328,22 +325,7 @@ def __getitem__(self, item):
             if self._mask[item]:
                 return self.dtype.na_value
             return self._data[item]
-<<<<<<< HEAD
         return type(self)(self._data[item], self._mask[item])
-=======
-
-<<<<<<< HEAD
-        return type(self)(self._data[item],
-<<<<<<< HEAD
-                          mask=_bitarray_to_numpy(self._mask)[item],
-                          dtype=self.dtype)
->>>>>>> 2ff4b0907... First pass at implementation (needs refactor)
-=======
-                          mask=_bitarray_to_numpy(self._mask)[item])
->>>>>>> e085674ac... Reverted changes; created new module for mask
-=======
-        return type(self)(self._data[item], self._mask[item])
->>>>>>> 384287e71... Revert unnecessary changes from master
 
     def _coerce_to_ndarray(self):
         """
@@ -352,7 +334,8 @@ def _coerce_to_ndarray(self):
 
         # TODO(jreback) make this better
         data = self._data.astype(object)
-        data[self._mask] = self._na_value
+        mask = np.array(self._mask, copy=False)
+        data[mask] = self._na_value
         return data
 
     __array_priority__ = 1000  # higher than ndarray so ops dispatch to us
@@ -581,6 +564,7 @@ def cmp_method(self, other):
             else:
                 mask = self._mask | mask
 
+            mask = np.array(mask, copy=False)
             result[mask] = op_name == 'ne'
             return result
 

pandas/core/arrays/integer.py

@@ -0,0 +1,30 @@
+_MaskArrayType = None
+
+
+def get_mask_array_type():
+    """Set the mask array type to use, we need to do
+    this after all modules are imported as the implementations
+    e.g. pyarrow depend on pandas being importable
+    """
+    global _MaskArrayType
+
+    if _MaskArrayType is not None:
+        return _MaskArrayType
+
+    # if ArrowBoolArray is available use it
+    # otherwise use the NumpyMask
+    try:
+        from pandas.core.arrays.mask._pyarrow import ArrowBoolArray
+
+        MaskArray = ArrowBoolArray
+
+    except ImportError:
+        from pandas.core.arrays.mask._numpy import NumpyBoolArray
+
+        MaskArray = NumpyBoolArray
+
+    _MaskArrayType = MaskArray
+    return _MaskArrayType
+
+
+__all__ = ['get_mask_array_type']

pandas/core/arrays/mask/__init__.py

@@ -0,0 +1,128 @@
+"""A boolean mask interace.
+
+This module provides an interface to a numpy / pyarrow boolean mask.
+This is limited as not all of the implementations can hold NA, so
+for consistency this is an internal.
+"""
+
+import copy
+
+import numpy as np
+
+from pandas import compat
+from pandas.api.extensions import ExtensionDtype
+from pandas.core.arrays.base import ExtensionArray
+from pandas.core.missing import isna
+
+
+class BoolDtype(ExtensionDtype):
+
+    type = np.bool_
+    kind = 'b'
+    name = 'bool'
+
+    @classmethod
+    def construct_from_string(cls, string):
+        if string == cls.name:
+            return cls()
+        else:
+            raise TypeError("Cannot construct a '{}' from "
+                            "'{}'".format(cls, string))
+
+    def _is_boolean(self):
+        return True
+
+    def __hash__(self):
+        return hash(str(self))
+
+    def __eq__(self, other):
+        # compare == to np.dtype('bool')
+        if isinstance(other, compat.string_types):
+            return other == self.name
+        elif other is self:
+            return True
+        elif isinstance(other, np.dtype):
+            return other == 'bool'
+        else:
+            return hash(self) == hash(other)
+
+
+class BoolArray(ExtensionArray):
+    """Common baseclass for both pyarrow and numpy masked arrays"""
+    _typ = "maskarray"
+
+    @classmethod
+    def _from_sequence(cls, scalars, dtype=None, copy=False):
+        return cls.from_scalars(scalars)
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    @property
+    def size(self):
+        return len(self)
+
+    def __eq__(self, other):
+        return np.array(self, copy=False) == np.array(other, copy=False)
+
+    def __len__(self):
+        return len(self._data)
+
+    def isna(self):
+        nas = isna(np.array(self._data, copy=False))
+        return type(self).from_scalars(nas)
+
+    def __invert__(self):
+        return type(self).from_scalars(
+            ~np.array(self._data, copy=False)
+        )
+
+    def __or__(self, other):
+        return type(self).from_scalars(np.array(
+            self, copy=False).__or__(np.array(other, copy=False)))
+
+    def __ior__(self, other):
+        return type(self).from_scalars(
+            np.array(self, copy=False) | np.array(other, copy=False))
+
+    def __and__(self, other):
+        return type(self).from_scalars(
+            np.array(self, copy=False).__and__(np.array(other, copy=False)))
+
+    def __iand__(self, other):
+        return type(self).from_scalars(
+            np.array(self, copy=False) & (np.array(other, copy=False)))
+
+    def view(self, dtype=None):
+        arr = np.array(self._data, copy=False)
+        if dtype is not None:
+            arr = arr.view(dtype=dtype)
+        return arr
+
+    def sum(self, axis=None):
+        return np.array(self, copy=False).sum()
+
+    def copy(self, deep=False):
+        if deep:
+            return type(self)(copy.deepcopy(self._data))
+        else:
+            return type(self)(copy.copy(self._data))
+
+    def any(self, axis=0, out=None):
+        return np.array(self._data, copy=False).any()
+
+    def all(self, axis=0, out=None):
+        return np.array(self._data, copy=False).all()
+
+    def _reduce(self, method, skipna=True, **kwargs):
+        if skipna:
+            arr = self[~self.isna()]
+        else:
+            arr = self
+
+        try:
+            op = getattr(arr, method)
+        except AttributeError:
+            raise TypeError
+        return op(**kwargs)