diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index 95777e05..a7721d55 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -660,41 +660,58 @@ Base.@propagate_inbounds function _getindex( return result end +# With all-Int `front` and a concrete numeric element type, collecting `u[front...]` +# over the selected columns matches the `stack` fallback (which throws on an empty +# selection) without copying `A.u[col_idxs]`. `stack` stacks array-valued elements +# and promotes abstract element types, so those keep the fallback. +function _is_numeric_row(A::AbstractVectorOfArray, front) + T = eltype(eltype(A.u)) + return all(idx -> idx isa Int, front) && isconcretetype(T) && T <: Number +end + 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) + if _is_numeric_row(A, front) + cols = col_idxs isa Colon ? vec(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 _is_numeric_row(A, front) + cols = col_idxs isa Colon ? vec(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..97ee9dee 100644 --- a/test/Core/basic_indexing.jl +++ b/test/Core/basic_indexing.jl @@ -394,3 +394,54 @@ 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] with scalar prefix indices" begin + same(x, y) = typeof(x) == typeof(y) && x == y + + n = 1000 + A = VectorOfArray([rand(5) for _ in 1:n]) + get_row(A, i) = A[i, :] + get_row(A, 2) # compile 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) + @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] + + for X in (A, D), cols in (1:0, Int[], falses(n)) + @test same(X[2, cols], Float64[]) + end + + 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 + + B = VectorOfArray([rand(3, 4) for _ in 1:5]) + @test same(B[1, 2, 1:0], Float64[]) + @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]) + + M = VectorOfArray(reshape([[1.0, 2.0] .+ 10j for j in 1:4], 2, 2)) + @test same(M[2, :], [12.0, 22.0, 32.0, 42.0]) + @test same(M[2, 2:3], [22.0, 32.0]) + M3 = VectorOfArray(reshape([[1.0 2.0; 3.0 4.0] .+ 10j for j in 1:8], 2, 2, 2)) + @test same(M3[1, 2, :], [10j + 2.0 for j in 1:8]) + + # Array-valued and abstract element types are still stacked and promoted. + S = [[SVector(1.0, 2, 3) .+ j, SVector(4.0, 5, 6) .+ j] for j in 1:3] + for X in ( + VectorOfArray(S), DiffEqArray(S, [0.0, 1, 2]), + VectorOfArray([[[1.0, 2], [3.0, 4]] for _ in 1:3]), + VectorOfArray([Real[1, 2.0], Real[3.0, 4]]), + ) + @test same(X[2, :], stack(getindex.(X.u, 2))) + @test same(X[2, 1:2], stack(getindex.(X.u[1:2], 2))) + end +end