diff --git a/src/chainedvector.jl b/src/chainedvector.jl index 14028bf..e8c46c8 100644 --- a/src/chainedvector.jl +++ b/src/chainedvector.jl @@ -242,6 +242,7 @@ end # custom index type used in eachindex struct ChainedVectorIndex{A} <: Integer + arrays::Vector{A} arrays_i::Int array::A array_i::Int @@ -258,17 +259,23 @@ Base.hash(x::ChainedVectorIndex, h::UInt) = hash(x.i, h) @inline Base.getindex(x::ChainedVectorIndex) = @inbounds x.array[x.array_i] +# An index from `eachindex(B)` may be used on another ChainedVector `A`, e.g. `similar(B)`. +# It addresses the chunk it holds only when used on `B`; elsewhere it's the linear index `x.i`. +@inline ownsindex(A::ChainedVector, x::ChainedVectorIndex) = x.arrays === A.arrays + function Base.checkbounds(::Type{Bool}, A::ChainedVector, ind::ChainedVectorIndex) - @assert ind.array === A.arrays[ind.arrays_i] "indexing ChainedVector with wrong ChainedVectorIndex" + ownsindex(A, ind) || return checkbounds(Bool, A, ind.i) return 1 <= ind.array_i <= length(ind.array) end Base.@propagate_inbounds function Base.getindex(A::ChainedVector, x::ChainedVectorIndex) + ownsindex(A, x) || return A[x.i] @boundscheck checkbounds(A, x) return @inbounds x.array[x.array_i] end Base.@propagate_inbounds function Base.setindex!(A::ChainedVector, v, x::ChainedVectorIndex) + ownsindex(A, x) || return setindex!(A, v, x.i) @boundscheck checkbounds(A, x) @inbounds x.array[x.array_i] = v return v @@ -278,7 +285,7 @@ function Base.getindex(A::ChainedVector{T}, inds::AbstractVector{<:ChainedVector len = length(inds) x = Vector{T}(undef, len) for i = 1:len - x[i] = inds[i][] + x[i] = A[inds[i]] end return x end @@ -298,8 +305,9 @@ function Base.nextind(A::ChainedVector, x::ChainedVectorIndex) i += 1 else chunk_i += 1 # make sure this goes out of bounds + i += 1 end - return ChainedVectorIndex(chunkidx, chunk, chunk_i, i) + return ChainedVectorIndex(x.arrays, chunkidx, chunk, chunk_i, i) end function Base.prevind(A::ChainedVector, x::ChainedVectorIndex) @@ -317,8 +325,9 @@ function Base.prevind(A::ChainedVector, x::ChainedVectorIndex) i -= 1 else chunk_i -= 1 # make sure this goes out of bounds + i -= 1 end - return ChainedVectorIndex(chunkidx, chunk, chunk_i, i) + return ChainedVectorIndex(x.arrays, chunkidx, chunk, chunk_i, i) end # efficient iteration via eachindex @@ -342,7 +351,7 @@ end chunkidx = chunk_i = 1 @inbounds chunk = arrays[chunkidx] # we already ran cleanup! so chunks are guaranteed non-empty - return ChainedVectorIndex(chunkidx, chunk, chunk_i, 1), (arrays, chunkidx, chunk, length(chunk), chunk_i + 1, 2) + return ChainedVectorIndex(arrays, chunkidx, chunk, chunk_i, 1), (arrays, chunkidx, chunk, length(chunk), chunk_i + 1, 2) end @inline function Base.iterate(x::IndexIterator, (arrays, chunkidx, chunk, chunklen, chunk_i, i)) @@ -353,7 +362,7 @@ end chunklen = length(chunk) chunk_i = 1 end - return ChainedVectorIndex(chunkidx, chunk, chunk_i, i), (arrays, chunkidx, chunk, chunklen, chunk_i + 1, i + 1) + return ChainedVectorIndex(arrays, chunkidx, chunk, chunk_i, i), (arrays, chunkidx, chunk, chunklen, chunk_i + 1, i + 1) end @inline function Base.iterate(A::ChainedVector) @@ -955,7 +964,17 @@ function Base.filter!(f, a::ChainedVector) return a end -Base.replace(f::Base.Callable, a::ChainedVector) = ChainedVector([replace(f, A) for A in a.arrays]) -Base.replace!(f::Base.Callable, a::ChainedVector) = (foreach(A -> replace!(f, A), a.arrays); return a) -Base.replace(a::ChainedVector, old_new::Pair...; count::Union{Integer,Nothing}=nothing) = ChainedVector([replace(A, old_new...; count=count) for A in a.arrays]) -Base.replace!(a::ChainedVector, old_new::Pair...; count::Integer=typemax(Int)) = (foreach(A -> replace!(A, old_new...; count=count), a.arrays); return a) +# Without a `count` limit, chunks are replaced independently. A limit spans chunks, +# so limited replacements use Base's elementwise methods. +Base.replace(f::Base.Callable, a::ChainedVector; count::Integer=typemax(Int)) = + count < length(a) ? invoke(replace, Tuple{Base.Callable, Any}, f, a; count=count) : + ChainedVector([replace(f, A) for A in a.arrays]) +Base.replace!(f::Base.Callable, a::ChainedVector; count::Integer=typemax(Int)) = + count < length(a) ? invoke(replace!, Tuple{Base.Callable, Any}, f, a; count=count) : + (foreach(A -> replace!(f, A), a.arrays); a) +Base.replace(a::ChainedVector, old_new::Pair...; count::Union{Integer,Nothing}=nothing) = + count === nothing ? ChainedVector([replace(A, old_new...) for A in a.arrays]) : + invoke(replace, Tuple{Any, Vararg{Pair}}, a, old_new...; count=count) +Base.replace!(a::ChainedVector, old_new::Pair...; count::Integer=typemax(Int)) = + count < length(a) ? invoke(replace!, Tuple{Any, Vararg{Pair}}, a, old_new...; count=count) : + (foreach(A -> replace!(A, old_new...), a.arrays); a) diff --git a/test/chainedvector.jl b/test/chainedvector.jl index 7d30cd9..e0cce1c 100644 --- a/test/chainedvector.jl +++ b/test/chainedvector.jl @@ -380,6 +380,17 @@ end @test map(x -> x == 1 ? 2.0 : x, x) == replace!(x, 1 => 2) @test isempty(x) + @test replace!(ChainedVector([[1,2], [1,2]]), 2=>20) == [1,20,1,20] + @test replace!(ChainedVector([[1,2], [1,2]]), 2=>20, count=1) == [1,20,1,2] + @test replace!(ChainedVector([[1,2], [1,2]]), 2=>20, count=2) == [1,20,1,20] + x = [1,2] + @test replace!(ChainedVector([x,[2,3]]), 2=>99) == [1,99,99,3] + @test x == [1,99] + # results stay ChainedVectors and promote like Base + @test replace(ChainedVector([[1,2], [1,2]]), 2=>2.5) isa ChainedVector{Float64} + @test replace(ChainedVector([[1,2], [1,2]]), 2=>2.5, count=1) == [1,2.5,1,2] + @test replace(ChainedVector([[1,missing], [missing]]), missing=>0) isa ChainedVector{Int} + # copyto! # ChainedVector dest: doffs, soffs, n x = ChainedVector([[1,2,3], [4,5,6], [7,8,9,10]]) @@ -623,6 +634,48 @@ end end end +@testset "replace[!] comparison with Vector" begin + + testvecs = ( + [[1, 2], [3, 2, 5]], + [[1, 2]], + [[2],[2],[2],[2,3]], + [[1,2,missing]], + [[missing,1],[missing,2,1]], + [[missing]] + ) + function missing_equal(a,b) + ismissing(a) && ismissing(b) && return true + ismissing(a) ⊻ ismissing(b) && return false + return all(skipmissing(a) .== skipmissing(b)) + end + gen_cv_v(x) = (c = ChainedVector(x); (c, collect(c))) + for f in (replace, replace!) + for x in testvecs + cv, v = gen_cv_v(x) + @test missing_equal(f(v, 2 => 22),f(cv, 2 => 22)) + @test missing_equal(v,cv) + + cv, v = gen_cv_v(x) + @test missing_equal(f(x -> x ÷ 2, v), f(x -> x ÷ 2, cv)) + @test missing_equal(v,cv) + + for c in (0, 1, 2, 3) + cv, v = gen_cv_v(x) + @test missing_equal(f(x -> x ÷ 2, v, count=c), f(x -> x ÷ 2, cv, count=c)) + @test missing_equal(v,cv) + + for p in ((2=>2,),(2 => 22,), (2 => 22, 3 => 33)) + cv, v = gen_cv_v(x) + @test missing_equal(f(v, p..., count=c), f(cv, p..., count=c)) + @test eltype(f(v, p..., count=c)) == eltype(f(cv, p..., count=c)) + @test missing_equal(v,cv) + end + end + end + end +end + @testset "iteration protocol on ChainedVector" begin for len in 0:6 @@ -733,8 +786,29 @@ end ind = prevind(x, ind) @test_throws BoundsError x[ind] # https://github.com/apache/arrow-julia/issues/418 + # indices from one ChainedVector index others by position, even under @inbounds y = ChainedVector([collect(1:i) for i = 10:100]) - @test_throws AssertionError x[first(eachindex(y))] + z = ChainedVector([x[1:7], x[8:end]]) # same values, different chunks + for (i, ind) in enumerate(eachindex(y)) + @test x[ind] == @inbounds(x[ind]) == z[ind] == x[i] == z[i] + @inbounds y[ind] = -i + end + @test y == -(1:length(y)) + w = similar(y) + for ind in eachindex(y) + @inbounds w[ind] = y[ind] + end + @test w == y + longer = ChainedVector([x.arrays; [[1]]]) + @test_throws BoundsError x[collect(eachindex(longer))[end]] + p = ChainedVector([[1, 2], [3, 4]]) + i = first(eachindex(ChainedVector([[10], [20, 30, 40]]))) + @test p[nextind(p, i)] == 2 + @test p[prevind(p, nextind(p, nextind(p, i)))] == 2 + # generic Base code indexes `similar(A)` with `eachindex(A)` + a = ChainedVector([[1, 2], [1, 2]]) + @test replace(x -> 10x, a; count=3) == [10, 20, 10, 2] + @test a == [1, 2, 1, 2] # https://github.com/JuliaData/SentinelArrays.jl/issues/74 x = ChainedVector([[true], [false], [true]]) @test BitVector(x) == [true, false, true] @@ -782,7 +856,7 @@ end end @testset "getindex with UnitRange" begin - x = ChainedVector([collect(1:i) for i = 10:100]) + x = ChainedVector([collect(1:i) for i = 1:10]) @test isempty(x[1:0]) @test x[1:1] == [1] @test x[1:end] == x