create NumpyBoolArray and ArrayBoolArrayk subclasses for MaskArray

Author: jreback

pandas/core/arrays/_mask.py

@@ -1,103 +1,136 @@
+"""A boolean mask interace.
+
+This module provide a numpy-like mask; it can be optionally backed
+by pyarrow.
+"""
+
 import numpy as np
+from pandas.core.base import PandasObject
 
 
-class NAMask():
+# TODO: The NumpyBoolArray should actually be a full EA.
+
+class NumpyBoolArray(PandasObject):
     """Generic class which can be used to represent missing data.
+    """
+    _typ = "maskarray"
 
-    Will use bitarray if available; otherwise will use numpy."""
+    @property
+    def dtype(self):
+        return np.dtype('bool')
 
-    def __init__(self, mask):
+    def __init__(self, mask, copy=True):
         """
         Parameters
         ----------
         mask : numpy array
             Mask of missing values.
         """
-
-        self._has_bitarray = False
-        try:
-            import bitarray
-            globals()['bitarray'] = bitarray
-            self._has_bitarray = True
-            self._data = self._numpy_to_bitarray(mask)
-        except (ImportError, ModuleNotFoundError):
-            self._data = mask.astype(bool, copy=False)
-
-    def _numpy_to_bitarray(self, arr):
-        bit_arr = bitarray()
-        bit_arr.pack(arr.astype(bool, copy=False))
-
-    def _bitarray_to_numpy(self, arr):
-        return np.fromstring(arr.unpack(), dtype=bool)
-
-    def __getitem__(self, item):
-        if self._has_bitarray:
-            raise NotImplementedError
-
-        return self._data[item]
+        mask = np.asarray(mask)
+        self._data = mask.astype(bool, copy=copy)
+
+    def __repr__(self):
+        from pandas.io.formats.printing import format_object_summary
+
+        template = (
+            u'{class_name}'
+            u'{data}\n'
+            u'Length: {length}'
+        )
+        # the short repr has no trailing newline, while the truncated
+        # repr does. So we include a newline in our template, and strip
+        # any trailing newlines from format_object_summary
+        data = format_object_summary(self, self._formatter(),
+                                     indent_for_name=False).rstrip(', \n')
+        class_name = u'<{}>\n'.format(self.__class__.__name__)
+        return template.format(class_name=class_name, data=data,
+                               length=len(self))
+
+    def _formatter(self, boxed=False):
+        return repr
+
+    def __getitem__(self, key):
+        return self._data[key]
 
     def __setitem__(self, key, value):
-        if self._has_bitarray:
-            raise NotImplementedError
-
         self._data[key] = value
 
-    def __array__(self):
-        if self._has_bitarray:
-            raise NotImplementedError
+    def __len__(self):
+        return len(self._data)
 
+    def __eq__(self, other):
+        return np.array(self, copy=False) == np.array(other, copy=False)
+
+    def __array__(self, dtype=None):
         return self._data
 
     def __iter__(self):
         for i in range(len(self._data)):
             yield self._data[i]
 
     def __invert__(self):
-        if self._has_bitarray:
-            raise NotImplementedError
-
         return type(self)(~self._data)
 
     def __or__(self, other):
-        if self._has_bitarray:
-            raise NotImplementedError
-
-        return self._data.__or__(other)
+        return np.array(
+            self, copy=False).__or__(np.array(other, copy=False))
 
     def __ior__(self, other):
-        if self._has_bitarray:
-            raise NotImplementedError
+        return np.array(self, copy=False) | np.array(other, copy=False)
 
-        return self._data | other
+    def __and__(self, other):
+        return np.array(self, copy=False).__and__(np.array(other, copy=False))
+
+    def __iand__(self, other):
+        return np.array(self, copy=False) & (np.array(other, copy=False))
 
     @property
     def nbytes(self):
-        if self._has_bitarray:
-            return self._data.buffer_info()[1]
-
         return self._data.nbytes
 
     @property
     def size(self):
-        if self._has_bitarray:
-            raise NotImplementedError
+        return len(self)
+
+    @property
+    def ndim(self):
+        return 1
 
-        return self._data.size        
+    @property
+    def shape(self):
+        return (len(self), )
 
-    def astype(self, dtype, copy=False):
-        if self._has_bitarray:
-            raise NotImplementedError
+    def reshape(self, shape, **kwargs):
+        return np.array(self, copy=False).reshape(shape, **kwargs)
 
-        return self._data.astype(dtype, copy=copy)
+    def astype(self, dtype, copy=False):
+        return np.array(self, copy=False).astype(dtype, copy=copy)
 
     def any(self):
         return self._data.any()
 
+    def all(self):
+        return self._data.all()
+
     def copy(self):
         return type(self)(self._data.copy())
 
-    def sum(self):
-        if self._has_bitarray:
-            raise NotImplementedError
+    def sum(self, axis=None):
+        return np.array(self, copy=False).sum()
+
+    def take(self, indicies, **kwargs):
+        return np.array(self, copy=False).take(indicies)
+
+
+def get_mask_array_type():
+    # protect imports
+
+    # if ArrowBoolArray is available use it
+    # otherwise use the NumpyMask
+    try:
+        from pandas.core.arrays.bool import ArrowBoolArray
+        MaskArray = ArrowBoolArray
+    except ImportError:
+        MaskArray = NumpyBoolArray
 
-        return self._data.sum()
+    return MaskArray

pandas/tests/extension/arrow/bool.py renamed to pandas/core/arrays/bool.py

@@ -12,6 +12,7 @@
 import pyarrow as pa
 
 import pandas as pd
+from pandas.api.types import is_scalar
 from pandas.api.extensions import (
     ExtensionArray, ExtensionDtype, register_extension_dtype, take)
 
@@ -41,17 +42,20 @@ def _is_boolean(self):
 
 
 class ArrowBoolArray(ExtensionArray):
-    def __init__(self, values):
+    def __init__(self, values, copy=False):
         if not isinstance(values, pa.ChunkedArray):
-            raise ValueError
+
+            # TODO: hack a minute
+            values = np.asarray(values).astype(np.bool_, copy=False)
+            values = pa.chunked_array([values])
+        else:
+            if copy:
+                values = values.copy()
 
         assert values.type == pa.bool_()
         self._data = values
         self._dtype = ArrowBoolDtype()
 
-    def __repr__(self):
-        return "ArrowBoolArray({})".format(repr(self._data))
-
     @classmethod
     def from_scalars(cls, values):
         arr = pa.chunked_array([pa.array(np.asarray(values))])
@@ -67,15 +71,24 @@ def _from_sequence(cls, scalars, dtype=None, copy=False):
         return cls.from_scalars(scalars)
 
     def __getitem__(self, item):
-        if pd.api.types.is_scalar(item):
+        if is_scalar(item):
             return self._data.to_pandas()[item]
         else:
             vals = self._data.to_pandas()[item]
             return type(self).from_scalars(vals)
 
+    def __setitem__(self, key, value):
+        # TODO: hack-a-minute
+        data = np.array(self._data)
+        data[key] = value
+        self._data = pa.array(data)
+
     def __len__(self):
         return len(self._data)
 
+    def view(self, dtype):
+        return np.array(self._data, copy=False).view(dtype)
+
     def astype(self, dtype, copy=True):
         # needed to fix this astype for the Series constructor.
         if isinstance(dtype, type(self.dtype)) and dtype == self.dtype:
@@ -94,11 +107,36 @@ def nbytes(self):
                    for x in chunk.buffers()
                    if x is not None)
 
+    @property
+    def size(self):
+        return len(self)
+
+    def __repr__(self):
+        from pandas.io.formats.printing import format_object_summary
+
+        template = (
+            u'{class_name}'
+            u'{data}\n'
+            u'Length: {length}'
+        )
+        # the short repr has no trailing newline, while the truncated
+        # repr does. So we include a newline in our template, and strip
+        # any trailing newlines from format_object_summary
+        data = format_object_summary(self, self._formatter(),
+                                     indent_for_name=False).rstrip(', \n')
+        class_name = u'<{}>\n'.format(self.__class__.__name__)
+        return template.format(class_name=class_name, data=data,
+                               length=len(self))
+
+    def __eq__(self, other):
+        return np.array(self, copy=False) == np.array(other, copy=False)
+
     def isna(self):
         nas = pd.isna(self._data.to_pandas())
         return type(self).from_scalars(nas)
 
-    def take(self, indices, allow_fill=False, fill_value=None):
+    def take(self, indices, allow_fill=False, fill_value=None, axis=None):
+        # TODO: had to add axis here
         data = self._data.to_pandas()
 
         if allow_fill and fill_value is None:
@@ -114,6 +152,9 @@ def copy(self, deep=False):
         else:
             return type(self)(copy.copy(self._data))
 
+    def sum(self, axis=None):
+        return np.array(self, copy=False).sum()
+
     def _concat_same_type(cls, to_concat):
         chunks = list(itertools.chain.from_iterable(x._data.chunks
                                                     for x in to_concat))
@@ -125,6 +166,22 @@ def __invert__(self):
             ~self._data.to_pandas()
         )
 
+    def __or__(self, other):
+        return np.array(
+            self, copy=False).__or__(np.array(other, copy=False))
+
+    def __ior__(self, other):
+        return np.array(self, copy=False) | np.array(other, copy=False)
+
+    def __and__(self, other):
+        return np.array(self, copy=False).__and__(np.array(other, copy=False))
+
+    def __iand__(self, other):
+        return np.array(self, copy=False) & (np.array(other, copy=False))
+
+    def __array__(self, dtype=None):
+        return np.array(self._data, copy=False)
+
     def _reduce(self, method, skipna=True, **kwargs):
         if skipna:
             arr = self[~self.isna()]

pandas/core/arrays/integer.py

@@ -19,11 +19,8 @@
 
 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):
@@ -287,12 +284,14 @@ def dtype(self):
         return _dtypes[str(self._data.dtype)]
 
     def __init__(self, values, mask, copy=False):
+        MaskArray = get_mask_array_type()
+
         if not (isinstance(values, np.ndarray)
                 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, MaskArray)) and
+                is_bool_dtype(mask.dtype)):
             raise TypeError("mask should be boolean numpy array. Use "
                             "the 'integer_array' function instead")
 
@@ -301,7 +300,7 @@ def __init__(self, values, mask, copy=False):
             mask = mask.copy()
 
         self._data = values
-        self._mask = NAMask(mask)
+        self._mask = MaskArray(mask, copy=False)
 
     @classmethod
     def _from_sequence(cls, scalars, dtype=None, copy=False):
@@ -328,22 +327,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):
         """