diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index af4da6e9..8a449dba 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -664,37 +664,51 @@ Base.@propagate_inbounds function _getindex( A::AbstractVectorOfArray, ::NotSymbolic, I::Union{Int, AbstractArray{Int}, AbstractArray{Bool}, Colon}... ) - return if last(I) isa Int - A.u[last(I)][Base.front(I)...] - else - stack(getindex.(A.u[last(I)], tuple.(Base.front(I))...)) + if last(I) isa Int + return A.u[last(I)][Base.front(I)...] + end + col_idxs = last(I) + front = Base.front(I) + # Fast path: when every prefix index is a scalar Int, the result is a + # plain Vector of scalars, one per selected column. A direct + # comprehension over A.u avoids both the `A.u[col_idxs]` copy of the + # outer container and the tuple-broadcast machinery of the fallback. + if all(idx -> idx isa Int, front) + cols = col_idxs isa Colon ? A.u : view(A.u, col_idxs) + return [u[front...] for u in cols] end + return stack(getindex.(A.u[col_idxs], tuple.(front)...)) end Base.@propagate_inbounds function _getindex( A::AbstractDiffEqArray, ::NotSymbolic, I::Union{Int, AbstractArray{Int}, AbstractArray{Bool}, Colon}... ) - return if last(I) isa Int - A.u[last(I)][Base.front(I)...] - else - col_idxs = last(I) - # Only preserve DiffEqArray type if all prefix indices are Colons (selecting whole inner arrays) - if all(idx -> idx isa Colon, Base.front(I)) - # For Colon, select all columns - if col_idxs isa Colon - col_idxs = eachindex(A.u) - end - # For DiffEqArray, we need to preserve the time values and type - # Create a vector of sliced arrays instead of stacking into higher-dim array - u_slice = [A.u[col][Base.front(I)...] for col in col_idxs] - # Return as DiffEqArray with sliced time values - return DiffEqArray(u_slice, A.t[col_idxs], parameter_values(A), symbolic_container(A)) - else - # Prefix indices are not all Colons - do the same as VectorOfArray - # (stack the results into a higher-dimensional array) - return stack(getindex.(A.u[col_idxs], tuple.(Base.front(I))...)) + if last(I) isa Int + return A.u[last(I)][Base.front(I)...] + end + col_idxs = last(I) + front = Base.front(I) + # Only preserve DiffEqArray type if all prefix indices are Colons (selecting whole inner arrays) + if all(idx -> idx isa Colon, front) + # For Colon, select all columns + if col_idxs isa Colon + col_idxs = eachindex(A.u) end + # For DiffEqArray, we need to preserve the time values and type + # Create a vector of sliced arrays instead of stacking into higher-dim array + u_slice = [A.u[col][front...] for col in col_idxs] + # Return as DiffEqArray with sliced time values + return DiffEqArray(u_slice, A.t[col_idxs], parameter_values(A), symbolic_container(A)) + elseif all(idx -> idx isa Int, front) + # Same fast path as the AbstractVectorOfArray fallback above: a + # scalar prefix selects a plain Vector of scalars, one per column. + cols = col_idxs isa Colon ? A.u : view(A.u, col_idxs) + return [u[front...] for u in cols] + else + # Prefix indices are not all Colons or all Ints - do the same as VectorOfArray + # (stack the results into a higher-dimensional array) + return stack(getindex.(A.u[col_idxs], tuple.(front)...)) end end Base.@propagate_inbounds function _getindex( diff --git a/test/Core/basic_indexing.jl b/test/Core/basic_indexing.jl index f3b05cbd..74cfc18b 100644 --- a/test/Core/basic_indexing.jl +++ b/test/Core/basic_indexing.jl @@ -394,3 +394,39 @@ x = VectorOfArray(StructArray{SVector{1, Float64}}(ntuple(_ -> [1.0, 2.0], 1))) y = 2 * x @. x = y @test all(all.(y .== x)) + +@testset "A[i, cols]/A[i, :] scalar-row indexing is O(1) per column, ragged unaffected" begin + # A[i, :] and A[i, range] (a scalar row index, selecting many columns) do + # not grow the outer container: the result is read straight from A.u + # without first copying A.u[col_idxs] or going through tuple-broadcast + # machinery, so allocation does not scale with `n` beyond the output. + n = 1000 + A = VectorOfArray([rand(5) for _ in 1:n]) + get_row(A, i) = A[i, :] + get_row(A, 2) # compile away type instability before measuring + @test (@allocated get_row(A, 2)) < 10_000 # n Float64s alone is 8000 bytes + @test A[2, :] == [A.u[j][2] for j in 1:n] + @test A[1, 1:10] == [A.u[j][1] for j in 1:10] + + t = collect(range(0, 1, length = n)) + D = DiffEqArray([rand(5) for _ in 1:n], t) + get_row(D, 2) # compile away type instability before measuring + @test (@allocated get_row(D, 2)) < 10_000 + @test D[2, :] == [D.u[j][2] for j in 1:n] + @test D[1, 1:10] == [D.u[j][1] for j in 1:10] + + # Ragged behavior is unaffected by the fast path: `R[i, :]`/`R[i, range]` + # still throws (only `R[i, j]` zero-pads), and in-bounds scalar-row reads + # across ragged columns are unaffected. + R = VectorOfArray([[1.0, 2.0], [3.0], [4.0, 5.0, 6.0]]) + @test_throws BoundsError R[2, :] + @test R[1, 1:3] == [1.0, 3.0, 4.0] + @test R[1, 1] == 1.0 + @test R[3, 2] == 0.0 # zero-padded, untouched by this fast path + + # Mixed Int/Colon prefixes (not all-Int, so outside the fast path) still + # take the stack-based fallback and are unaffected. + B = VectorOfArray([rand(3, 4) for _ in 1:5]) + @test B[1, :, 2:3] == stack([B.u[c][1, :] for c in 2:3]) + @test B[:, 2, 2:3] == stack([B.u[c][:, 2] for c in 2:3]) +end