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

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(

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