Add setindex! for ScaledArray and support missing values (#27)

Author: junyuan-chen

src/DiffinDiffsBase.jl

@@ -22,6 +22,7 @@ using Tables: AbstractColumns, istable, columnnames, getcolumn
 
 import Base: ==, +, -, *, isless, show, parent, view, diff
 import Base: eltype, firstindex, lastindex, getindex, iterate, length, sym_in
+import Missings: allowmissing, disallowmissing
 import StatsBase: coef, vcov, confint, nobs, dof_residual, responsename, coefnames, weights,
     coeftable
 import StatsModels: concrete_term, schema, termvars, lag, lead

src/ScaledArrays.jl

@@ -11,26 +11,31 @@ end
 An array type that stores data as indices of a range.
 
 # Fields
-- `refs::RA<:AbstractArray{<:Any, N}`: an array of indices.
-- `pool::P<:AbstractRange{T}`: a range that covers all possible values stored by the array.
+- `refs::RA<:AbstractArray{R,N}`: an array of indices.
+- `pool::P<:AbstractRange`: a range that covers all possible values stored by the array.
 - `invpool::Dict{T,R}`: a map from array elements to indices of `pool`.
 """
 mutable struct ScaledArray{T,R,N,RA,P} <: AbstractArray{T,N}
     refs::RA
     pool::P
     invpool::Dict{T,R}
-    ScaledArray{T,R,N,RA,P}(rs::RefArray{RA}, pool::P, invpool::Dict{T,R}) where
-        {T, R, N, RA<:AbstractArray{R, N}, P<:AbstractRange{T}} =
-            new{T,R,N,RA,P}(rs.a, pool, invpool)
+    function ScaledArray{T,R,N,RA,P}(rs::RefArray{RA}, pool::P, invpool::Dict{T,R}) where
+            {T, R, N, RA<:AbstractArray{R, N}, P<:AbstractRange}
+        eltype(P) == nonmissingtype(T) || throw(ArgumentError(
+            "expect element type of pool being $(nonmissingtype(T)); got $(eltype(P))"))
+        return new{T,R,N,RA,P}(rs.a, pool, invpool)
+    end
 end
 
-ScaledArray(rs::RefArray{RA}, pool::P, invpool::Dict{T,R}) where
-    {T,R,RA<:AbstractArray{R},P} = ScaledArray{T,R,ndims(RA),RA,P}(rs, pool, invpool)
+ScaledArray(rs::RefArray{RA}, pool::P, invpool::Dict{T,R}) where {T,R,RA<:AbstractArray{R},P} =
+    ScaledArray{T,R,ndims(RA),RA,P}(rs, pool, invpool)
 
 const ScaledVector{T,R} = ScaledArray{T,R,1}
 const ScaledMatrix{T,R} = ScaledArray{T,R,2}
 
-scale(sa::ScaledArray) = step(sa.pool)
+const ScaledArrOrSub = Union{ScaledArray, SubArray{<:Any, <:Any, <:ScaledArray}}
+
+scale(sa::ScaledArrOrSub) = step(DataAPI.refpool(sa))
 
 function _validmin(min, xmin, isstart::Bool)
     if min === nothing
@@ -178,6 +183,14 @@ ScaledArray(sa::ScaledArray, step=nothing; reftype::Type=eltype(refarray(sa)),
 Base.size(sa::ScaledArray) = size(sa.refs)
 Base.IndexStyle(::Type{<:ScaledArray{T,R,N,RA}}) where {T,R,N,RA} = IndexStyle(RA)
 
+Base.similar(sa::ScaledArray{T,R}, dims::Dims=size(sa)) where {T,R} =
+    ScaledArray(RefArray(ones(R, dims)), DataAPI.refpool(sa), Dict{T,R}())
+
+Base.similar(sa::SubArray{<:Any, <:Any, <:ScaledArray{T,R}}, dims::Dims=size(sa)) where {T,R} =
+    ScaledArray(RefArray(ones(R, dims)), DataAPI.refpool(sa), Dict{T,R}())
+
+Base.similar(sa::ScaledArrOrSub, dims::Int...) = similar(sa, dims)
+
 DataAPI.refarray(sa::ScaledArray) = sa.refs
 DataAPI.refvalue(sa::ScaledArray, n::Integer) = getindex(DataAPI.refpool(sa), n)
 DataAPI.refpool(sa::ScaledArray) = sa.pool
@@ -202,19 +215,51 @@ DataAPI.invrefpool(ssa::SubArray{<:Any, <:Any, <:ScaledArray}) =
     return @inbounds pool[n]
 end
 
-@inline function Base.getindex(sa::ScaledArray, I...)
-    refs = DataAPI.refarray(sa)
-    @boundscheck checkbounds(refs, I...)
-    @inbounds ns = refs[I...]
+Base.@propagate_inbounds function Base.getindex(sa::ScaledArrOrSub, I::AbstractVector)
+    newrefs = DataAPI.refarray(sa)[I]
     pool = DataAPI.refpool(sa)
-    N = length(pool)
-    @boundscheck checkindex(Bool, 0:N, ns) || throw_boundserror(pool, ns)
-    return @inbounds pool[ns]
+    invpool = DataAPI.invrefpool(sa)
+    return ScaledArray(RefArray(newrefs), pool, invpool)
+end
+
+Base.@propagate_inbounds function Base.setindex!(sa::ScaledArray, val, ind::Int)
+    invpool = DataAPI.invrefpool(sa)
+    n = get(invpool, val, nothing)
+    if n === nothing
+        pool = DataAPI.refpool(sa)
+        r = first(pool):step(pool):val
+        last(r) > last(pool) && (sa.pool = r)
+        n = length(r)
+        invpool[val] = n
+    end
+    refs = DataAPI.refarray(sa)
+    refs[ind] = n
+    return sa
+end
+
+Base.@propagate_inbounds function Base.setindex!(sa::ScaledArray, ::Missing, ind::Int)
+    refs = DataAPI.refarray(sa)
+    z = zero(eltype(refs))
+    invpool = DataAPI.invrefpool(sa)
+    invpool[missing] = z
+    refs[ind] = z
+    return sa
+end
+
+allowmissing(sa::ScaledArray{T,R}) where {T,R} =
+    ScaledArray(RefArray(sa.refs), sa.pool, convert(Dict{Union{T,Missing},R}, sa.invpool))
+
+function disallowmissing(sa::ScaledArray{T,R}) where {T,R}
+    T1 = nonmissingtype(T)
+    any(x->iszero(x), sa.refs) && throw(ArgumentError("cannot convert missing to $T1"))
+    delete!(sa.invpool, missing)
+    return ScaledArray(RefArray(sa.refs), sa.pool, convert(Dict{T1,R}, sa.invpool))
 end
 
-function Base.:(==)(x::ScaledArray, y::ScaledArray)
+function Base.:(==)(x::ScaledArrOrSub, y::ScaledArrOrSub)
     size(x) == size(y) || return false
-    first(x.pool) == first(y.pool) && step(x.pool) == step(y.pool) && return x.refs == y.refs
+    first(DataAPI.refpool(x)) == first(DataAPI.refpool(y)) &&
+        scale(x) == scale(y) && return DataAPI.refarray(x) == DataAPI.refarray(y)
     eq = true
     for (p, q) in zip(x, y)
         # missing could arise

src/operations.jl

@@ -94,7 +94,15 @@ function cellrows(cols::VecColumnTable, refrows::IdDict)
     ncol = length(cols)
     ncell = length(refrows)
     rows = Vector{Vector{Int}}(undef, ncell)
-    columns = AbstractVector[Vector{eltype(c)}(undef, ncell) for c in cols]
+    columns = Vector{AbstractVector}(undef, ncol)
+    for i in 1:ncol
+        c = cols[i]
+        if typeof(c) <: ScaledArray || typeof(c) <: SubArray{<:Any,1,<:ScaledArray}
+            columns[i] = similar(c, ncell)
+        else
+            columns[i] = Vector{eltype(c)}(undef, ncell)
+        end
+    end
     refs = Vector{keytype(refrows)}(undef, ncell)
     r = 0
     @inbounds for (k, v) in refrows

src/tables.jl

@@ -96,7 +96,7 @@ end
 
 Base.values(cols::VecColumnTable) = _columns(cols)
 Base.haskey(cols::VecColumnTable, key::Symbol) = haskey(_lookup(cols), key)
-Base.haskey(cols::VecColumnTable, i::Int) = 0 < i < length(_names(cols))
+Base.haskey(cols::VecColumnTable, i::Int) = 0 < i <= length(_names(cols))
 
 function Base.:(==)(x::VecColumnTable, y::VecColumnTable)
     size(x) == size(y) || return false

test/ScaledArrays.jl

@@ -50,24 +50,51 @@ using DiffinDiffsBase: RefArray, validpool, scaledlabel
     @test sa5.refs == sa4.refs
     @test sa5.refs !== sa4.refs
 
+    sa6 = ScaledArray([missing, 1, 2], 1)
+    @test eltype(sa6) == Union{Int,Missing}
+    @test sa6.refs == 0:2
+    @test eltype(sa6.refs) == Int32
+    @test ismissing(sa6[1])
+
     @test_throws ArgumentError ScaledArray(sa5, stop=7, usepool=false)
+    @test_throws ArgumentError ScaledArray(RefArray(1:3), 1.0:0.5:3.0, Dict{Int,Int}())
 
     @test size(sa) == (10,)
     @test IndexStyle(typeof(sa)) == IndexLinear()
 
+    @test similar(sa, 4) == ScaledArray(RefArray(ones(Int, 4)), sa.pool, Dict{Date,Int}())
+    @test similar(sa, (4,)) == similar(view(sa, 1:5), 4) == similar(sa, 4)
+
     @test refarray(sa) === sa.refs
     @test refvalue(sa, 1) == Date(1)
     @test refpool(sa) === sa.pool
     @test invrefpool(sa) === sa.invpool
 
-    ssa = view(sa, 3:4)
-    @test refarray(ssa) == view(sa.refs, 3:4)
+    ssa = view(sa, 3:5)
+    @test refarray(ssa) == view(sa.refs, 3:5)
     @test refvalue(ssa, 1) == Date(1)
     @test refpool(ssa) === sa.pool
     @test invrefpool(ssa) === sa.invpool
+    @test ssa == sa[3:5]
 
     @test sa[1] == Date(1)
     @test sa[1:2] == sa[[1,2]] == sa[(1:10).<3] == Date.(1:2)
+    @test sa[1:2] isa ScaledArray
+    @test ssa[1:2] isa ScaledArray
+
+    sa[1] = Date(10)
+    @test sa[1] == Date(10)
+    sa[1:2] .= Date(100)
+    @test sa[1:2] == [Date(100), Date(100)]
+    @test last(sa.pool) == Date(100)
+    @test_throws MethodError sa[1] = missing
+    sa = allowmissing(sa)
+    sa[1] = missing
+    @test ismissing(sa[1])
+    @test_throws ArgumentError disallowmissing(sa)
+    sa[1] = Date(1)
+    sa = disallowmissing(sa)
+    @test eltype(sa) == Date
 end
 
 @testset "scaledlabel" begin

test/operations.jl

@@ -42,6 +42,16 @@ end
         sortslices(unique(hcat(hrs.wave, hrs.wave_hosp), dims=1), dims=1)
     @test propertynames(cells) == [:wave, :wave_hosp]
     @test rows[1] == intersect(findall(x->x==7, hrs.wave), findall(x->x==8, hrs.wave_hosp))
+
+    df = DataFrame(hrs)
+    df.wave = ScaledArray(df.wave, 1)
+    df.wave_hosp = ScaledArray(df.wave_hosp, 1)
+    cols1 = subcolumns(df, (:wave, :wave_hosp))
+    cells1, rows1 = cellrows(cols1, rows_dict)
+    @test rows1 == rows
+    @test cells1 == cells
+    @test cells1.wave isa ScaledArray
+    @test cells1.wave_hosp isa ScaledArray
 end
 
 @testset "settime" begin

test/runtests.jl

@@ -11,6 +11,7 @@ using DiffinDiffsBase: @fieldequal, unpack, @unpack, checktable, hastreat, parse
     _totermset!, parse_didargs!, _treatnames, _parse_bycells!, _parse_subset, _nselected,
     treatindex, checktreatindex
 using LinearAlgebra: Diagonal
+using Missings: allowmissing, disallowmissing
 using PooledArrays: PooledArray
 using StatsBase: Weights, UnitWeights
 using StatsModels: termvars