Add comprehension fast path for A[ints..., cols] indexing

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(

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