From 63b1de63d32688133e550db9b201101128f3a842 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Thu, 8 Oct 2026 07:30:12 -0400 Subject: [PATCH] Add fast path to stack(::AbstractVectorOfArray) for numeric inner arrays MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit stack(VA) called stack(stack.(VA.u); dims): the inner stack.() applies stack to every element of VA.u before the outer stack sees them. That step exists to collapse nested AbstractVectorOfArray (or other non-plain-array) elements into plain arrays first, but when VA.u's elements are already plain numeric arrays, stack.() is a redundant elementwise copy — one extra allocation per inner array for no change in the result. Skip it when eltype(VA.u) <: AbstractArray{<:Number} and stack VA.u directly; fall back to the original stack(stack.(VA.u); dims) path for everything else (nested AbstractVectorOfArray elements, etc.), so that case is unaffected. n=1000 inner arrays: 28.4us/2006 allocs -> 5.6us/3 allocs (1.12), 52.2us/1002 allocs (1.10, unmeasured after since the audit reported "same on 1.10" was not re-verified for this specific candidate; the allocation drop alone (104160 -> 16088 bytes) is confirmed on 1.10). Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Agent-Harness: Claude Code 2.1.294 Agent-Model: claude-sonnet-5 Agent-Session: https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv --- src/vector_of_array.jl | 7 +++++++ test/Core/interface_tests.jl | 8 ++++++++ 2 files changed, 15 insertions(+) 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]