Bitmask Backed MaskedArray
State: closed

POC using vendored nanoarrow



Labels:
name: Build
description: Library building on various platforms
name: NA - MaskedArrays
description: Related to pd.NA and nullable extension arrays
name: Arrow
description: pyarrow functionality

Comments:
Benchmarks

```sh
$ asv continuous upstream/main HEAD
      before           after         ratio
     [4f1a086f]       [173b4cbe]
     <main>           <bitmask-backed>
+        1.23±0μs         66.8±1μs    54.35  array.IntegerArray.time_constructor
+        707±10ns      1.41±0.04μs     2.00  array.BooleanArray.time_constructor
+         957±4μs      1.70±0.02ms     1.78  frame_ctor.FromScalar.time_frame_from_scalar_ea_float64
+         978±6μs      1.68±0.01ms     1.71  frame_ctor.FromScalar.time_frame_from_scalar_ea_float64_na
+      8.72±0.2μs      11.9±0.09μs     1.37  boolean.TimeLogicalOps.time_xor_scalar
+      8.96±0.2μs       11.9±0.3μs     1.33  boolean.TimeLogicalOps.time_or_array
+      8.60±0.3μs       11.4±0.2μs     1.33  boolean.TimeLogicalOps.time_and_array
+     2.19±0.03ms       2.92±0.5ms     1.33  multiindex_object.SetOperations.time_operation('monotonic', 'ea_int', 'intersection', False)
+      13.7±0.2ms       17.7±0.1ms     1.29  index_object.SetOperations.time_operation('monotonic', 'date_string', 'intersection')
+     1.26±0.07ms      1.53±0.02ms     1.21  index_object.SetOperations.time_operation('non_monotonic', 'ea_int', 'intersection')
+        297±30μs         368±30μs     1.24  series_methods.NanOps.time_func('sum', 1000000, 'int8')
+     1.24±0.01ms      1.48±0.03ms     1.19  index_object.SetOperations.time_operation('non_monotonic', 'ea_int', 'intersection')
+     2.47±0.02ms      2.92±0.03ms     1.18  reshape.ReshapeMaskedArrayDtype.time_transpose('Int64')
+      45.2±0.3μs       52.5±0.6μs     1.16  algos.isin.IsinAlmostFullWithRandomInt.time_isin(<class 'numpy.uint64'>, 12, 'outside')
+     2.51±0.01ms      2.89±0.05ms     1.15  reshape.ReshapeMaskedArrayDtype.time_transpose('Float64')
+        76.8±1μs         88.0±2μs     1.15  algos.isin.IsinWithRandomFloat.time_isin(<class 'numpy.float64'>, 2000, 'outside')
+        225±30μs          257±4μs     1.14  algos.isin.IsIn.time_isin_categorical('string[python]')
+         195±4ns          223±5ns     1.14  indexing_engines.MaskedNumericEngineIndexing.time_get_loc((<class 'pandas._libs.index.MaskedFloat64Engine'>, 'Float64'), 'monotonic_decr', True, 2000000)
+      8.03±0.2ms      9.13±0.08ms     1.14  frame_methods.Equals.time_frame_nonunique_unequal
+     2.05±0.03ms      2.33±0.04ms     1.14  index_object.SetOperations.time_operation('non_monotonic', 'ea_int', 'union')
+         478±3μs          539±1μs     1.13  series_methods.NanOps.time_func('min', 1000000, 'boolean')
+        95.3±1μs          108±2μs     1.13  algos.isin.IsinAlmostFullWithRandomInt.time_isin(<class 'numpy.uint64'>, 13, 'outside')
+        46.2±1μs       51.8±0.5μs     1.12  algos.isin.IsinAlmostFullWithRandomInt.time_isin(<class 'numpy.int64'>, 12, 'outside')
+      11.0±0.5μs       12.3±0.9μs     1.12  algorithms.SortIntegerArray.time_argsort(1000)
+       860±100ns        963±100ns     1.12  index_cached_properties.IndexCache.time_shape('TimedeltaIndex')
+         463±5μs          515±5μs     1.11  series_methods.NanOps.time_func('max', 1000000, 'boolean')
+        56.3±2μs       62.5±0.8μs     1.11  algos.isin.IsinWithArangeSorted.time_isin(<class 'numpy.float64'>, 2000)
+      8.60±0.2ms       9.51±0.2ms     1.11  frame_methods.Equals.time_frame_nonunique_equal
+     1.12±0.01ms      1.25±0.03ms     1.11  series_methods.Fillna.time_bfill('Int64')
+     2.07±0.02ms       2.30±0.1ms     1.11  multiindex_object.SetOperations.time_operation('non_monotonic', 'datetime', 'union', False)
+         196±1ns          217±3ns     1.11  indexing_engines.NumericEngineIndexing.time_get_loc_near_middle((<class 'pandas._libs.index.Float32Engine'>, <class 'numpy.float32'>), 'monotonic_decr', True, 2000000)
+        31.6±1μs       34.8±0.2μs     1.10  algos.isin.IsIn.time_isin_mismatched_dtype('Int64')
+     1.15±0.02ms      1.26±0.01ms     1.10  series_methods.Fillna.time_bfill('Float64')
+      4.13±0.05s       4.55±0.02s     1.10  stat_ops.FrameOps.time_op('kurt', 'Int64', 1)
+     2.71±0.03ms      2.98±0.02ms     1.10  indexing_engines.MaskedNumericEngineIndexing.time_get_loc((<class 'pandas._libs.index.MaskedInt16Engine'>, 'Int16'), 'monotonic_incr', True, 2000000)
-     2.45±0.02ms      2.22±0.01ms     0.91  stat_ops.FrameOps.time_op('sem', 'Int64', None)
-      13.5±0.6ms       12.3±0.1ms     0.91  categoricals.Constructor.time_regular
-     2.43±0.02ms      2.18±0.01ms     0.90  stat_ops.FrameOps.time_op('sem', 'Int64', 0)
-      16.3±0.6μs       14.8±0.2μs     0.90  series_methods.NanOps.time_func('std', 1000, 'int32')
-      2.82±0.4ms      2.52±0.05ms     0.89  series_methods.NanOps.time_func('var', 1000000, 'int32')
-      9.94±0.4ms       8.81±0.2ms     0.89  series_methods.NanOps.time_func('sem', 1000000, 'float64')
-      93.7±0.8μs       83.7±0.4μs     0.89  join_merge.ConcatIndexDtype.time_concat_series('Int64', 'non_monotonic', 0, True)
-      93.8±0.5μs       83.6±0.7μs     0.89  join_merge.ConcatIndexDtype.time_concat_series('Int64', 'non_monotonic', 0, False)
-      1.22±0.3μs       1.08±0.1μs     0.88  index_cached_properties.IndexCache.time_values('CategoricalIndex')
-      11.9±0.3ms      10.1±0.08ms     0.84  frame_methods.Fillna.time_fillna(False, 'Float64')
-         520±50μs          433±3μs     0.83  series_methods.NanOps.time_func('sum', 1000000, 'float64')
-        728±50μs          596±7μs     0.82  series_methods.NanOps.time_func('sum', 1000000, 'Int64')
-      4.90±0.3ms       4.03±0.2ms     0.82  indexing.MultiIndexing.time_loc_slice_plus_null_slice(False)
-      12.1±0.4ms      9.92±0.07ms     0.82  frame_methods.Fillna.time_fillna(False, 'Int64')
-         374±3μs          307±2μs     0.82  timeseries.DatetimeIndex.time_unique('repeated')
-     1.21±0.01ms          981±8μs     0.81  frame_methods.Fillna.time_bfill(True, 'Int64')
-        815±70ns         661±60ns     0.81  index_cached_properties.IndexCache.time_values('TimedeltaIndex')
-     1.20±0.01ms          972±7μs     0.81  frame_methods.Fillna.time_bfill(True, 'Float64')
-     11.6±0.03ms      9.01±0.08ms     0.78  indexing.InsertColumns.time_insert_middle
-      3.68±0.3ms      2.85±0.05ms     0.77  series_methods.NanOps.time_func('var', 1000000, 'int64')
-     3.91±0.02μs      3.03±0.02μs     0.78  series_methods.Any.time_any(1000, 'slow', 'boolean')
-     3.99±0.06μs      3.06±0.04μs     0.77  series_methods.All.time_all(1000, 'fast', 'boolean')
-            458M             349M     0.76  array.BooleanArrayMem.peakmem_array
-     4.74±0.09ms      3.28±0.03ms     0.69  series_methods.Fillna.time_fillna('Float64')
-     4.73±0.03ms      3.19±0.01ms     0.67  series_methods.Fillna.time_fillna('Int64')
-         324±6μs       29.9±0.4μs     0.09  series_methods.All.time_all(1000000, 'slow', 'boolean')
-         309±3μs       10.8±0.1μs     0.03  series_methods.All.time_all(1000000, 'fast', 'boolean')

SOME BENCHMARKS HAVE CHANGED SIGNIFICANTLY.
PERFORMANCE DECREASED.
```

