Add O(1) fast path to AbstractVectorOfArray getindex

src/vector_of_array.jl

@@ -955,10 +955,35 @@ end
 Base.@propagate_inbounds function Base.getindex(
         A::AbstractVectorOfArray{T, N}, I::Vararg{Int, N}
     ) where {T, N}
-    @boundscheck checkbounds(A, I...)
     col = I[N]
     inner_I = Base.front(I)
-    u_col = A.u[col]
+    u = A.u
+    # Fast path avoiding the O(n) `size(A)` scan below. `size(u_col, d) <=
+    # size(A, d)` always holds, so a 1-based bound against `size(u_col, d)` is
+    # sufficient for an in-bounds index into `A`. Must not use
+    # `checkbounds(u_col, ...)`: that respects a non-1-based inner axis (e.g.
+    # OffsetArray) and would wrongly accept an index `A`'s 1-based contract
+    # must reject. N == 1 stores scalars in `u` directly (no inner array),
+    # where `size(A)` is already O(1); skip the fast path there.
+    # `A`'s column dimension is likewise 1-based over `length(u)` regardless
+    # of `u`'s own axes (e.g. an OffsetVector container): `1 <= col <=
+    # length(u)` matches that contract, and `checkbounds(Bool, u, col)` is
+    # still required since `u`'s own axes need not be `1:length(u)`.
+    if N > 1 && 1 <= col <= length(u) && checkbounds(Bool, u, col)
+        u_col = @inbounds u[col]
+        in_inner_bounds = true
+        for d in 1:(N - 1)
+            if !(1 <= inner_I[d] <= size(u_col, d))
+                in_inner_bounds = false
+                break
+            end
+        end
+        if in_inner_bounds
+            return @inbounds u_col[inner_I...]
+        end
+    end
+    @boundscheck checkbounds(A, I...)
+    u_col = u[col]
     # Return zero for indices outside ragged storage
     for d in 1:(N - 1)
         if inner_I[d] > size(u_col, d)

test/Core/basic_indexing.jl

@@ -1,6 +1,27 @@
 using RecursiveArrayTools, Test
 using RecursiveArrayToolsShorthandConstructors
 
+# Counts calls to `size` on its wrapped vector, used to check that an inner
+# array is not scanned once per column of an `AbstractVectorOfArray`.
+const counter = Ref(0)
+struct CountingVec{T} <: AbstractVector{T}
+    data::Vector{T}
+end
+Base.size(c::CountingVec) = (counter[] += 1; size(c.data))
+Base.getindex(c::CountingVec, i::Int) = c.data[i]
+Base.IndexStyle(::Type{<:CountingVec}) = IndexLinear()
+
+# A vector whose own indices do not start at 1, e.g. `OffsetVec(v, -1)` has
+# axes `0:length(v)-1`, like `OffsetArrays.OffsetArray(v, 0:length(v)-1)`.
+struct OffsetVec{T} <: AbstractVector{T}
+    data::Vector{T}
+    offset::Int
+end
+Base.axes(o::OffsetVec) = (o.offset .+ (1:length(o.data)),)
+Base.size(o::OffsetVec) = size(o.data)
+Base.getindex(o::OffsetVec, i::Int) = o.data[i - o.offset]
+Base.IndexStyle(::Type{<:OffsetVec}) = IndexLinear()
+
 # Example Problem
 recs = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
 testa = cat(recs..., dims = 2)
@@ -394,3 +415,62 @@ x = VectorOfArray(StructArray{SVector{1, Float64}}(ntuple(_ -> [1.0, 2.0], 1)))
 y = 2 * x
 @. x = y
 @test all(all.(y .== x))
+
+@testset "getindex O(1) fast path keeps BoundsError/ragged semantics" begin
+    # count `size` calls on the inner arrays to check the fast path is O(1)
+    # without relying on timing
+    counter[] = 0
+    n = 1000
+    A = VectorOfArray([CountingVec(rand(2)) for _ in 1:n])
+    @test A[2, 7] == A.u[7][2]
+    @test counter[] == 1 # exactly the selected column's own size, not one per inner array
+
+    # ragged: zero-padding inside the max size is unchanged
+    R = VA[[1.0, 2.0], [3.0], [4.0, 5.0, 6.0]]
+    @test R[1, 1] == 1.0
+    @test R[2, 1] == 2.0
+    @test R[3, 2] == 0.0 # inner array 2 (length 1) padded with zero at row 3
+    @test R[1, 3] == 4.0
+
+    # truly out-of-range indices still throw, on both dimensions
+    @test_throws BoundsError R[4, 1] # row 4 exceeds the global max size (3)
+    @test_throws BoundsError R[1, 4] # column 4 exceeds the number of inner arrays
+    @test_throws BoundsError R[0, 1]
+    @test_throws BoundsError R[1, 0]
+end
+
+@testset "getindex fast path rejects a non-1-based inner axis (OffsetArray)" begin
+    # minimal non-1-based AbstractVector, avoiding an OffsetArrays dependency
+    O0 = VA[OffsetVec([1.0, 2.0, 3.0], -1), OffsetVec([4.0, 5.0, 6.0], -1)] # axes 0:2
+    O2 = VA[OffsetVec([1.0, 2.0, 3.0], 1), OffsetVec([4.0, 5.0, 6.0], 1)] # axes 2:4
+    @test_throws BoundsError O0[0, 1]
+    @test_throws BoundsError O2[4, 1]
+end
+
+@testset "getindex fast path rejects a non-1-based VA.u container" begin
+    # VectorOfArray(::AbstractVector{VT}) stores a non-1-based container (like
+    # OffsetVec) as-is; this is the same column-index mistake as the inner-axis
+    # case above, one level up.
+    cols = [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]]
+    C0 = VectorOfArray(OffsetVec(cols, -1)) # axes 0:3
+    C2 = VectorOfArray(OffsetVec(cols, 1)) # axes 2:5
+
+    @test_throws BoundsError C0[1, 0]
+    @test_throws BoundsError C0[2, 0]
+    @test_throws BoundsError C0[CartesianIndex(1, 0)]
+    @test_throws BoundsError C2[1, 5]
+    @test_throws BoundsError C2[2, 5]
+
+    # a column inside both `1:length(u)` and `u`'s own axes still returns the
+    # stored value
+    @test C0[1, 1] == 3.0
+    @test C0[1, 3] == 7.0
+    @test C2[1, 2] == 1.0
+    @test C2[1, 4] == 5.0
+
+    # a column inside `1:length(u)` but outside `u`'s own axes still throws,
+    # matching master (indexing `u` itself fails natively there)
+    @test_throws BoundsError C0[1, 4]
+    @test_throws BoundsError C2[1, 1]
+    @test_throws BoundsError C2[1, 6]
+end