add columnindex for DataFrameRow, speed up columnindex, fix a bug in …

…indexing (#2380)

NEWS.md

@@ -40,7 +40,9 @@
   inserted in `insertcols!` ([#2365](https://github.com/JuliaData/DataFrames.jl/pull/2365))
 * add `isapprox` method to check for approximate equality between two dataframes
   ([#2373](https://github.com/JuliaData/DataFrames.jl/pull/2373))
-
+* add `columnindex` for `DataFrameRow`
+  ([#2380](https://github.com/JuliaData/DataFrames.jl/pull/2380))
+
 ## Deprecated
 
 * `DataFrame!` is now deprecated ([#2338](https://github.com/JuliaData/DataFrames.jl/pull/2338))

src/other/index.jl

@@ -364,7 +364,9 @@ Base.@propagate_inbounds function parentcols(ind::SubIndex, idx::Symbol)
     parentcol = ind.parent[idx]
     @boundscheck begin
         remap = ind.remap
-        remap[parentcol] == 0 && throw(ArgumentError("$idx not found"))
+        if parentcol > length(remap) || remap[parentcol] <= 0
+            throw(ArgumentError("$idx not found"))
+        end
     end
     return parentcol
 end

src/other/tables.jl

@@ -6,7 +6,19 @@ Tables.rows(df::AbstractDataFrame) = eachrow(df)
 Tables.rowtable(df::AbstractDataFrame) = Tables.rowtable(Tables.columntable(df))
 Tables.namedtupleiterator(df::AbstractDataFrame) =
     Tables.namedtupleiterator(Tables.columntable(df))
-Tables.columnindex(df::AbstractDataFrame, idx::AbstractString) =
+
+function Tables.columnindex(df::Union{AbstractDataFrame, DataFrameRow}, idx::Symbol)
+    ind = index(df)
+    if ind isa Index
+        return get(ind.lookup, idx, 0)
+    else
+        parent_ind = ind.parent
+        loc = get(parent_ind.lookup, idx, 0)
+        return loc == 0 || loc > length(ind.remap) ? 0 : max(0, ind.remap[loc])
+    end
+end
+
+Tables.columnindex(df::Union{AbstractDataFrame, DataFrameRow}, idx::AbstractString) =
     columnindex(df, Symbol(idx))
 
 Tables.schema(df::AbstractDataFrame) = Tables.Schema(propertynames(df), eltype.(eachcol(df)))

test/indexing.jl

@@ -1822,6 +1822,19 @@ end
         @test_throws MethodError df[1, 2:3] = DataFrame([11 12])
         @test_throws MethodError df[1, [false, true, true]] = DataFrame([11 12])
     end
+
+    @testset "cornercase of view indexing" begin
+        df = DataFrame(reshape(1:12, 4, :))
+        dfr = df[1, 3:2]
+        for idx in [:x1, :x2, :x3, :x4]
+            @test_throws ArgumentError dfr[idx]
+        end
+
+        sdf = @view df[1:1, 3:2]
+        for idx in [:x1, :x2, :x3, :x4]
+            @test_throws ArgumentError sdf[1, idx]
+        end
+    end
 end
 
 end # module

test/tables.jl

@@ -197,10 +197,23 @@ end
 
 @testset "columnindex" begin
     df = DataFrame(rand(3,4))
-    @test columnindex.(Ref(df), names(df)) == 1:4
-    @test columnindex(df, :a) == 0
-    # @test_throws ErrorException columnindex(df, 1)
-    # @test_throws ErrorException columnindex(df, "x1")
+
+    for x in (df, view(df, 1, :), view(df, 1:1, :))
+        @test columnindex.(Ref(x), names(df)) == 1:4
+        @test columnindex.(Ref(x), propertynames(df)) == 1:4
+        @test columnindex(x, :a) == 0
+        @test columnindex(x, "a") == 0
+        @test_throws MethodError columnindex(x, 1)
+    end
+
+    for x in (view(df, 1, 4:3), view(df, 1:1, 4:3))
+        @test all(==(0), columnindex.(Ref(x), [names(df); "a"]))
+        @test all(==(0), columnindex.(Ref(x), [propertynames(df); :a]))
+    end
+    for x in (view(df, 1, [4, 3]), view(df, 1:1, [4, 3]))
+        @test columnindex.(Ref(x), [names(df); "a"]) == [0, 0, 2, 1, 0]
+        @test columnindex.(Ref(x), [propertynames(df); :a]) == [0, 0, 2, 1, 0]
+    end
 end
 
 @testset "eachrow and eachcol integration" begin