Skip to content

Commit

Permalink
Make partition support non-1-indexed vector. (JuliaLang#40830)
Browse files Browse the repository at this point in the history
For a high dimension array, at least `OffsetArrays.jl` return a 1-indexed vector after reshape, so the current implementation seems ok.

I believe that the view of Base's Range types has been mapped to `getindex` correctly, so I think we don't need a separate routine.

Co-authored-by: Johnny Chen <johnnychen94@hotmail.com>
  • Loading branch information
2 people authored and Amit Shirodkar committed Jun 9, 2021
1 parent ba621b8 commit 770a9cd
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 6 deletions.
12 changes: 6 additions & 6 deletions base/iterators.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1171,15 +1171,15 @@ function length(itr::PartitionIterator)
return cld(l, itr.n)
end

function iterate(itr::PartitionIterator{<:AbstractRange}, state=1)
state > length(itr.c) && return nothing
r = min(state + itr.n - 1, length(itr.c))
function iterate(itr::PartitionIterator{<:AbstractRange}, state = firstindex(itr.c))
state > lastindex(itr.c) && return nothing
r = min(state + itr.n - 1, lastindex(itr.c))
return @inbounds itr.c[state:r], r + 1
end

function iterate(itr::PartitionIterator{<:AbstractArray}, state=1)
state > length(itr.c) && return nothing
r = min(state + itr.n - 1, length(itr.c))
function iterate(itr::PartitionIterator{<:AbstractArray}, state = firstindex(itr.c))
state > lastindex(itr.c) && return nothing
r = min(state + itr.n - 1, lastindex(itr.c))
return @inbounds view(itr.c, state:r), r + 1
end

Expand Down
5 changes: 5 additions & 0 deletions test/iterators.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

using Base.Iterators
using Random
using Base: IdentityUnitRange

@test Base.IteratorSize(Any) isa Base.SizeUnknown

Expand Down Expand Up @@ -848,3 +849,7 @@ end
@test cumprod(x + 1 for x in 1:3) == [2, 6, 24]
@test accumulate(+, (x^2 for x in 1:3); init=100) == [101, 105, 114]
end

@testset "proper patition for non-1-indexed vector" begin
@test partition(IdentityUnitRange(11:19), 5) |> collect == [11:15,16:19] # IdentityUnitRange
end
9 changes: 9 additions & 0 deletions test/offsetarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

isdefined(Main, :OffsetArrays) || @eval Main include("testhelpers/OffsetArrays.jl")
using .Main.OffsetArrays
import .Main.OffsetArrays: IdOffsetRange
using DelimitedFiles
using Random
using LinearAlgebra
Expand Down Expand Up @@ -786,3 +787,11 @@ end
@test b[i] == a[r[i]]
end
end

@testset "proper patition for non-1-indexed vector" begin
@test Iterators.partition(OffsetArray(1:10,10), 5) |> collect == [1:5,6:10] # OffsetVector
@test Iterators.partition(OffsetArray(collect(1:10),10), 5) |> collect == [1:5,6:10] # OffsetVector
@test Iterators.partition(OffsetArray(reshape(1:9,3,3), (3,3)), 5) |> collect == [1:5,6:9] #OffsetMatrix
@test Iterators.partition(OffsetArray(reshape(collect(1:9),3,3), (3,3)), 5) |> collect == [1:5,6:9] #OffsetMatrix
@test Iterators.partition(IdOffsetRange(2:7,10), 5) |> collect == [12:16,17:17] # IdOffsetRange
end

0 comments on commit 770a9cd

Please sign in to comment.