Improve getindex performance on ChainedVector with UnitRanges (#106)

* Improve getindex performance on ChainedVector with UnitRanges * Update src/chainedvector.jl Co-authored-by: Jerry Ling <[email protected]> --------- Co-authored-by: Jerry Ling <[email protected]>
JuliaData · Oct 24, 2024 · 3f83eb8 · 3f83eb8
1 parent b9a7ef7
commit 3f83eb8
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 1 deletion.
diff --git a/src/chainedvector.jl b/src/chainedvector.jl
@@ -177,14 +177,40 @@ 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)
+    arrays = x.arrays
+    res = similar(arrays[1], len)
+    chunk = j = ind = 1
+    chunklen = length(arrays[1])
+
+    # 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)

diff --git a/test/chainedvector.jl b/test/chainedvector.jl
@@ -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