Make indexing of eachrow return the object of the same type on a view…

… of the parent (#3037)

NEWS.md

@@ -38,9 +38,12 @@
 
 ## Performance
 
-* Speed up `permute!` and `invpermute!` (and therefore sorting) 2x-8x 
+* Speed up `permute!` and `invpermute!` (and therefore sorting) 2x-8x
   for large tables by using cycle notation
   ([#3035](https://github.com/JuliaData/DataFrames.jl/pull/3035))
+* Make one-dimensional multi-element indexing of `DataFrameRows` return
+  `DataFrameRows`
+  ([#3037](https://github.com/JuliaData/DataFrames.jl/pull/3037))
 
 # DataFrames.jl v1.3.2 Patch Release Notes
 

src/abstractdataframe/iteration.jl

@@ -78,6 +78,8 @@ Base.IndexStyle(::Type{<:DataFrameRows}) = Base.IndexLinear()
 Base.size(itr::DataFrameRows) = (size(parent(itr), 1), )
 
 Base.@propagate_inbounds Base.getindex(itr::DataFrameRows, i::Int) = parent(itr)[i, :]
+Base.@propagate_inbounds Base.getindex(itr::DataFrameRows, idx) =
+    eachrow(@view parent(itr)[idx, :])
 
 # separate methods are needed due to dispatch ambiguity
 Base.getproperty(itr::DataFrameRows, col_ind::Symbol) =

test/iteration.jl

@@ -200,4 +200,27 @@ end
     @test findall(col -> eltype(col) <: Int, cols) == [1, 3]
 end
 
+@testset "multirow indexing of DataFrameRows" begin
+    df = DataFrame(a=1:4, b=11:14)
+    er = eachrow(df)
+    for sel in (1:2, Not(3:4), [true, true, false, false])
+        er2 = er[sel]
+        @test er2 isa DataFrames.DataFrameRows
+        @test parent(er2) isa SubDataFrame
+        @test parent(parent(er2)) === df
+        @test collect(er2) == [er[1], er[2]]
+    end
+    er2 = er[:]
+    @test er2 == er
+    @test er2 isa DataFrames.DataFrameRows
+    er2 = er[2:1]
+    @test length(er2) == 0
+    @test isempty(parent(er2))
+
+    # this is still allowed
+    er2 = er[1:2, 1]
+    @test er2 isa Vector{DataFrameRow}
+    @test er2 == er[1:2]
+end
+
 end # module