Skip to content

Avoid recompiling the term walk for every table type #355

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
matthieugomez wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Avoid recompiling the term walk for every table type #355

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/errormessages.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ end
"""
Return a nice-ish error message if the Symbol `name` isn't a column name in `table`, otherwise a zero-length string.
"""
function checkcol(table, name::Symbol)
function checkcol(@nospecialize(table), name::Symbol)
i = Tables.columnindex(table, name)
if i == 0 # if no such column
names = Tables.columnnames(table)
Expand All @@ -36,7 +36,7 @@ end
Check that each name in the given model `f` exists in the data source `t` and return a message if not. Return a zero string otherwise.
`t` is something that implements the `Tables` interface.
"""
function checknamesexist(f::FormulaTerm, t)
function checknamesexist(f::FormulaTerm, @nospecialize(t))
if ! Tables.istable(t)
throw(ArgumentError( "$(typeof(t)) isn't a valid Table type" ))
end
Expand Down
38 changes: 27 additions & 11 deletions src/modelframe.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,24 +53,40 @@ end
_missing_omit(x::AbstractVector{T}) where T = copyto!(similar(x, nonmissingtype(T)), x)
_missing_omit(x::AbstractVector, rows) = _missing_omit(view(x, rows))

function missing_omit(d::T) where T<:ColumnTable
# loops instead of `map` over the `NamedTuple` so nothing is specialized on
# the table type (see `_modelcols_each` in terms.jl)
function missing_omit(@nospecialize(d::ColumnTable))
nonmissings = trues(length(first(d)))
for col in d
for col in values(d)
_nonmissing!(nonmissings, col)
end
d_nonmissing = if all(nonmissings)
map(_missing_omit, d)
cols = Vector{Any}(undef, length(d))
if all(nonmissings)
for (i, col) in enumerate(values(d))
cols[i] = _missing_omit(col)
end
else
rows = findall(nonmissings)
map(Base.Fix2(_missing_omit, rows), d)
for (i, col) in enumerate(values(d))
cols[i] = _missing_omit(col, rows)
end
end
d_nonmissing = NamedTuple{keys(d)}(Tuple(cols))
d_nonmissing, nonmissings
end

missing_omit(data::T, formula::AbstractTerm) where T<:ColumnTable =
missing_omit(NamedTuple{tuple(termvars(formula)...)}(data))
function missing_omit(@nospecialize(data::ColumnTable), formula::AbstractTerm)
vars = termvars(formula)
# equivalent to NamedTuple{Tuple(vars)}(data), without Base's generator over
# the table type
cols = Vector{Any}(undef, length(vars))
for (i, v) in enumerate(vars)
cols[i] = getfield(data, v)
end
missing_omit(NamedTuple{Tuple(vars)}(Tuple(cols)))
end

function ModelFrame(f::FormulaTerm, data::ColumnTable;
function ModelFrame(f::FormulaTerm, @nospecialize(data::ColumnTable);
model::Type{M}=StatisticalModel, contrasts=Dict{Symbol,Any}()) where M

msg = checknamesexist( f, data )
Expand All @@ -89,7 +105,7 @@ end
ModelFrame(f::FormulaTerm, data; model=StatisticalModel, contrasts=Dict{Symbol,Any}()) =
ModelFrame(f, columntable(data); model=model, contrasts=contrasts)

StatsAPI.modelmatrix(f::FormulaTerm, data; kwargs...) = modelmatrix(f.rhs, data; kwargs...)
StatsAPI.modelmatrix(f::FormulaTerm, @nospecialize(data); kwargs...) = modelmatrix(f.rhs, data; kwargs...)

"""
modelmatrix(t::AbstractTerm, data; hints=Dict(), mod=StatisticalModel)
Expand All @@ -113,7 +129,7 @@ keyword arguments are passed to [`apply_schema`](@ref).
[`modelcols`](@ref) pipeline directly

