diff --git a/src/errormessages.jl b/src/errormessages.jl index 7edbeb46..1214125b 100644 --- a/src/errormessages.jl +++ b/src/errormessages.jl @@ -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) @@ -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 diff --git a/src/modelframe.jl b/src/modelframe.jl index c9fed27f..b9770736 100644 --- a/src/modelframe.jl +++ b/src/modelframe.jl @@ -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 ) @@ -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) @@ -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))")) @@ -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) || diff --git a/src/schema.jl b/src/schema.jl index d05b0ad7..a7b3d203 100644 --- a/src/schema.jl +++ b/src/schema.jl @@ -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) @@ -170,7 +177,7 @@ 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)) @@ -178,7 +185,7 @@ function concrete_term(t::Term, dt::ColumnTable, hint) 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)) @@ -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 diff --git a/src/temporal_terms.jl b/src/temporal_terms.jl index 3a8f25a3..d8b6f2e7 100644 --- a/src/temporal_terms.jl +++ b/src/temporal_terms.jl @@ -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 diff --git a/src/terms.jl b/src/terms.jl index 3617e28b..7e1ebc1e 100644 --- a/src/terms.jl +++ b/src/terms.jl @@ -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], :] """ @@ -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