Fix performance issues with specset (#11)

* Determine equality of StatsStep arguments by ===

* Add copyargs and avoid unnecessary copying in proceed

Author: junyuan-chen

Project.toml

@@ -9,7 +9,6 @@ Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Missings = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
-SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsModels = "3eaba693-59b7-5ba5-a881-562e759f1c8d"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
@@ -21,7 +20,6 @@ DataFrames = "0.22"
 MacroTools = "0.5"
 Missings = "0.4"
 Reexport = "0.2, 1"
-SplitApplyCombine = "1.1"
 StatsBase = "0.33"
 StatsModels = "0.6.18"
 Tables = "1.2"

src/DiffinDiffsBase.jl

@@ -5,7 +5,6 @@ using Combinatorics: combinations
 using MacroTools: @capture, isexpr, postwalk
 using Missings: disallowmissing
 using Reexport
-using SplitApplyCombine: groupfind, groupview
 using StatsBase: Weights, uweights
 @reexport using StatsModels
 using Tables: istable, getcolumn, columntable, columnnames

src/StatsProcedures.jl

@@ -94,16 +94,14 @@ See also [`proceed`](@ref).
 """
 combinedargs(::StatsStep, ::Any) = ()
 
+copyargs(::StatsStep) = ()
+
 function (step::StatsStep{A,F})(@nospecialize(ntargs::NamedTuple);
         verbose::Bool=false) where {A,F}
     haskey(ntargs, :verbose) && (verbose = ntargs.verbose)
     verbose && printstyled("Running ", step, "\n", color=:green)
     ret = F.instance(groupargs(step, ntargs)..., combinedargs(step, (ntargs,))...)
-    if ret isa Tuple{<:NamedTuple, Bool}
-        return merge(ntargs, ret[1])
-    else
-        error("unexpected $(typeof(ret)) returned from $(F.name.mt.name) associated with StatsStep $A")
-    end
+    return merge(ntargs, ret)
 end
 
 (step::StatsStep)(; verbose::Bool=false) = step(NamedTuple(), verbose=verbose)
@@ -197,6 +195,7 @@ _sharedby(s::SharedStatsStep) = s.ids
 _f(s::SharedStatsStep) = _f(s.step)
 groupargs(s::SharedStatsStep, @nospecialize(ntargs::NamedTuple)) = groupargs(s.step, ntargs)
 combinedargs(s::SharedStatsStep, v::AbstractArray) = combinedargs(s.step, v)
+copyargs(s::SharedStatsStep) = copyargs(s.step)
 
 ==(x::SharedStatsStep, y::SharedStatsStep) =
     x.step == y.step && x.ids == y.ids
@@ -320,6 +319,10 @@ function pool(ps::AbstractStatsProcedure...)
     return PooledStatsProcedure(ps, shared)
 end
 
+# A shortcut for the simple case
+pool(p::AbstractStatsProcedure) =
+    PooledStatsProcedure((p,), [SharedStatsStep(s, 1) for s in p])
+
 length(p::PooledStatsProcedure) = length(p.steps)
 eltype(::Type{PooledStatsProcedure}) = SharedStatsStep
 firstindex(p::PooledStatsProcedure) = firstindex(p.steps)
@@ -439,6 +442,11 @@ function show(io::IO, ::MIME"text/plain", sp::StatsSpec{T}) where T
     _show_args(io, sp)
 end
 
+function _count!(objcount::IdDict, obj)
+    count = get(objcount, obj, 0)
+    objcount[obj] = count + 1
+end
+
 """
     proceed(sps::AbstractVector{<:StatsSpec}; kwargs...)
 
@@ -464,50 +472,65 @@ function proceed(sps::AbstractVector{<:StatsSpec};
         verbose::Bool=false, keep=nothing, keepall::Bool=false)
     nsps = length(sps)
     nsps == 0 && throw(ArgumentError("expect a nonempty vector"))
+
+    gids = IdDict{AbstractStatsProcedure, Vector{Int}}()
+    objcount = IdDict{Any, Int}()
     traces = Vector{NamedTuple}(undef, nsps)
-    for i in 1:nsps
-        traces[i] = sps[i].args
+    @inbounds for i in 1:nsps
+        push!(get!(Vector{Int}, gids, _procedure(sps[i])()), i)
+        args = sps[i].args
+        foreach(x->_count!(objcount, x), args)
+        traces[i] = args
     end
-    gids = groupfind(r->_procedure(r)(), sps)
+    
     steps = pool((p for p in keys(gids))...)
+    tasks = IdDict{Tuple, Vector{Int}}()
     ntask_total = 0
-    for step in steps
+    @inbounds for step in steps
         ntask = 0
-        verbose && printstyled("Running ", step, "...")
-        taskids = vcat((gids[steps.procs[i]] for i in _sharedby(step))...)
-        tasks = groupview(r->groupargs(step, r), view(traces, taskids))
-        for (ins, subtb) in pairs(tasks)
-            ret = _f(step)(ins..., combinedargs(step, subtb)...)
-            if ret isa Tuple{<:NamedTuple, Bool}
-                ret, share = ret
-            else
-                fname = typeof(_f(step)).name.mt.name
-                stepname = typeof(step.step).parameters[1]
-                error("unexpected type $(typeof(ret)) of object returned from $fname associated with StatsStep $stepname")
+        verbose && print("Running ", step, "...")
+        # Group arguments by ===
+        for i in _sharedby(step)
+            taskids = gids[steps.procs[i]]
+            for j in taskids
+                push!(get!(Vector{Int}, tasks, groupargs(step, traces[j])), j)
             end
-            ntask += 1
-            ntask_total += 1
-            if share
-                for i in eachindex(subtb)
-                    subtb[i] = merge(subtb[i], ret)
-                end
-            else
-                for i in eachindex(subtb)
-                    subtb[i] = merge(subtb[i], deepcopy(ret))
+        end
+
+        for (gargs, ids) in tasks
+            # Handle potential in-place operations on mutable objects
+            nids = length(ids)
+            icopy = copyargs(step)
+            if length(icopy) > 0
+                gargs = Any[gargs...]
+                for i in copyargs(step)
+                    a = gargs[i]
+                    objcount[a] > nids && (gargs[i] = deepcopy(a))
                 end
             end
+
+            ret = _f(step)(gargs..., combinedargs(step, view(traces, ids))...)
+            for id in ids
+                foreach(x->_count!(objcount, x), ret)
+                traces[id] = merge(traces[id], ret)
+            end
         end
+        ntask = length(tasks)
+        ntask_total += ntask
+        empty!(tasks)
         nprocs = length(_sharedby(step))
-        verbose && printstyled("Finished ", ntask, ntask > 1 ? " tasks" : " task", " for ",
+        verbose && print("Finished ", ntask, ntask > 1 ? " tasks" : " task", " for ",
             nprocs, nprocs > 1 ? " procedures\n" : " procedure\n")
     end
+
     nprocs = length(steps.procs)
     verbose && printstyled("All steps finished (", ntask_total,
         ntask_total > 1 ? " tasks" : " task", " for ", nprocs,
         nprocs > 1 ? " procedures)\n" : " procedure)\n", bold=true, color=:green)
-    for i in 1:nsps
+    @inbounds for i in 1:nsps
         traces[i] = result(_procedure(sps[i]), traces[i])
     end
+
     if keepall
         return traces
     elseif keep === nothing
@@ -517,11 +540,11 @@ function proceed(sps::AbstractVector{<:StatsSpec};
         if keep isa Symbol
             keep = (keep,)
         else
-            eltype(keep) == Symbol ||
-                throw(ArgumentError("expect Symbol or collections of Symbols for the value of option `keep`"))
+            eltype(keep) == Symbol || throw(ArgumentError(
+                "expect Symbol or collections of Symbols for the value of option `keep`"))
         end
         in(:result, keep) || (keep = (keep..., :result))
-        for i in 1:nsps
+        @inbounds for i in 1:nsps
             names = ((n for n in keep if haskey(traces[i], n))...,)
             traces[i] = NamedTuple{names}(traces[i])
         end

src/procedures.jl

@@ -25,7 +25,7 @@ function checkdata(data, subset::Union{AbstractVector, Nothing},
     end
 
     sum(esample) == 0 && error("no nonmissing data")
-    return (esample=esample,), false
+    return (esample=esample,)
 end
 
 """
@@ -97,7 +97,7 @@ function checkvars!(data, tr::AbstractTreatment, pr::AbstractParallel,
 
     overlap!(esample, tr_rows, tr, pr, treatname, data)
     sum(esample) == 0 && error("no nonmissing data")
-    return (esample=esample, tr_rows=tr_rows), false
+    return (esample=esample, tr_rows=tr_rows)
 end
 
 """
@@ -109,6 +109,7 @@ const CheckVars = StatsStep{:CheckVars, typeof(checkvars!)}
 
 required(::CheckVars) = (:data, :tr, :pr, :yterm, :treatname, :esample)
 default(::CheckVars) = (treatintterms=(), xterms=())
+copyargs(::CheckVars) = (6,)
 
 """
     makeweights(args...)
@@ -119,19 +120,18 @@ See also [`MakeWeights`](@ref).
 function makeweights(data, esample::BitVector, weightname::Symbol)
     weights = Weights(convert(Vector{Float64}, view(getcolumn(data, weightname), esample)))
     all(isfinite, weights) || error("data column $weightname contain not-a-number values")
-    (weights=weights,), true
+    return (weights=weights,)
 end
 
 function makeweights(data, esample::BitVector, weightname::Nothing)
     weights = uweights(sum(esample))
-    (weights=weights,), true
+    return (weights=weights,)
 end
 
 """
     MakeWeights <: StatsStep
 
 Call [`DiffinDiffsBase.makeweights`](@ref) to create a generic `Weights` vector.
-The returned object named `weights` may be shared across multiple specifications.
 """
 const MakeWeights = StatsStep{:MakeWeights, typeof(makeweights)}
 

test/StatsProcedures.jl

@@ -1,34 +1,31 @@
 using DiffinDiffsBase: _f, _get, groupargs, _get_default,
     _sharedby, _show_args, _args_kwargs, _parse!, pool, proceed
-import DiffinDiffsBase: required, default, transformed, combinedargs
+import DiffinDiffsBase: required, default, transformed, combinedargs, copyargs
 
-testvoidstep(a::String) = NamedTuple(), false
+testvoidstep(a::String) = NamedTuple()
 const TestVoidStep = StatsStep{:TestVoidStep, typeof(testvoidstep)}
 required(::TestVoidStep) = (:a,)
 
-testregstep(a::String, b::String) = (c=a*b,), false
+testregstep(a::String, b::String) = (c=a*b,)
 const TestRegStep = StatsStep{:TestRegStep, typeof(testregstep)}
 default(::TestRegStep) = (a="a", b="b")
 
-testlaststep(a::String, c::String) = (result=a*c,), false
+testlaststep(a::String, c::String) = (result=a*c,)
 const TestLastStep = StatsStep{:TestLastStep, typeof(testlaststep)}
 default(::TestLastStep) = (a="a",)
 transformed(::TestLastStep, ntargs::NamedTuple) = (ntargs.c,)
 
-testarraystep(a::String) = (result=[a],), false
-const TestArrayStep = StatsStep{:TestArrayStep, typeof(testarraystep)}
-default(::TestArrayStep) = (a="a",)
-
-testcombinestep(a::String, bs::String...) = (c=collect(bs),), true
+testcombinestep(a::String, bs::String...) = (c=collect(bs),)
 const TestCombineStep = StatsStep{:TestCombineStep, typeof(testcombinestep)}
 default(::TestCombineStep) = (a="a",)
 combinedargs(::TestCombineStep, ntargs) = [nt.b for nt in ntargs]
 
-testinvalidstep(a::String, b::String) = b, false
-const TestInvalidStep = StatsStep{:TestInvalidStep, typeof(testinvalidstep)}
-default(::TestInvalidStep) = (a="a", b="b")
+testarraystep(a::String, c::Array) = (result=c,)
+const TestArrayStep = StatsStep{:TestArrayStep, typeof(testarraystep)}
+required(::TestArrayStep) = (:a, :c)
+copyargs(::TestArrayStep) = (2,)
 
-const TestUnnamedStep = StatsStep{:TestUnnamedStep, typeof(testinvalidstep)}
+const TestUnnamedStep = StatsStep{:TestUnnamedStep, typeof(testregstep)}
 
 @testset "StatsStep" begin
     @testset "_get" begin
@@ -64,14 +61,14 @@ const TestUnnamedStep = StatsStep{:TestUnnamedStep, typeof(testinvalidstep)}
         @test TestLastStep()((a="a", b="a", c="ab")) ==
             (a="a", b="a", c="ab", result="aab")
 
-        @test TestArrayStep()((a="a",)) == (a="a", result=["a"])
         @test TestCombineStep()((a="a", b="b")) == (a="a", b="b", c=["b"])
+        c = ["c"]
+        ret = TestArrayStep()((a="a", c=c,))
+        @test ret.result === c
 
         @test sprint(show, TestVoidStep()) == "TestVoidStep"
         @test sprint(show, MIME("text/plain"), TestVoidStep()) ==
             "TestVoidStep (StatsStep that calls testvoidstep)"
-        
-        @test_throws ErrorException TestInvalidStep()()
     end
 end
 
@@ -86,8 +83,8 @@ const NP = TestProcedure{:NullProcedure,Tuple{}}
 const np = NP()
 const CP = TestProcedure{:CombineProcedure,Tuple{TestCombineStep,TestArrayStep}}
 const cp = CP()
-const EP = TestProcedure{:InvalidProcedure,Tuple{TestInvalidStep}}
-const ep = EP()
+const AP = TestProcedure{:ArrayProcedure,Tuple{TestArrayStep}}
+const ap = AP()
 
 @testset "AbstractStatsProcedure" begin
     @test length(rp) == 3
@@ -266,8 +263,10 @@ testformatter(nt::NamedTuple) = (haskey(nt, :name) ? nt.name : "", nt.p, (a=nt.a
     s5 = StatsSpec("s5", IP, (a="a", b="b"))
     s6 = StatsSpec("s6", CP, (a="a", b="b1"))
     s7 = StatsSpec("s7", CP, (a="a", b="b2"))
-    s8 = StatsSpec("s8", NP, NamedTuple())
-    s9 = StatsSpec("s9", EP, (a="a", b="b"))
+    c = ["c"]
+    s8 = StatsSpec("s8", AP, (a="a", c=c))
+    s9 = StatsSpec("s9", AP, (a="a1", c=c))
+    s10 = StatsSpec("s10", NP, NamedTuple())
 
     @test proceed([s1]) == ["aab"]
     @test proceed([s1,s2], verbose=true) == ["aab", "aab"]
@@ -290,19 +289,23 @@ testformatter(nt::NamedTuple) = (haskey(nt, :name) ? nt.name : "", nt.p, (a=nt.a
     @test proceed([s1,s4], keepall=true) ==
         [(a="a", b="b", c="ab", result="aab"), (a="a", b="b", c="ab")]
 
-    @test proceed([s6], keepall=true) == [(a="a", b="b1", c=["b1"], result=["a"])]
+    @test proceed([s6], keepall=true) == [(a="a", b="b1", c=["b1"], result=["b1"])]
     ret = proceed([s6,s7], keepall=true)
-    @test ret == [(a="a", b="b1", c=["b1", "b2"], result=["a"]),
-        (a="a", b="b2", c=["b1", "b2"], result=["a"])]
+    @test ret == [(a="a", b="b1", c=["b1", "b2"], result=["b1", "b2"]),
+        (a="a", b="b2", c=["b1", "b2"], result=["b1", "b2"])]
     @test ret[1].c === ret[2].c
-    @test ret[1].result !== ret[2].result
+    @test ret[1].result === ret[2].result
+
+    ret = proceed([s8], keepall=true)
+    @test ret[1].c === ret[1].result
+    ret = proceed([s8,s9], keepall=true)
+    @test ret[1].c !== ret[1].result
 
-    @test proceed([s8]) == [nothing]
-    @test proceed([s8], keepall=true) == NamedTuple[NamedTuple()]
-    @test proceed([s8], keep=:result) == NamedTuple[NamedTuple()]
+    @test proceed([s10]) == [nothing]
+    @test proceed([s10], keepall=true) == NamedTuple[NamedTuple()]
+    @test proceed([s10], keep=:result) == NamedTuple[NamedTuple()]
 
     @test_throws ArgumentError proceed(StatsSpec[])
-    @test_throws ErrorException proceed([s9])
 end
 
 @testset "_parse!" begin