"""
function StatsAPI.modelmatrix(t::Union{AbstractTerm, TupleTerm}, data;
function StatsAPI.modelmatrix(t::Union{AbstractTerm, TupleTerm}, @nospecialize(data);
hints=Dict{Symbol,Any}(), mod::Type{M}=StatisticalModel) where M
Tables.istable(data) ||
throw(ArgumentError("expected data in a Table, got $(typeof(data))"))
Expand Down Expand Up @@ -141,7 +157,7 @@ keyword arguments are passed to [`apply_schema`](@ref).
[`modelcols`](@ref) pipeline directly

"""
function StatsAPI.response(f::FormulaTerm, data;
function StatsAPI.response(f::FormulaTerm, @nospecialize(data);
hints=Dict{Symbol,Any}(),
mod::Type{M}=StatisticalModel) where M
Tables.istable(data) ||
Expand Down
23 changes: 15 additions & 8 deletions src/schema.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,15 +115,22 @@ y(continuous)
```
"""
schema(data, hints=Dict{Symbol,Any}()) = schema(columntable(data), hints)
schema(dt::D, hints=Dict{Symbol,Any}()) where {D<:ColumnTable} =
schema(Term.(collect(fieldnames(D))), dt, hints)
schema(@nospecialize(dt::ColumnTable), hints=Dict{Symbol,Any}()) =
schema(Term.(collect(keys(dt))), dt, hints)
schema(ts::AbstractVector{<:AbstractTerm}, data, hints::Dict{Symbol}) =
schema(ts, columntable(data), hints)

# handle hints:
schema(ts::AbstractVector{<:AbstractTerm}, dt::ColumnTable,
hints::Dict{Symbol}=Dict{Symbol,Any}()) =
sch = Schema(t=>concrete_term(t, dt, hints) for t in ts)
# a loop rather than a generator so that no closure captures the table (see
# `_modelcols_each` in terms.jl)
function schema(ts::AbstractVector{<:AbstractTerm}, @nospecialize(dt::ColumnTable),
hints::Dict{Symbol}=Dict{Symbol,Any}())
sch = Dict{Term,AbstractTerm}()
for t in ts
sch[t] = concrete_term(t, dt, hints)
end
return Schema(sch)
end

schema(f::TermOrTerms, data, hints::Dict{Symbol}) =
schema(filter(needs_schema, terms(f)), data, hints)
Expand Down Expand Up @@ -170,15 +177,15 @@ a(continuous)
"""
concrete_term(t::Term, d, hints::Dict{Symbol}) = concrete_term(t, d, get(hints, t.sym, nothing))

function concrete_term(t::Term, dt::ColumnTable, hint)
function concrete_term(t::Term, @nospecialize(dt::ColumnTable), hint)
msg = checkcol(dt, t.sym)
if msg != ""
throw(ArgumentError(msg))
end
return concrete_term(t, getproperty(dt, t.sym), hint)
end

function concrete_term(t::Term, dt::ColumnTable, hints::Dict{Symbol})
function concrete_term(t::Term, @nospecialize(dt::ColumnTable), hints::Dict{Symbol})
msg = checkcol(dt, t.sym)
if msg != ""
throw(ArgumentError(msg))
Expand All @@ -191,7 +198,7 @@ concrete_term(t::Term, d) = concrete_term(t, d, nothing)

# if the "hint" is already an AbstractTerm, use that
# need this specified to avoid ambiguity
concrete_term(t::Term, d::ColumnTable, hint::AbstractTerm) = hint
concrete_term(t::Term, @nospecialize(d::ColumnTable), hint::AbstractTerm) = hint
concrete_term(t::Term, x, hint::AbstractTerm) = hint

# second possible fix for #97
Expand Down
2 changes: 1 addition & 1 deletion src/temporal_terms.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ end
ShiftedArrays.lead(t::T, n=1) where {T<:AbstractTerm} = LeadLagTerm{T,typeof(lead)}(t, n)
ShiftedArrays.lag(t::T, n=1) where {T<:AbstractTerm} = LeadLagTerm{T,typeof(lag)}(t, n)

function modelcols(ll::LeadLagTerm{<:Any, F}, d::Tables.ColumnTable) where F
function modelcols(ll::LeadLagTerm{<:Any, F}, @nospecialize(d::Tables.ColumnTable)) where F
original_cols = modelcols(ll.term, d)
return F.instance(original_cols, ll.nsteps)
end
Expand Down
63 changes: 42 additions & 21 deletions src/terms.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -492,22 +492,44 @@ julia> modelcols(MatrixTerm(ts), d)
9.0 0.05079 0.0 1.0
```
"""
modelcols(ts::TupleTerm, d::NamedTuple) = modelcols.(ts, Ref(d))
modelcols(ts::TupleTerm, @nospecialize(d::NamedTuple)) = Tuple(_modelcols_each(ts, d))

# Data arguments are `@nospecialize`d and term collections are walked with plain
# loops rather than comprehensions or broadcasts: a closure over `d` (as in
# `[modelcols(t, d) for t in ts]` or `modelcols.(ts, Ref(d))`) has the table's
# `NamedTuple` type baked into its own type, so every table with a new set of
# column names or types would force a fresh round of inference through the
# whole pipeline. The loop bodies dispatch at run time on the concrete column
# types instead, which are shared across tables.
function _modelcols_each(ts, @nospecialize(d))
out = Vector{Any}(undef, length(ts))
for (i, t) in enumerate(ts)
out[i] = modelcols(t, d)
end
return _narrow(out)
end
# same element type as the comprehension would have produced
_narrow(v::Vector{Any}) = [x for x in v]

modelcols(t::Term, d::NamedTuple) = getproperty(d, t.sym)
modelcols(t::ConstantTerm, d::NamedTuple) = t.n
modelcols(t::Term, @nospecialize(d::NamedTuple)) = getproperty(d, t.sym)
modelcols(t::ConstantTerm, @nospecialize(d::NamedTuple)) = t.n

modelcols(ft::FunctionTerm, d::NamedTuple) =
modelcols(ft::FunctionTerm, @nospecialize(d::NamedTuple)) =
Base.Broadcast.materialize(lazy_modelcols(ft, d))

lazy_modelcols(ft::FunctionTerm, d::NamedTuple) =
Base.Broadcast.broadcasted(ft.f, lazy_modelcols.(ft.args, Ref(d))...)
lazy_modelcols(x, d) = modelcols(x, d)
function lazy_modelcols(ft::FunctionTerm, @nospecialize(d::NamedTuple))
args = Vector{Any}(undef, length(ft.args))
for (i, a) in enumerate(ft.args)
args[i] = lazy_modelcols(a, d)
end
return Base.Broadcast.broadcasted(ft.f, args...)
end
lazy_modelcols(x, @nospecialize(d)) = modelcols(x, d)


modelcols(t::ContinuousTerm, d::NamedTuple) = copy.(d[t.sym])
modelcols(t::ContinuousTerm, @nospecialize(d::NamedTuple)) = copy.(d[t.sym])

modelcols(t::CategoricalTerm, d::NamedTuple) = t.contrasts[d[t.sym], :]
modelcols(t::CategoricalTerm, @nospecialize(d::NamedTuple)) = t.contrasts[d[t.sym], :]


"""
Expand Down Expand Up @@ -539,25 +561,24 @@ end

# two options here: either special-case ColumnTable (named tuple of vectors)
# vs. vanilla NamedTuple, or reshape and use normal broadcasting
modelcols(t::InteractionTerm, d::NamedTuple) =
kron_insideout(*, (modelcols(term, d) for term in t.terms)...)
modelcols(t::InteractionTerm, @nospecialize(d::NamedTuple)) =
kron_insideout(*, _modelcols_each(t.terms, d)...)

function modelcols(t::InteractionTerm, d::ColumnTable)
row_kron_insideout(*, (modelcols(term, d) for term in t.terms)...)
end
modelcols(t::InteractionTerm, @nospecialize(d::ColumnTable)) =
row_kron_insideout(*, _modelcols_each(t.terms, d)...)

modelcols(t::InterceptTerm{true}, d::NamedTuple) = ones(size(first(d)))
modelcols(t::InterceptTerm{false}, d) = Matrix{Float64}(undef, size(first(d),1), 0)
modelcols(t::InterceptTerm{true}, @nospecialize(d::NamedTuple)) = ones(size(first(d)))
modelcols(t::InterceptTerm{false}, @nospecialize(d)) = Matrix{Float64}(undef, size(first(d),1), 0)

modelcols(t::FormulaTerm, d::NamedTuple) = (modelcols(t.lhs,d), modelcols(t.rhs, d))
modelcols(t::FormulaTerm, @nospecialize(d::NamedTuple)) = (modelcols(t.lhs,d), modelcols(t.rhs, d))

function modelcols(t::MatrixTerm, d::ColumnTable)
mat = reduce(hcat, [modelcols(tt, d) for tt in t.terms])
function modelcols(t::MatrixTerm, @nospecialize(d::ColumnTable))
mat = reduce(hcat, _modelcols_each(t.terms, d))
reshape(mat, size(mat, 1), :)
end

modelcols(t::MatrixTerm, d::NamedTuple) =
reduce(vcat, [modelcols(tt, d) for tt in t.terms])
modelcols(t::MatrixTerm, @nospecialize(d::NamedTuple)) =
reduce(vcat, _modelcols_each(t.terms, d))

vectorize(x::Tuple) = collect(x)
vectorize(x::AbstractVector) = x
Expand Down
Loading