I ran into an interesting performance issue when collecting a zip. Before you say “don’t collect!”, the next step is to feed the collected array into sort, which throws a method error if you feed it a zip iterator. The issue involves collecting a zip where one of the arguments to zip has type Vector{Union{Nothing, Int}}. I thought that small unions were supposed to be efficient now, but this appears to be an exception to that rule. Or maybe there’s more going on here than I understand. Here are the input arrays:
using BenchmarkTools
A = rand(['a', 'b', 'c'], 10_000);
x = ['b', 'c', 'a'];
y = indexin(A, x); # eltype == Union{Nothing, Int}
y_pure = convert(Vector{Int}, y); # eltype == Int
As you can see, collecting the zip with y_pure is much faster and only uses 4 allocations, rather than 10,000 allocations. Does anyone understand why collecting the zip with the Union type performs so much worse?
So a variable inferred to be of such type needs to be boxed, and as so, it is allocated in the heap. I think there are cases in which such boxing can be optimized away but does not seem to be this case. I would recommend analyzing the result of @code_warntype or other related macros, to have a better idea of why this happens.
Thanks! That makes sense. I looked at the results of @code_warntype, but they don’t mean a whole lot to me. At any rate, it doesn’t look like there are any type inference problems. Here are the outputs, if anyone is interested:
Yes. I checked the difference between them, and the problem is just that the collected Vector is of a not isbitstype and so each position of it is a pointer to a newly allocated small object in the heap. Memory allocation is not incredibly slow, but comparing the allocation of so many small objects against just writing Ints to a ‘compact’ vector of Ints sums up a lot of difference. I had a similar problem recently, I have an heuristic that runs a loop at least one million times, and it does not a lot of work each iteration. I have reduced its runtime by about 6x by avoiding making any allocations inside the loop.
I think what is meant by ‘small unions are efficient’ (someone correct me if I am wrong) is that they are not much worse than just the allocation needed to box the union, and in some cases (that will not create a new object with fields that are Unions, i.e., that the heap object will not escape the scope) they can be optimized away.
I don’t think this is the case. The memory layout of the “data” portion of the array is identical to that of isbitstype as long as the element type is of small union:
Besides, it doesn’t explain why collect(Tuple{Char, Int64}, zip($A, $y)) is also slow.
I think the problem is the allocation of intermediate tuples due to union splitting not working well with tuple-of-unions. With manual implementations
function collect_int_char1(xs)
ys = Vector{Tuple{Char,Int}}(undef, length(xs))
a, b = xs.is
for i in eachindex(ys, xs.is...)
x1 = @inbounds a[i]
x2 = @inbounds b[i]
if x1 isa Char
if x2 isa Int
@inbounds ys[i] = (x1, x2)
end
end
end
return ys
end
function collect_int_char2(xs)
ys = Vector{Tuple{Char,Int}}(undef, length(xs))
a, b = xs.is
for i in eachindex(ys, xs.is...)
x1 = @inbounds a[i]
x2 = @inbounds b[i]
x = (x1, x2)
if x isa Tuple{Char,Int}
@inbounds ys[i] = x
end
end
return ys
end