Updated benchmarks post. The direct constructor is definitely slower, though not sure how much that can be optimized beyond what is there now (cython shows only python interaction on np.ravel and instance checks)

Indexing showed a lot of small improvements? Wondering if those are just noise. Groupby generally shows a slow-down now too I hope to still bridge
Any idea how this affects build sizes?

I've tagged this as build to see how the wheel builders react.
Updated benchmarks. Still closer, but there are definitely ~10-20% regressions on some groupby / indexing operations. The main thing is the time to convert a bitmask into a NumPy array.

Taking the groupby code as an example, the methods there rely on numpy semantics, so `to_numpy()` is a must for each column. Calling that for 1 million boolean bits takes about ~500us. The groupby benchmarks that show a regression typically have 10 columns, so that 10 col x 500us yields a 5ms difference in the benchmarks, which is enough to trigger the regression.
> Any idea how this affects build sizes?

The size of the arrays shared library increases from 136K to 300K on a release build for me
For our groupby code, _in theory_ we could pass the bitmask and have a branch there that checks this instead of the mask? (like we currently have the `uses_mask` pattern in groupby.pyx)
That could be an option. The thing I think preventing this though is some of the Python level processing before we get into the Cython calls like this:

https://github.com/pandas-dev/pandas/blob/2c4c072ade78b96a9eb05097a5fcf4347a3768f3/pandas/core/groupby/ops.py#L317
Actually https://github.com/pandas-dev/pandas/pull/54506/commits/e8b7819d41582b5228361548afbaafcfe2dd7165 might have been the trick. With that groupby barely shows a regression. Going to just let the entire test suite run tonight and post updated results, but feeling hopeful after that change
OK this is looking a bit better. Updated benchmarks show a limited amount of regressions, and there are now some cases where leveraging the bitmask can actually help performance a lot

