Summary

The following tests are meant to give some insights into the different approaches and should be taken with a grain of salt.

What is being tested here is not exactly generic filtering, but rather just applying a threshold, which has the notable feature that computing the condition is pretty fast. Very different results would be obtained if the condition would imply expensive computations.

Basically a 2-pass accelerated (either with Numba or Cython -- as long as you know the types in advance) is going to be the fastest and more memory efficient, except for very large inputs, for which single pass Numba/Cython are faster (at the cost of larger temporary memory usage). Using np.where()/np.nonzero(), instead of the mask directly, may lead to slightly faster computation, and typically does not hurt (except perhaps for the larger temporary memory footprint) as long as the size of the output is less than 50% of the input. The np.fromiter() approach is much slower, but does not produce large temporary objects.

Definitions

1. Using generators:

def filter_fromiter(arr, k):
    return np.fromiter((x for x in arr if x < k), dtype=arr.dtype)

2. Using boolean mask slicing:

def filter_mask(arr, k):
    return arr[arr < k]

3. Using np.where():

def filter_where(arr, k):
    return arr[np.where(arr < k)]

4. Using np.nonzero()

def filter_nonzero(arr, k):
    return arr[np.nonzero(arr < k)]

5. Using a Cython-based custom implementation(s):
   • single-pass filter_cy()
   • two-passes filter2_cy()

%%cython -c-O3 -c-march=native -a
#cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True


cimport numpy as cnp
cimport cython as ccy

import numpy as np
import cython as cy


cdef long NUM = 1048576
cdef long MAX_VAL = 1048576
cdef long K = 1048576 // 2


cdef int smaller_than_cy(long x, long k=K):
    return x < k


cdef size_t _filter_cy(long[:] arr, long[:] result, size_t size, long k):
    cdef size_t j = 0
    for i in range(size):
        if smaller_than_cy(arr[i]):
            result[j] = arr[i]
            j += 1
    return j


cpdef filter_cy(arr, k):
    result = np.empty_like(arr)
    new_size = _filter_cy(arr, result, arr.size, k)
    result.resize(new_size)
    return result


cdef size_t _filtered_size(long[:] arr, size_t size, long k):
    cdef size_t j = 0
    for i in range(size):
        if smaller_than_cy(arr[i]):
            j += 1
    return j


cpdef filter2_cy(arr, k):
    cdef size_t new_size = _filtered_size(arr, arr.size, k)
    result = np.empty(new_size, dtype=arr.dtype)
    new_size = _filter_cy(arr, result, arr.size, k)
    return result

import functools


filter_np_cy = functools.partial(filter_cy, k=K)
filter_np_cy.__name__ = 'filter_np_cy'


filter2_np_cy = functools.partial(filter2_cy, k=K)
filter2_np_cy.__name__ = 'filter2_np_cy'

6. Using a Numba-based custom implementation
   • single-pass filter_np_nb()
   • two-passes filter2_np_nb()

import numba as nb
import functools


@nb.njit
def filter_func(x, k):
    return x < k


@nb.njit
def filter_nb(arr, result, k):
    j = 0
    for i in range(arr.size):
        if filter_func(arr[i], k):
            result[j] = arr[i]
            j += 1
    return j


def filter_np_nb(arr, k=K):
    result = np.empty_like(arr)
    j = filter_nb(arr, result, k)
    result.resize(j, refcheck=False)
    return result


@nb.njit
def filter2_nb(arr, k):
    j = 0
    for i in range(arr.size):
        if filter_func(arr[i], k):
            j += 1
    result = np.empty(j, dtype=arr.dtype)
    j = 0
    for i in range(arr.size):
        if filter_func(arr[i], k):
            result[j] = arr[i]
            j += 1
    return result


filter2_np_nb = functools.partial(filter2_nb, k=K)
filter2_np_nb.__name__ = 'filter2_np_nb'

Timing Benchmarks

The generator-based filter_fromiter() method is much slower than the others (by approx. 2 orders of magnitude and it is therefore omitted in the charts).

The timing would depend on both the input array size and the percent of filtered items.

As a function of input size

The first graph addresses the timings as a function of the input size (for ~50% filtered out elements):

In general, the Numba based approaches are consistently the fastest, closely followed by the Cython approach. Within them, the two-passes approaches are typically fastest, except for very large inputs where the single-pass approach tend to take over. Within NumPy, the np.where()-based and np.nonzero()-based approaches are basically the same (except for very small inputs for which np.nonzero() seems to be slightly slower), and they are both faster than the boolean mask slicing, except for very small inputs (below ~100 elements) where the boolean mask slicing is faster. Moreover, for very small inputs, the Cython based solution are slower than NumPy-based ones.