add Iterators.partition for DataFrameRows (#3299)

Author: bkamins

NEWS.md

@@ -2,6 +2,8 @@
 
 ## New functionalities
 
+* Add `Iterators.partition` support for `DataFrameRows`
+  ([#3299](https://github.com/JuliaData/DataFrames.jl/pull/3299))
 * `DataFrameRows` and `DataFrameColumns` now support
   `nrow`, `ncol`, and `Tables.subset`
   ([#3311](https://github.com/JuliaData/DataFrames.jl/pull/3311))

src/abstractdataframe/iteration.jl

@@ -95,6 +95,66 @@ Compat.hasproperty(itr::DataFrameRows, s::AbstractString) = haskey(index(parent(
 # Private fields are never exposed since they can conflict with column names
 Base.propertynames(itr::DataFrameRows, private::Bool=false) = propertynames(parent(itr))
 
+"""
+    Iterators.partition(dfr::DataFrameRows, n::Integer)
+
+Iterate over `DataFrameRows` `dfr` `n` rows at a time, returning each block
+as a `DataFrameRows` over a view of rows of parent of `dfr`.
+
+# Examples
+
+```jldoctest
+julia> collect(Iterators.partition(eachrow(DataFrame(x=1:5)), 2))
+3-element Vector{DataFrames.DataFrameRows{SubDataFrame{DataFrame, DataFrames.Index, UnitRange{Int64}}}}:
+ 2×1 DataFrameRows
+ Row │ x     
+     │ Int64 
+─────┼───────
+   1 │     1
+   2 │     2
+ 2×1 DataFrameRows
+ Row │ x     
+     │ Int64 
+─────┼───────
+   1 │     3
+   2 │     4
+ 1×1 DataFrameRows
+ Row │ x     
+     │ Int64 
+─────┼───────
+   1 │     5
+```
+"""
+function Iterators.partition(dfr::DataFrameRows, n::Integer)
+    n < 1 && throw(ArgumentError("cannot create partitions of length $n"))
+    return Iterators.PartitionIterator(dfr, Int(n))
+end
+
+# use autodetection of eltype
+Base.IteratorEltype(::Type{<:Iterators.PartitionIterator{<:DataFrameRows}}) =
+    Base.EltypeUnknown()
+
+# we do not need to be overly specific here as we rely on autodetection of eltype
+# this method is needed only to override the fallback for `PartitionIterator`
+Base.eltype(::Type{<:Iterators.PartitionIterator{<:DataFrameRows}}) =
+    DataFrameRows
+
+Base.IteratorSize(::Type{<:Iterators.PartitionIterator{<:DataFrameRows}}) =
+    Base.HasLength()
+
+function Base.length(itr::Iterators.PartitionIterator{<:DataFrameRows})
+    l = nrow(parent(itr.c))
+    return cld(l, itr.n)
+end
+
+function Base.iterate(itr::Iterators.PartitionIterator{<:DataFrameRows}, state::Int=1)
+    df = parent(itr.c)
+    last_idx = nrow(df)
+    state > last_idx && return nothing
+    r = min(state + itr.n - 1, last_idx)
+    return eachrow(view(df, state:r, :)), r + 1
+end
+
 # Iteration by columns
 
 const DATAFRAMECOLUMNS_DOCSTR = """

test/dataframe.jl

@@ -2311,13 +2311,34 @@ end
         @test all(v -> v isa SubDataFrame, res)
         @test_throws ArgumentError Iterators.partition(df, false)
         @test_throws ArgumentError Iterators.partition(df, -1)
+
+        dfr = eachrow(df)
+        p = Iterators.partition(dfr, 2)
+        @test p isa Iterators.PartitionIterator
+        @test Tables.partitions(p) === p
+        @test eltype(p) === DataFrames.DataFrameRows
+        @test Base.IteratorEltype(typeof(p)) === Base.EltypeUnknown()
+        @test length(p) == 3
+        @test Base.IteratorSize(typeof(p)) === Base.HasLength()
+        res = collect(p)
+        @test res == eachrow.([DataFrame(x=1:2), DataFrame(x=3:4), DataFrame(x=5)])
+        @test all(v -> v isa DataFrames.DataFrameRows, 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
+
+    p = Iterators.partition(eachrow(DataFrame()), 1)
+    @test p isa Iterators.PartitionIterator
+    @test Tables.partitions(p) === p
+    @test isempty(p)
+    @test length(p) == 0
+    @test eltype(collect(p)) <: DataFrames.DataFrameRows
 end
 
 end # module