Hoping to figure out what's going on with the reshape regressions then this is pretty close to ready. Happy to address any review comments in the meantime
Thanks a lot for exploring this, Will!

Didn't look in detail at the code changes, but some quick general questions:

- Does this support "2d" masks? Or if not, do you think that would be not too complicated to add later? 
- With pyarrow becoming a default dependency, some of the things that you implemented here in cython, _might_ potentially also be done using a pyarrow.BooleanArray (without nulls) and compute functions on that? (although the previous question about 2D would definitely complicate that)
Yes this supports 2D masks. Essentially the data is stored in a C-contiguous fashion, and the storage of the `array_shape` along with any `reshape` calls within the methods allow you to reconstruct a 2D array.  That seems to cover the majority of existing tests, though it may need to become more robust over time
I think the remaining bottlenecks can be grouped into two major categories:
  1. ~~Things that flow through nanops, which use numpy functions that requiere the conversion from bitmask -> ndarray~~
  2. Reshape / concat operations that require the mask itself to be concatenated. I don't think there is a high performance option for concating bitmask arrays (unless @paleolimbot knows) so would need to implement that here or upstream

(edit: 1. was helped out by a more efficient invert operator)
> I don't think there is a high performance option for concating bitmask arrays

nanoarrow definitely doesn't have that...off the top of my head, the best option might be to check for `all(lengths % 8)`: if true, just concatenate the bytes; if false, loop over fixed size chunks of bytes (or int32, which can be faster) and unpack/repack. In nanoarrow that could be `ArrowBitmapAppendBits()`?

