[C++] Filter compute function seems slow compared to numpy nonzero + take

[asfimport]

From https://stackoverflow.com/questions/64581590/is-there-a-more-efficient-way-to-select-rows-from-a-pyarrow-table-based-on-conte

I made a smaller, simplified example:

arr = pa.array(np.random.randn(1_000_000))

# mask with only few True values
mask1 = np.zeros(len(arr), dtype=bool)
mask1[np.random.randint(len(arr), size=100)] = True
mask1_pa = pa.array(mask1)

# mask with larger proportion of True values
mask2 = np.zeros(len(arr), dtype=bool)
mask2[np.random.randint(len(arr), size=10_000)] = True
mask2_pa = pa.array(mask2)

Doing timings of doing a Arrow Filter kernel vs using numpy to convert the mask into indices and then using a Take kernel:

# mask 1
In [3]: %timeit arr.filter(mask1_pa)
132 µs ± 4.44 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [4]: %%timeit
   ...: indices = np.nonzero(mask1)[0]
   ...: arr.take(indices)
114 µs ± 2.62 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

# mask 2
In [8]: %timeit arr.filter(mask2_pa)
711 µs ± 63.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [9]: %%timeit
   ...: indices = np.nonzero(mask2)[0]
   ...: arr.take(indices)
333 µs ± 6.32 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

So in the first case, both are quite similar in timing. But in the second case, the numpy+take version is faster.

I know this might depend on a lot on the actual proportion of True values and how they are positioned in the array (random vs concentrated) etc, so there is probably not a general rule of what should be faster.
But, it still seems a potential indication that things can be optimized in the Filter kernel.

Reporter: Joris Van den Bossche / @jorisvandenbossche

Related issues:

• [C++] Investigate a bit run reader that would only return runs of set bits (is fixed by)
• [C++] Implement vector kernel that transforms a boolean mask into selection indices (is related to)

_{Note: This issue was originally created as ARROW-10423. Please see the migration documentation for further details.}

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
And note I timed it with released versions from conda-forge (so not an issue of forgetting that I am using a debug build ;-))

[asfimport]

Antoine Pitrou / @pitrou:
Sounds like Arrow might not preallocate the result array? It should be easy to compute its size based on a popcount of the filter.

[asfimport]

Wes McKinney / @wesm:
We should also expose the equivalent of numpy.nonzero as a kernel (the code to do this already exists, it just hasn't been "kernelized")

[asfimport]

Jason Sachs:

We should also expose the equivalent of numpy.nonzero as a kernel (the code to do this already exists, it just hasn't been "kernelized")

yes please =)

[asfimport]

Antoine Pitrou / @pitrou:
Here is the current status here. It seems Arrow is now much faster than Numpy at filtering:

>>> %timeit arr.filter(mask1_pa)
20.4 µs ± 11.4 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
>>> %timeit arr.to_numpy()[mask1]
96.8 µs ± 605 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

>>> %timeit arr.filter(mask2_pa)
134 µs ± 334 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
>>> %timeit arr.to_numpy()[mask2]
284 µs ± 545 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

And it's roughly as fast in taking:

>>> %timeit indices = np.nonzero(mask1)[0]; arr.take(indices)
159 µs ± 308 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
>>> %timeit indices = np.nonzero(mask1)[0]; arr.to_numpy()[indices]
142 µs ± 168 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

>>> %timeit indices = np.nonzero(mask2)[0]; arr.take(indices)
274 µs ± 1.83 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
>>> %timeit indices = np.nonzero(mask2)[0]; arr.to_numpy()[indices]
285 µs ± 3.24 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

[asfimport]

Antoine Pitrou / @pitrou:
It appears the performance gap was fixed in ARROW-10696.