Update DictRow: add row index as a field for rownumber

JuliaData · Apr 23, 2023 · 5f091b0 · 5f091b0
1 parent aefedf6
commit 5f091b0
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/src/dicts.jl b/src/dicts.jl
@@ -108,12 +108,13 @@ schema(x::DictRowTable) = Schema(getfield(x, :names), [getfield(x, :types)[nm] f
 struct DictRow <: AbstractRow
     names::Vector{Symbol}
     row::Dict{Symbol, Any}
+    rownumber::Int
 end
 
-rownumber(x::DictRow) = getfield(x, :row)
+rownumber(x::DictRow) = getfield(x, :rownumber)
 columnnames(x::DictRow) = getfield(x, :names)
-getcolumn(x::DictRow, i::Int) = get(rownumber(x), columnnames(x)[i], missing)
-getcolumn(x::DictRow, nm::Symbol) = get(rownumber(x), nm, missing)
+getcolumn(x::DictRow, i::Int) = get(getfield(x, :row), columnnames(x)[i], missing)
+getcolumn(x::DictRow, nm::Symbol) = get(getfield(x, :row), nm, missing)
 
 Base.IteratorSize(::Type{DictRowTable}) = Base.HasLength()
 Base.length(x::DictRowTable) = length(getfield(x, :values))
@@ -122,7 +123,7 @@ Base.eltype(::Type{DictRowTable}) = DictRow
 
 function Base.iterate(x::DictRowTable, st=1)
     st > length(x) && return nothing
-    return DictRow(x.names, x.values[st]), st + 1
+    return DictRow(x.names, x.values[st], st), st + 1
 end
 
 function subset(x::DictRowTable, inds; viewhint::Union{Bool,Nothing}=nothing, view::Union{Bool,Nothing}=nothing)
@@ -132,7 +133,7 @@ function subset(x::DictRowTable, inds; viewhint::Union{Bool,Nothing}=nothing, vi
     end
     values = viewhint === true ? Base.view(getfield(x, :values), inds) : getfield(x, :values)[inds]
     if inds isa Integer
-        return DictRow(getfield(x, :names), values)
+        return DictRow(getfield(x, :names), values, inds)
     else
         values isa AbstractVector || throw(ArgumentError("`Tables.subset`: invalid `inds` argument, expected `RowTable` output, got $(typeof(ret))"))
         return DictRowTable(getfield(x, :names), getfield(x, :types), values)

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -83,7 +83,8 @@ using Test, Tables, OrderedCollections, TableTraits, DataValues, QueryOperators,
     @test Tables.columnaccess(rows)
     @test Tables.columns(rows) === nt
     @test Tables.materializer(rows) === Tables.materializer(nt)
-
+    @test Tables.rownumber(Tables.IteratorRow(row)) == 1
+
     @test Tables.rowmerge(row; b=200, hey="hello") == (a=1, b=200, hey="hello")
     @test Tables.rowmerge(row, (hey="hello", a=200)) == (a=200, b=4, hey="hello")
     @test Tables.rowmerge(row, (hey="hello", a=200), row, (x=:x, y=:y, hey="bye")) == (a=1, b=4, hey="bye", x=:x, y=:y)
@@ -826,6 +827,10 @@ end
     @test isequal(ct.c, [3, missing, missing, 10, 10])
     @test isequal(ct.d, [missing, 5, 7, missing, 11])
 
+    for (i,row) in enumerate(drt)
+        @test Tables.rownumber(row) == i
+    end
+
     dct = Tables.dictcolumntable(rt)
     @test isequal(dct.a, [1, missing, 6, 8, 8])
     @test isequal(ct.b, Union{Int, Float64, Missing}[2, 4.0, missing, 9, 9])