WIP: Bitarray backed Int EAs

pandas/core/arrays/integer.py

@@ -2,6 +2,7 @@
 import warnings
 import copy
 import numpy as np
+from bitarray import bitarray
 
 from pandas._libs.lib import infer_dtype
 from pandas.util._decorators import cache_readonly
@@ -167,10 +168,12 @@ def coerce_to_array(values, dtype, mask=None, copy=False):
     else:
         assert len(mask) == len(values)
 
+    # Work with bitarrays from here on out
+    if isinstance(mask, np.ndarray):
+        mask = _numpy_to_bitarray(mask)
+
     if not values.ndim == 1:
         raise TypeError("values must be a 1D list-like")
-    if not mask.ndim == 1:
-        raise TypeError("mask must be a 1D list-like")
 
     # infer dtype if needed
     if dtype is None:
@@ -195,6 +198,23 @@ def coerce_to_array(values, dtype, mask=None, copy=False):
     return values, mask
 
 
+def _numpy_to_bitarray(arr):
+    """
+    Efficiently convert a NumPy array to a bitarray object.
+    """
+    barr = bitarray()
+    barr.pack(arr.astype(bool, copy=False).tostring())
+
+    return barr
+
+
+def _bitarray_to_numpy(arr):
+    """
+    Efficiently convert a bitarray object to a NumPy array.
+    """
+    return np.fromstring(arr.unpack(), dtype=bool)
+
+
 class IntegerArray(ExtensionArray, ExtensionOpsMixin):
     """
     We represent an IntegerArray with 2 numpy arrays
@@ -235,8 +255,9 @@ def __getitem__(self, item):
             if self._mask[item]:
                 return self.dtype.na_value
             return self._data[item]
+
         return type(self)(self._data[item],
-                          mask=self._mask[item],
+                          mask=_bitarray_to_numpy(self._mask)[item],
                           dtype=self.dtype)
 
     def _coerce_to_ndarray(self):
@@ -317,7 +338,10 @@ def __setitem__(self, key, value):
             mask = mask[0]
 
         self._data[key] = value
-        self._mask[key] = mask
+        # Coerce to numpy array to leverage advanced indexing, then coerce back
+        arr = _bitarray_to_numpy(self._mask)
+        arr[key] = mask
+        self._mask = _numpy_to_bitarray(arr)
 
     def __len__(self):
         return len(self._data)
@@ -343,10 +367,10 @@ def __repr__(self):
 
     @property
     def nbytes(self):
-        return self._data.nbytes + self._mask.nbytes
+        return self._data.nbytes + self._mask.buffer_info()[1]
 
     def isna(self):
-        return self._mask
+        return _bitarray_to_numpy(self._mask)
 
     @property
     def _na_value(self):
@@ -444,7 +468,7 @@ def value_counts(self, dropna=True):
             # TODO(extension)
             # appending to an Index *always* infers
             # w/o passing the dtype
-            array = np.append(array, [self._mask.sum()])
+            array = np.append(array, [_bitarray_to_numpy(self._mask).sum()])
             index = Index(np.concatenate(
                 [index.values,
                  np.array([np.nan], dtype=object)]), dtype=object)
@@ -513,7 +537,8 @@ def _maybe_mask_result(self, result, mask, other, op_name):
         # may need to fill infs
         # and mask wraparound
         if is_float_dtype(result):
-            mask |= (result == np.inf) | (result == -np.inf)
+            arr = _numpy_to_bitarray((result == np.inf) | (result == -np.inf))
+            mask |= arr
 
         # if we have a float operand we are by-definition
         # a float result

pandas/core/series.py

@@ -988,7 +988,10 @@ def _set_with(self, key, value):
                 else:
                     return self._set_values(key, value)
             elif key_type == 'boolean':
-                self._set_values(key.astype(np.bool_), value)
+                try:
+                    self._set_values(key.astype(np.bool_), value)
+                except AttributeError:
+                    self._set_values(key, value)
             else:
                 self._set_labels(key, value)
 

pandas/util/testing.py

@@ -15,6 +15,7 @@
 from functools import wraps
 from contextlib import contextmanager
 
+from bitarray import bitarray
 from numpy.random import randn, rand
 import numpy as np
 
@@ -1172,6 +1173,13 @@ def assert_extension_array_equal(left, right):
     assert left.dtype == right.dtype
     left_na = left.isna()
     right_na = right.isna()
+
+    # TODO - maybe generate dedicated method for bitarray comparison?
+    if isinstance(left_na, bitarray):
+        left_na = np.fromstring(left_na.unpack(), dtype=bool)
+    if isinstance(right_na, bitarray):
+        right_na = np.fromstring(right_na.unpack(), dtype=bool)
+
     assert_numpy_array_equal(left_na, right_na)
 
     left_valid = left[~left_na].astype(object)