Filling an array with a generator that returns tuples

I’m trying to create a function that takes a number of iterables and fills an array columnwise with each combination of those. The following function has the correct output, but the number of allocations is unexpected.

using Base.Iterators
using BenchmarkTools

function product_stack(iters...)
    output = zeros(length(iters), (*).(length.(iters)...))
    for (index, combination) in enumerate(product(iters...))
        output[:, index] .= combination
    end
    return output
end
julia> @btime product_stack(1:3, 1:3)
  7.512 μs (91 allocations: 3.69 KiB)
2×9 Matrix{Float64}:
 1.0  2.0  3.0  1.0  2.0  3.0  1.0  2.0  3.0
 1.0  1.0  1.0  2.0  2.0  2.0  3.0  3.0  3.0

Is it possible to do this with a single allocation?

Extra question:
I was hoping something like this might work, but unfortunately it does not:

function product_stack(iters...)
    output = zeros(length(iters), (*).(length.(iters)...))
    output .= product(iters...)
    return output
end

Why not?

This is much faster:

function product_stack2(iters...)
    p = flatten(product(iters...))
    reshape(collect(p), length(iters), :)
end

@btime product_stack2(1:3,1:3)
  101.695 ns (3 allocations: 320 bytes)
2×9 Matrix{Int64}:
 1  2  3  1  2  3  1  2  3
 1  1  1  2  2  2  3  3  3

The second example doesn’t work because when you say output .= x, this replaces all elements of output with the same x element. In your case, x is an iterator and each element of output is a Float64, they can’t match. You can modify it as follows but it won’t be as fast as the version above:

function product_stack(iters...)
    output = zeros(Int64, length(iters), (*).(length.(iters)...))
    output[1:end] .= flatten(product(iters...))
    return output
end
@btime product_stack(1:3,1:3)
  1.560 μs (17 allocations: 1.02 KiB)
1 Like

Wow, thank you for your answer!
That broadened my horizon regarding iterators and especially reshape.