> some of the things that you implemented here in cython, might potentially also be done using a pyarrow.BooleanArray (without nulls) and compute functions on that?

Arrow C++ definitely dedicates a lot of space to SIMD-optimized packing/unpacking although not sure the packing would work with the mutability you've implemented here, nor am I sure the extent to which the packing/unpacking is optimized in the compute kernels. At least for 32-bit unpacking, the unpacking you added to nanoarrow outperforms a cast from boolean to int32 by 8-12x.

<details>

``` r
# pak::pak("apache/arrow-nanoarrow/r")
library(nanoarrow)

lgls <- nanoarrow:::vec_gen(logical(), 1e6)
bool_array <- as_nanoarrow_array(lgls)
bool_array_arrow <- arrow::as_arrow_array(bool_array)

bench::mark(
  convert_array(bool_array, integer()),
  as.vector(bool_array_arrow$cast(arrow::int32())),
  relative = TRUE
)
#> # A tibble: 2 × 6
#>   expression                             min median `itr/sec` mem_alloc `gc/sec`
#>   <bch:expr>                           <dbl>  <dbl>     <dbl>     <dbl>    <dbl>
#> 1 convert_array(bool_array, integer())   1     1         7.99      6.49     48.2
#> 2 as.vector(bool_array_arrow$cast(arr…  23.0   8.57      1         1         1
```

