Type stability in map call

Hi there, I’m a bit new to Julia, and I can’t seem to figure out how to make my code type-stable. I’m writing some code in the context of feature selection, and the specific function I’m looking at now is meant to create formulas for all predicted variables (ys), given a set of selected right-hand-side variables (rhs). Here is the code:

using StatsModels
function compose(y::Symbol, terms::Vector{Symbol})::FormulaTerm
    return term(y) ~ foldl(+, term.(terms))
function get_formulas(ys::Vector{Symbol}, rhs::Vector{Symbol})::Vector{FormulaTerm}
    return map(y::Symbol -> compose(y, rhs), ys)

Despite the fixed return type of the compose function and the flawless ‘conversion’ of the map result to the fixed return type of the get_formulas function, the compiler does not think the return type of the map call is stable. Here is the output by @code_warntype:

julia> @code_warntype get_formulas([:y1, :y2], [:x1, :x2])

1 ─ %1 = Core.apply_type(Main.Vector, Main.FormulaTerm)::Core.Const(Vector{FormulaTerm})
│   %2 = Main.:(var"#13#14")::Core.Const(var"#13#14")
│   %3 = Core.typeof(rhs)::Core.Const(Vector{Symbol})
│   %4 = Core.apply_type(%2, %3)::Core.Const(var"#13#14"{Vector{Symbol}})
│        (#13 = %new(%4, rhs))
│   %6 = #13::var"#13#14"{Vector{Symbol}}
│   %7 = Main.map(%6, ys)::Vector{_A} where _A
│   %8 = Base.convert(%1, %7)::Any
│   %9 = Core.typeassert(%8, %1)::Vector{FormulaTerm}
└──      return %9

where ::Vector{_A} where _A (at %7) and ::Any (at %8) are colored red.

Could anyone point me in the direction of type stability?

Things I’ve already tried:

  • I tried simply adding ::Vector{FormulaTerm} after the map call, but that results in an error:
TypeError: in typeassert, expected Vector{FormulaTerm}, got a value of type Vector{FormulaTerm{Term, Tuple{Term, Term}}}
  • Changing that type assert to ::Vector{F} where F<: FormulaTerm gets rid of the error, but does not get rid of the type instability.

I may be wrong, but seems like the function is type-stable? From my understanding, if the return type of the function is inferred correctly, it will be good for performance. Though, now that I think about it, it does seem if if all the statements are typed correctly in the body of the function, there might be increased performance.

The type of the map result is colored red, doesn’t that indicate type instability? I don’t understand why the compiler doesn’t know what type the result of a mapping of a function that can only return one type is.

map typically returns a vector whose eltype is as small as possible for the resulting elements. Since TermVector is not concrete, map can select between many possible subtypes.

In your case if you want a type stable result you can do

TermVector[compose(y, rhs) for y in ys]