diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index af4da6e9..6b4bb60a 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -1045,6 +1045,13 @@ function Base.append!( end function Base.stack(VA::AbstractVectorOfArray; dims = :) + # `stack.(VA.u)` exists to collapse any nested `AbstractVectorOfArray` + # (or other non-numeric-array) elements of `VA.u` into plain arrays before + # the outer `stack`. When the elements are already plain numeric arrays, + # that inner `stack.()` call is a redundant elementwise copy. + if eltype(VA.u) <: AbstractArray{<:Number} + return stack(VA.u; dims) + end return stack(stack.(VA.u); dims) end diff --git a/test/Core/interface_tests.jl b/test/Core/interface_tests.jl index c4aa6ae8..a5df1452 100644 --- a/test/Core/interface_tests.jl +++ b/test/Core/interface_tests.jl @@ -208,6 +208,14 @@ testva = VA[VA[ones(2, 2), 2ones(2, 2)], 3ones(2, 2, 2)] @test stack(testva) == [1.0 1.0; 1.0 1.0;;; 2.0 2.0; 2.0 2.0;;;; 3.0 3.0; 3.0 3.0;;; 3.0 3.0; 3.0 3.0] +# stack on plain-numeric inner arrays skips `stack.(VA.u)` (an elementwise +# copy of every inner array before the outer stack): only the final output +# array is allocated, instead of also allocating one copy per inner array. +numeric_stack = VectorOfArray([rand(2) for _ in 1:1000]) +stack(numeric_stack) # compile away type instability before measuring +@test (@allocated stack(numeric_stack)) < 20_000 # final 2x1000 Float64 array alone is 16000 bytes +@test stack(numeric_stack) == reduce(hcat, numeric_stack.u) + # convert array from VectorOfArray/DiffEqArray t = 1:8 recs = [rand(10, 7) for i in 1:8]