PERF improve performance of is_lexsorted

[MarcoGorelli]

If result is False, there's no need to keep the outer loop going, right?

Timing result, with failure taken from test_is_lexsorted:

%load_ext cython

%%cython -a

cimport cython
from cython cimport Py_ssize_t
from numpy cimport int64_t, import_array, ndarray, PyArray_DATA
import numpy as np

from libc.stdlib cimport free, malloc
from libc.stdio cimport printf
import_array()


@cython.wraparound(False)
@cython.boundscheck(False)
def main_is_lexsorted(list_of_arrays: list) -> bint:
    cdef:
        Py_ssize_t i
        Py_ssize_t n, nlevels
        int64_t k, cur, pre
        ndarray arr
        bint result = True

    nlevels = len(list_of_arrays)
    n = len(list_of_arrays[0])

    cdef int64_t **vecs = <int64_t**>malloc(nlevels * sizeof(int64_t*))
    for i in range(nlevels):
        arr = list_of_arrays[i]
        assert arr.dtype.name == 'int64'
        vecs[i] = <int64_t*>PyArray_DATA(arr)

    # Assume uniqueness??
    with nogil:
        for i in range(1, n):
            for k in range(nlevels):
                cur = vecs[k][i]
                pre = vecs[k][i -1]
                if cur == pre:
                    continue
                elif cur > pre:
                    break
                else:
                    result = False
                    break
    free(vecs)
    return result

@cython.wraparound(False)
@cython.boundscheck(False)
def branch_is_lexsorted(list_of_arrays: list) -> bint:
    cdef:
        Py_ssize_t i
        Py_ssize_t n, nlevels
        int64_t k, cur, pre
        ndarray arr
        bint result = True

    nlevels = len(list_of_arrays)
    n = len(list_of_arrays[0])

    cdef int64_t **vecs = <int64_t**>malloc(nlevels * sizeof(int64_t*))
    for i in range(nlevels):
        arr = list_of_arrays[i]
        assert arr.dtype.name == 'int64'
        vecs[i] = <int64_t*>PyArray_DATA(arr)

    # Assume uniqueness??
    with nogil:
        for i in range(1, n):
            for k in range(nlevels):
                cur = vecs[k][i]
                pre = vecs[k][i -1]
                if cur == pre:
                    continue
                elif cur > pre:
                    break
                else:
                    result = False
                    break
            if not result:
                break
    free(vecs)
    return result

%%timeit
main_is_lexsorted(failure)
# 16.6 µs ± 411 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
 
%%timeit
branch_is_lexsorted(failure)
# 15.8 µs ± 113 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

More extreme example:

failure = [np.ones((100_000,), dtype='int64')]
failure[0][1] = 0 

%%timeit
main_is_lexsorted(failure)
# 556 µs ± 26.8 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%%timeit
branch_is_lexsorted(failure)
# 8.07 µs ± 72.3 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

[datapythonista]

This makes sense to me. But probably even better instead of the double break, to simply return when we know it's False. And we don't even need the result variable, since the return at the end of the function should always be return True.

Suggested change

	result = False
	break
	if not result:
	break
	free(vecs)
	return False

[MarcoGorelli]

looks like this isn't allowed

Error compiling Cython file:
------------------------------------------------------------
...
                else:
                    result = False
                    break
            if not result:
                free(vecs)
                return False
               ^
------------------------------------------------------------
pandas/_libs/algos.pyx:184:16: Returning Python object not allowed without gil

[datapythonista]

Good to know, I was wondering if there was a Cython limitation. Thanks for trying it.

[mroeschke]

LGTM. Off to @datapythonista

[jreback]

this diff is showing a lot of changes