Possible speedup in searchsortedfirst

[IlianPihlajamaa]

julia> versioninfo()
Julia Version 1.10.0
Commit 3120989f39 (2023-12-25 18:01 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: 12 × Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, skylake)
  Threads: 1 on 12 virtual cores

I get a good speedup by the following change (here searchsortedfirst1 is the function copied from Base.Sort).

import Base.Sort: lt, DirectOrdering, ForwardOrdering, ReverseOrdering
import Base:require_one_based_indexing

function searchsortedfirst1(a::AbstractRange{<:Real}, x::Real, o::DirectOrdering)::keytype(a)
    require_one_based_indexing(a)
    f, h, l = first(a), step(a), last(a)
    if !lt(o, f, x)
        1
    elseif h == 0 || lt(o, l, x)
        length(a) + 1
    else
        n = round(Integer, (x - f) / h + 1)
        lt(o, a[n], x) ? n + 1 : n
    end
end


function searchsortedfirst2(a::AbstractRange{<:Real}, x::Real, o::DirectOrdering)::keytype(a)
    require_one_based_indexing(a)
    f, h, l = first(a), step(a), last(a)
    if !lt(o, f, x)
        1
    elseif h == 0 || lt(o, l, x)
        length(a) + 1
    else
        ceil(Integer, (x - f) / h + 1) # this bit is different
    end
end

using BenchmarkTools

function main()
    a = range(0, 1, length=10)
    x = rand()

    n1 = searchsortedfirst1(a, x, ForwardOrdering())
    n2 = searchsortedfirst2(a, x, ForwardOrdering())
    @show n1 == n2

    # test for point exactly on gridpoint
    n1 = searchsortedfirst1(a, a[5], ForwardOrdering()) 
    n2 = searchsortedfirst2(a, a[5], ForwardOrdering())
    @show n1 == n2

    @btime searchsortedfirst1($a, $x, ForwardOrdering()) 
    @btime searchsortedfirst2($a, $x, ForwardOrdering())

    a = range(1, 0, length=10)
    x = rand()

    n1 = searchsortedfirst1(a, x, ReverseOrdering())
    n2 = searchsortedfirst2(a, x, ReverseOrdering())
    @show n1 == n2

    # test for point exactly on gridpoint
    n1 = searchsortedfirst1(a, a[5], ReverseOrdering()) 
    n2 = searchsortedfirst2(a, a[5], ReverseOrdering())
    @show n1 == n2

    @btime searchsortedfirst1($a, $x, ReverseOrdering()) 
    @btime searchsortedfirst2($a, $x, ReverseOrdering())
end

main()

n1 == n2 = true
n1 == n2 = true
27.236 ns (0 allocations: 0 bytes)
13.614 ns (0 allocations: 0 bytes)
n1 == n2 = true
n1 == n2 = true
27.035 ns (0 allocations: 0 bytes)
17.735 ns (0 allocations: 0 bytes)

[mikmoore]

It looks like you should open a PR for this altered implementation.

However, my initial suspicion is that this will very rarely lead to the wrong result. I expect that the lt(o, a[n], x) ? n + 1 : n switch in the original is necessary to correct for very rare rounding "errors" that would lead to off-by-one results. I imagine (hope?) we have some tests for such cases, which a PR would be evaluated against.

[IlianPihlajamaa]

You are right, of course. I found such a pathological case:

a = 0.7346369851144842:2.097212148566592e-10:1.1674647790886943
idx = 568251599

n1 = searchsortedfirst1(a, a[idx], ForwardOrdering()) 
n2 = searchsortedfirst2(a, a[idx], ForwardOrdering())
n1 == n2 # false