Skip to content

Commit

Permalink
Improve getindex performance on ChainedVector with UnitRanges
Browse files Browse the repository at this point in the history
  • Loading branch information
quinnj committed Oct 24, 2024
1 parent b9a7ef7 commit 74e984d
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 1 deletion.
27 changes: 26 additions & 1 deletion src/chainedvector.jl
Original file line number Diff line number Diff line change
Expand Up @@ -177,14 +177,39 @@ Base.@propagate_inbounds function Base.getindex(x::ChainedVector{T, A}, inds::Ab
arrays = x.arrays
for i = 1:len
@inbounds ind2 = inds[i]
# @boundscheck checkbounds(x, ind2)
chunk, chunklen, j = linearindex(x, chunk, chunklen, j, ind, ind2)
@inbounds res[i] = arrays[chunk][j]
ind = ind2
end
return res
end

Base.@propagate_inbounds function Base.getindex(x::ChainedVector{T, A}, inds::UnitRange{Int}) where {T, A}
isempty(inds) && return similar(x, 0)
len = length(inds)
res = similar(x.arrays[1], len)
chunk = j = ind = 1
chunklen = length(x.arrays[1])
arrays = x.arrays
# linearindex first item
chunk, chunklen, j = linearindex(x, chunk, chunklen, j, ind, inds[1])
@inbounds arraychunk = arrays[chunk]
i = 1
# now we can copy entire chunks and avoid further linear indexing
while true
N = min(len, chunklen - j + 1)
unsafe_copyto!(res, i, arraychunk, j, N)
len -= N
len == 0 && break
i += N
chunk += 1
chunklen = length(arrays[chunk])
j = 1
arraychunk = arrays[chunk]
end
return res
end

Base.@propagate_inbounds function Base.isassigned(A::ChainedVector, i::Integer)
@boundscheck checkbounds(A, i)
chunk, ix = index(A, i)
Expand Down
11 changes: 11 additions & 0 deletions test/chainedvector.jl
Original file line number Diff line number Diff line change
Expand Up @@ -734,3 +734,14 @@ end
@test deleteat!(v2, m2) == deleteat!(s2, m2)
end
end

@testset "getindex with UnitRange" begin
x = ChainedVector([collect(1:i) for i = 10:100])
@test isempty(x[1:0])
@test x[1:1] == [1]
@test x[1:end] == x
y = collect(x)
for i = 1:length(x), j = 1:length(x)
@test x[i:j] == y[i:j]
end
end

0 comments on commit 74e984d

Please sign in to comment.