Suboptimal inference for incompletely specified named tuples

[Keno]

jullia> f(x) = NamedTuple{(:x,)}((x,))
julia> code_typed(f, Tuple{Any})
1-element Array{Any,1}:
 CodeInfo(
1 1 ─ %1 = (Core.tuple)(x)::Tuple{Any}                                      │
  │   %2 = (NamedTuple{(:x,),T} where T<:Tuple)(%1)::Any                    │
  └──      return %2                                                        │
) => Any

julia> code_typed(f, Tuple{Union{Int64, Float64}})
1-element Array{Any,1}:
 CodeInfo(
1 1 ─ %1 = (Core.tuple)(x)::Tuple{Union{Float64, Int64}}                    │
  │   %2 = (NamedTuple{(:x,),T} where T<:Tuple)(%1)::Any                    │
  └──      return %2                                                        │
) => Any

It would be nice if these could at least figure out that it's gonna be a NamedTuple{(:x,), <:Any}, or even if the second one distributed over the union.

[JeffBezanson]

This is due to the constructors being @generated. Currently inference doesn't look at fallback (non-generated) implementations for abstract argument types, but it could.

[andyferris]

Yeah we're coming up against this one, too. It would be nice to infer the use Union{Missing, T} inside NamedTuples. xref JuliaData/TypedTables.jl#34.

I didn't follow your explanation, either. I would have guessed we would arrive at the Expr(:new, ...) line here (which should be inferrable), and the fallback branch wouldn't be invoked. (Or does generation get skipped if the signature contains tuples with Union inside?)

[andyferris]

OK, the reflection tools seem to be indicating that they consider the input types to be "abstract" or non-leaf types.

julia> code_lowered(NamedTuple{(:a, :b),Tuple{Int64,Union{Missing, Bool}}}, Tuple{Tuple{Int64, Union{Missing, Bool}}})
ERROR: Could not expand generator for `@generated` method (::Type{NamedTuple{names,T}})(args::T) where {names, T<:Tuple} in Core at boot.jl:539. This can happen if the provided argument types (Tuple{Tuple{Int64,Union{Missing, Bool}}}) are not leaf types, but the `generated` argument is `true`.
Stacktrace:
 [1] error(::String, ::Method, ::String, ::String, ::Type, ::String, ::String) at ./error.jl:42
 [2] (::getfield(Base, Symbol("##12#13")){Bool,DataType})(::Method) at ./reflection.jl:709
 [3] iterate at ./generator.jl:47 [inlined]
 [4] _collect(::Array{Union{Method, MethodInstance},1}, ::Base.Generator{Array{Union{Method, MethodInstance},1},getfield(Base, Symbol("##12#13")){Bool,DataType}}, ::Base.EltypeUnknown, ::Base.HasShape{1}) at ./array.jl:640
 [5] collect_similar(::Array{Union{Method, MethodInstance},1}, ::Base.Generator{Array{Union{Method, MethodInstance},1},getfield(Base, Symbol("##12#13")){Bool,DataType}}) at ./array.jl:569
 [6] map(::Function, ::Array{Union{Method, MethodInstance},1}) at ./abstractarray.jl:2014
 [7] #code_lowered#11 at ./reflection.jl:704 [inlined]
 [8] code_lowered(::Any, ::Any) at ./reflection.jl:704
 [9] top-level scope at none:0

Would it be fair to say that this heuristic is a little outdated for simple Union types (and tuples thereof), where code can be very fast and inference and code generation can be useful?

[JeffBezanson]

This is not a heuristic. We decided not to ever invoke generators on abstract types, because for that to be sound they have to be monotonic: you have to make sure that your generated code infers to a subtype of what we infer in any case where the arguments are a supertype. Of course that is possible, but in practice it is very difficult to ensure, and you get a segfault if you get it wrong.

[andyferris]

Thanks for the explanation.

So what is the best way forward in this case? As far as I can tell, the main problem with NamedTuple methods will be with the constructors; unlike tuples, named tuples with union elements will actually be concrete types (so all those generated functions you wrote for merge and so-on work regardless). Interstingly... it seems to me that if we had named tuple being a regular struct with a tuple inside and getproperty overloading, the constructor wouldn't need to be generated at all (it would just call new with the input tuple). Is there something the current NamedTuple implementation can do that a struct couldn't do, nowadays?

And the other option was to enable inference on the if !@generated branch?

[tpapp]

I wonder if the following MWE is this issue: define functions

f(::Type{Union{Missing, T}}) where T <: Real = rand() < 0.5 ? zero(T) : missing
_fs(::Type{Tuple{}}) = ()
_fs(::Type{T}) where T <: Tuple = (f(Base.tuple_type_head(T)), _fs(Base.tuple_type_tail(T))...)
fs(::Type{NamedTuple{N,T}}) where {N,T} = NamedTuple{N,T}(_fs(T))

and apply as (output trimmed)

julia> VERSION
v"1.1.0-DEV.671"

julia> T = NamedTuple{(:a, :b), Tuple{Union{Missing, Int}, Union{Missing, Float64}}}
NamedTuple{(:a, :b),Tuple{Union{Missing, Int64},Union{Missing, Float64}}}

julia> @code_warntype fs(T)
Body::Any
   24 ─ %66 = (Core.tuple)(%32, %64)::Tuple{Union{Missing, Int64},Union{Missing, Float64}}
   └───       goto #25                                   ││            
   25 ─ %68 = (NamedTuple{(:a, :b),Tuple{Union{Missing, Int64},Union{Missing, Float64}}})(%66)::Any
   └───       return %68                                 │             

Also, is the current workaround something like

fs_fixed(::Type{NamedTuple{N,T}}) where {N,T} = NamedTuple{N,T}(_fs(T))::NamedTuple{N,T}

[andyferris]

This is the best workaround I can come up with at short notice:

let
    for n in 1:32
        Ts = [Symbol("T$i") for i in 1:n]
        xs = [:(x[$i]) for i in 1:n]
        NT = :(Core.NamedTuple{names, Tuple{$(Ts...)}})
        eval(quote
            $NT(x::Tuple{$(Ts...)}) where {names, $(Ts...)} = $(Expr(:new, NT, xs...))
        end)
    end
end

It seems to work ok, in particular it fixes JuliaData/TypedTables.jl#34. Of course if I need more than 32 elements AND Union types it won't work, and it's not exactly elegant.

Would a workaround PR like this be welcome, or should we think of a more fundamental fix? (Another approach I explored was using the inner constructor as a thunk that makes a simple wrapper around the tuple to make it concrete, ending up calling a generated function that is not an inner constructor but calls $(Expr(:new, ...)) anyway).

EDIT: Hmm... my "workaround" seems to fix inference but generated code is horrible.

[meggart]

Let me just throw in another MWE that seems to be caused by this:

getfirstvals(a,b) = (a=a[1],b=b[1])
a=zeros(Union{Float64,Missing},10)
b=zeros(Union{Float64,Missing},10)

getfirstvals(a,b)

This does not infer as well and any fix or workaround would be welcome.

[Keno]

Fixed by @JeffBezanson in #30577.