add Iterators.partition

NEWS.md

@@ -1,3 +1,10 @@
+# DataFrames.jl v1.5 Release Notes
+
+## New functionalities
+
+* Add `Iterators.partition` support
+   ([#3212](https://github.com/JuliaData/DataFrames.jl/pull/3212))
+
 # DataFrames.jl v1.4.2 Patch Release Notes
 
 ## Bug fixes

Project.toml

@@ -1,6 +1,6 @@
 name = "DataFrames"
 uuid = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
-version = "1.4.2"
+version = "1.5.0"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

src/abstractdataframe/abstractdataframe.jl

@@ -3291,3 +3291,62 @@ function insertcols!(df::AbstractDataFrame; after::Bool=false,
     _drop_all_nonnote_metadata!(parent(df))
     return df
 end
+
+"""
+    Iterators.partition(df::AbstractDataFrame, n)
+
+Iterate over `df` data frame `n` rows at a time, returning each block
+as a `SubDataFrame`.
+
+# Examples
+
+```jldoctest
+julia> collect(Iterators.partition(DataFrame(x=1:5), 2))
+3-element Vector{SubDataFrame{DataFrame, DataFrames.Index, UnitRange{Int64}}}:
+ 2×1 SubDataFrame
+ Row │ x
+     │ Int64
+─────┼───────
+   1 │     1
+   2 │     2
+ 2×1 SubDataFrame
+ Row │ x
+     │ Int64
+─────┼───────
+   1 │     3
+   2 │     4
+ 1×1 SubDataFrame
+ Row │ x
+     │ Int64
+─────┼───────
+   1 │     5
+```
+"""
+function Iterators.partition(df::AbstractDataFrame, n::Integer)
+    n < 1 && throw(ArgumentError("cannot create partitions of length $n"))
+    return Iterators.PartitionIterator(df, Int(n))
+end
+
+# use autodetection of eltype
+Base.IteratorEltype(::Type{<:Iterators.PartitionIterator{<:AbstractDataFrame}}) =
+    Base.EltypeUnknown()
+
+# we do not need to be overly specific here as we rely on autodetection of eltype
+Base.eltype(::Type{<:Iterators.PartitionIterator{<:AbstractDataFrame}}) =
+    AbstractDataFrame
+
+IteratorSize(::Type{<:Iterators.PartitionIterator{<:AbstractDataFrame}}) =
+    Base.HasLength()
+
+function Base.length(itr::Iterators.PartitionIterator{<:AbstractDataFrame})
+    l = nrow(itr.c)
+    return cld(l, itr.n)
+end
+
+function Base.iterate(itr::Iterators.PartitionIterator{<:AbstractDataFrame}, state=1)
+    last_idx = nrow(itr.c)
+    state > last_idx && return nothing
+    r = min(state + itr.n - 1, last_idx)
+    return view(itr.c, state:r, :), r + 1
+end
+

test/dataframe.jl

@@ -2254,4 +2254,27 @@ end
     @test !isempty(DataFrame(a=1))
 end
 
+@testset "Iterators.partition" begin
+    for df in (DataFrame(x=1:5), view(DataFrame(x=1:6, y=11:16), 1:5, 1:1))
+        p = Iterators.partition(df, 2)
+        @test p isa Iterators.PartitionIterator
+        @test Tables.partitions(p) === p
+        @test eltype(p) === AbstractDataFrame
+        @test Base.IteratorEltype(typeof(p)) === Base.EltypeUnknown()
+        @test length(p) == 3
+        @test Base.IteratorSize(typeof(p)) === Base.HasLength()
+        res = collect(p)
+        @test res == [DataFrame(x=1:2), DataFrame(x=3:4), DataFrame(x=5)]
+        @test all(v -> v isa SubDataFrame, res)
+        @test_throws ArgumentError Iterators.partition(df, false)
+        @test_throws ArgumentError Iterators.partition(df, -1)
+    end
+    p = Iterators.partition(DataFrame(), 1)
+    @test p isa Iterators.PartitionIterator
+    @test Tables.partitions(p) === p
+    @test isempty(p)
+    @test length(p) == 0
+    @test eltype(collect(p)) <: SubDataFrame
+end
+
 end # module