<sup>Created on 2023-08-21 with [reprex v2.0.2](https://reprex.tidyverse.org)</sup>

</details>
Yea nice ideas. I think another option aside from unpack/repack is to let this class manage multiple underlying bitmaps and handle dispatching to each. 
OK this is mostly green now. I also added a benchmark to confirm the memory savings. The ratio gets better as you increase the number of objects in the test

```sh
       before           after         ratio
     [9268055a]       [34ff6988]
     <main>           <bitmask-backed>
-            457M             348M     0.76  array.BooleanArrayMem.peakmem_array

SOME BENCHMARKS HAVE CHANGED SIGNIFICANTLY.
PERFORMANCE INCREASED.
```

Overall the remaining benchmarks show more improvements than regressions:

https://github.com/pandas-dev/pandas/pull/54506#issuecomment-1676514037

For some of the larger regressions I think these can still be improved. Some notes:

1. The largest regression is in the IntegerArray constructor, though that benchmark isn't really apples-to-apples. Previously mask in the setup function was in its "final form", and the constructor just does an assignment. Now the constructor has to convert the ndarray to a BitMaskArray; we could alternately change the setup function to use a BitMaskArray to be equivalent, but kept as the worst case scenario for now
2. frame_ctor.FromScalar I think would also benefit from updates to missing.pyx, which creates an ndarray array today for the mask. If we had those create a BitMask array instead things would be faster
3. boolean.TimeLogicalOps is slower because we have to convert between BitMaskArrays and NumPy arrays when comparing values. I think the module can be refactored to defer NumPy conversions as much as possible
4. reshape.ReshapeMaskedArrayDtype regressions are again a cost of BitMaskArray <--> NumPy conversion. Note that `stack` reshaping is actually faster now, because of the BitMaskConcatenate method that has been implemented. I think more time would need to be spent for equivalent T / unstack fastpaths
To reduce some of the diff I implemented the buffer protocol in the new `BitmaskArray` class. This is in theory faster than a `to_numpy()` call and has the upside of working with our extensions that accept memoryview arguments.

Right now in the buffer protocol I am converting from bits to bytes, which is still relatively expensive. Both [PEP 3118 ](https://peps.python.org/pep-3118/#additions-to-the-struct-string-syntax) and the [NumPy array interface](https://numpy.org/doc/stable/reference/arrays.interface.html#python-side) mention the ability to work with bits, but they have different documented syntax (`Nt` versus `tN` where N is the number of bits) and I'm not clear that they were ever implemented? Maybe something @scoder or @seberg would know about?
I haven't had a chance to run a full benchmark yet but after implementing the buffer protocol for the bitmask array I was able to replace a lot of the `to_numpy()` calls (note - numpy arrays allow a memoryview indexer but pandas does not exhibit correct behavior - something we might want to patch on our end)

*~After the buffer protocol change performance seemed to skyrocket? I'm not sure I believe these numbers yet but small samples show improvements like below. I'm pretty skeptical of these since even non-EA types are showing a boost?~*

EDIT: just seemed like flaky asv results
OK I think this is reviewable now. As a quick summary:
    - Benchmarks are [here](https://github.com/pandas-dev/pandas/pull/54506#issuecomment-1676514037) - there are some regressions, some improvements. Slightly more regressions, though the constructor tests aren't really fair. Any/all checks on the masks are much faster
    - Any calls of `to_numpy()` are opportunities for future enhancement; these essentially convert our bits to bytes, which costs memory and time.
    - The buffer protocol also requires a conversion from bits to bytes. In theory [PEP-3118](https://peps.python.org/pep-3118/) and the NumPy [array interface protocol](https://numpy.org/doc/stable/reference/arrays.interface.html) document bits as a type, but I'm not sure either actually implemented support

If anyone else from @pandas-dev/pandas-core had time to try a few benchmarks it would be great too. Want to make sure the improvements here would work across all platforms
can the BitMaskArray support 2D?
Yea 2D is working with this. The "support" is limited to storing the bits in a C-style contiguous fashion. Slicing and non axis=0 operations right now require a `to_numpy()` call, but if you don't do that this works with 2D natively
To illustrate:

```python
>>> from pandas._libs.arrays import BitmaskArray
>>> import numpy as np
>>> arr = np.array([True, True, False, False]).reshape(2, -1)
>>> bma = BitmaskArray(arr)
>>> bma.any()
True
>>> bma.all()
False
>>> bma.shape
(2, 2)
>>> bma[0]
array([ True,  True])
>>> bma[1]
array([False, False])
>>> bma[:, 0]
array([ True, False])
>>> bma[:, 1]
array([ True, False])
>>> "{0:08b}".format(bma.bytes[0])
'00000011'
>>> bma.to_numpy()
array([[ True,  True],
       [False, False]])
```

This is all handled via `BitmaskArray.__getitem__` - the BitmaskArray has a few cases that it can fastpath and not have to convert to NumPy on (i.e. 1D scalar indexing, full slice) but generally for more complicated indexers it converts to NumPy and returns that as a fallback
> Both [PEP 3118 ](https://peps.python.org/pep-3118/#additions-to-the-struct-string-syntax)and the [NumPy array interface](https://numpy.org/doc/stable/reference/arrays.interface.html#python-side) mention the ability to work with bits, but they have different documented syntax (Nt versus tN where N is the number of bits) and I'm not clear that they were ever implemented? Maybe something @scoder or @seberg would know about?

I am not aware that NumPy actually supports this in any useful way at least.  `memoryview()` doesn't care if a format is unknown.  I would say the buffer protocol docs are more authoritative but it might be good to see if there is prior art that actually uses this.
From searching through https://discuss.python.org/t/question-pep-3118-format-strings-and-the-buffer-protocol/31264/8 I found a link to https://github.com/python/cpython/issues/47382#issuecomment-1561909807 which I think confirms that bit support isn't implemented in CPython

Of course not a blocker to this PR but something to consider in the future that would make this even more useful
Closing for now, but I think this is feasible if we ever wanted to come back to it