Make partition support non-1-indexed vector. (#40830)

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 <[email protected]>
JuliaLang · May 21, 2021 · 64940ec · 64940ec
1 parent db8d096
commit 64940ec
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 6 deletions.
diff --git a/base/iterators.jl b/base/iterators.jl
@@ -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
 

diff --git a/test/iterators.jl b/test/iterators.jl
@@ -2,6 +2,7 @@
 
 using Base.Iterators
 using Random
+using Base: IdentityUnitRange
 
 @test Base.IteratorSize(Any) isa Base.SizeUnknown
 
@@ -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
diff --git a/test/offsetarray.jl b/test/offsetarray.jl
@@ -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
@@ -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