Parallel argsort via c++ built-in parallel mode.

argsort.h

// #include <algorithm> // std::sort, std::stable_sort
#include <parallel/algorithm>

void argsort(unsigned long *pidx, const double *parr, unsigned long length)
{
    // Sort indexes based on comparing values in v
    //   Because Cython pass data blob as pointer, so access the array as pointer style.
    unsigned long *pidx2 = pidx + length;
    // double *parr2 = parr + length;
    __gnu_parallel::sort(pidx, pidx2,
                         [parr](size_t i1, size_t i2)
                         { return *(parr + i1) < *(parr + i2); });
}

parallel-argsort.md

compile:

python3 setup.py build_ext -if

using:

from pargsort import pargsort
import numpy as np

arr = np.random.uniform(0,100,100000000)

%timeit pargsort(arr)
%timeit np.argsort(arr)

Unkonw problem: When processing a very large number of numbers, such as arr = np.random.uniform(0,100,100000000) (an extra zero), the function will degrade to single core.

pargsort.pyx

import numpy as np
cimport numpy as np
import cython
cimport cython


cdef extern from "argsort.h":
    cdef void argsort(unsigned long *pidx, double *parr, unsigned long length) nogil


@cython.boundscheck(False) # turn off bounds-checking for entire function
@cython.wraparound(False)  # turn off negative index wrapping for entire function
def pargsort(np.ndarray arr):
    """
    In-place parallel sort for numpy types
    """
    cdef unsigned long l = arr.shape[0]
    cdef double[:] arr2 = arr.astype(np.double)
    cdef unsigned long[:] idx2 = np.arange(l, dtype=np.uint)
    argsort(&idx2[0], &arr2[0], l) # c++ only consider memview as a pointer, so pass the whole is unnecessary.
    return np.asarray(idx2)

setup.py

from setuptools import Extension, setup
from Cython.Build import cythonize

ext_modules = [
    Extension(
        "pargsort",
        ["pargsort.pyx"],
        language="c++",
        extra_compile_args=["-fopenmp"],
        extra_link_args=["-fopenmp"],
    )
]

setup(
    name="pargsort",
    ext_modules=cythonize(ext_modules, compiler_directives={"language_level": "3"}),
)