Skip to content

Commit

Permalink
Make findall faster for AbstractArrays (JuliaLang#37177)
Browse files Browse the repository at this point in the history
The `findall` fallback is quite slow when predicate is a small function
compared with generating a logical index using `broadcast` and calling `findall`
on it to compute integer indices. The gain is most visible when predicate is true
for a large proportion of entries, but it's there even when all of them are false.
The drawback of this approach is that it requires allocating a vector of `length(a)/8`
bytes whatever the number of returned indices.
  • Loading branch information
nalimilan committed Dec 18, 2020
1 parent 9921e8c commit d20ca48
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion base/array.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2142,6 +2142,10 @@ julia> findall(x -> x >= 0, d)
"""
findall(testf::Function, A) = collect(first(p) for p in pairs(A) if testf(last(p)))

# Broadcasting is much faster for small testf, and computing
# integer indices from logical index using findall has a negligible cost
findall(testf::Function, A::AbstractArray) = findall(testf.(A))

"""
findall(A)
Expand Down Expand Up @@ -2439,7 +2443,8 @@ function findall(pred::Fix2{typeof(in),<:Union{Array{<:Real},Real}}, x::Array{<:
end
# issorted fails for some element types so the method above has to be restricted
# to element with isless/< defined.
findall(pred::Fix2{typeof(in)}, x::Union{AbstractArray, Tuple}) = _findin(x, pred.x)
findall(pred::Fix2{typeof(in)}, x::AbstractArray) = _findin(x, pred.x)
findall(pred::Fix2{typeof(in)}, x::Tuple) = _findin(x, pred.x)

# Copying subregions
function indcopy(sz::Dims, I::Vector)
Expand Down

0 comments on commit d20ca48

Please sign in to comment.