make lastindex type stable by always returning RaggedEnd

Author: JoshuaLampert

src/vector_of_array.jl

@@ -402,14 +402,16 @@ function Base.lastindex(VA::AbstractVectorOfArray{T,N,A}) where {T,N,A}
     return lastindex(VA.u)
 end
 
+# Always return RaggedEnd for type stability. Use dim=0 to indicate a plain index stored in offset.
+# _resolve_ragged_index and _column_indices handle the dim=0 case to extract the actual index value.
 @inline function Base.lastindex(VA::AbstractVectorOfArray, d::Integer)
     if d == ndims(VA)
-        return lastindex(VA.u)
+        return RaggedEnd(0, Int(lastindex(VA.u)))
     elseif d < ndims(VA)
-        isempty(VA.u) && return 0
-        return RaggedEnd(Int(d))
+        isempty(VA.u) && return RaggedEnd(0, 0)
+        return RaggedEnd(Int(d), 0)
     else
-        return 1
+        return RaggedEnd(0, 1)
     end
 end
 
@@ -534,13 +536,27 @@ end
 @inline function _column_indices(VA::AbstractVectorOfArray, idx::AbstractArray{Bool})
     findall(idx)
 end
+@inline function _column_indices(VA::AbstractVectorOfArray, idx::RaggedEnd)
+    # RaggedEnd with dim=0 means it's just a plain index stored in offset
+    idx.dim == 0 ? idx.offset : idx
+end
 
 @inline _resolve_ragged_index(idx, ::AbstractVectorOfArray, ::Any) = idx
 @inline function _resolve_ragged_index(idx::RaggedEnd, VA::AbstractVectorOfArray, col)
-    return lastindex(VA.u[col], idx.dim) + idx.offset
+    if idx.dim == 0
+        # Special case: dim=0 means the offset contains the actual index value
+        return idx.offset
+    else
+        return lastindex(VA.u[col], idx.dim) + idx.offset
+    end
 end
 @inline function _resolve_ragged_index(idx::RaggedRange, VA::AbstractVectorOfArray, col)
-    stop_val = lastindex(VA.u[col], idx.dim) + idx.offset
+    stop_val = if idx.dim == 0
+        # dim == 0 is the sentinel for an already-resolved plain index stored in offset
+        idx.offset
+    else
+        lastindex(VA.u[col], idx.dim) + idx.offset
+    end
     return Base.range(idx.start; step = idx.step, stop = stop_val)
 end
 @inline function _resolve_ragged_index(idx::AbstractRange{<:RaggedEnd}, VA::AbstractVectorOfArray, col)
@@ -583,16 +599,100 @@ end
 @inline _has_ragged_end(x, xs...) = _has_ragged_end(x) || _has_ragged_end(xs)
 
 @inline function _ragged_getindex(A::AbstractVectorOfArray, I...)
-    cols = last(I)
-    prefix = Base.front(I)
+    n = ndims(A)
+    # Special-case when user provided one fewer index than ndims(A): last index is column selector.
+    if length(I) == n - 1
+        raw_cols = last(I)
+        # If the raw selector is a RaggedEnd/RaggedRange referring to inner dims, reinterpret as column selector.
+        cols = if raw_cols isa RaggedEnd && raw_cols.dim != 0
+            lastindex(A.u) + raw_cols.offset
+        elseif raw_cols isa RaggedRange && raw_cols.dim != 0
+            stop_val = lastindex(A.u) + raw_cols.offset
+            Base.range(raw_cols.start; step = raw_cols.step, stop = stop_val)
+        else
+            _column_indices(A, raw_cols)
+        end
+        prefix = Base.front(I)
+        if cols isa Int
+            resolved_prefix = _resolve_ragged_indices(prefix, A, cols)
+            inner_nd = ndims(A.u[cols])
+            n_missing = inner_nd - length(resolved_prefix)
+            padded = if n_missing > 0
+                if all(idx -> idx === Colon(), resolved_prefix)
+                    (resolved_prefix..., ntuple(_ -> Colon(), n_missing)...)
+                else
+                    (resolved_prefix..., (lastindex(A.u[cols], length(resolved_prefix) + i) for i in 1:n_missing)...,)
+                end
+            else
+                resolved_prefix
+            end
+            return A.u[cols][padded...]
+        else
+            return VectorOfArray([
+                begin
+                    resolved_prefix = _resolve_ragged_indices(prefix, A, col)
+                    inner_nd = ndims(A.u[col])
+                    n_missing = inner_nd - length(resolved_prefix)
+                    padded = if n_missing > 0
+                        if all(idx -> idx === Colon(), resolved_prefix)
+                            (resolved_prefix..., ntuple(_ -> Colon(), n_missing)...)
+                        else
+                            (resolved_prefix..., (lastindex(A.u[col], length(resolved_prefix) + i) for i in 1:n_missing)...,)
+                        end
+                    else
+                        resolved_prefix
+                    end
+                    A.u[col][padded...]
+                end for col in cols
+            ])
+        end
+    end
+
+    # Otherwise, use the full-length interpretation (last index is column selector; missing columns default to Colon()).
+    if length(I) == n
+        cols = last(I)
+        prefix = Base.front(I)
+    else
+        cols = Colon()
+        prefix = I
+    end
     if cols isa Int
+        if all(idx -> idx === Colon(), prefix)
+            return A.u[cols]
+        end
         resolved = _resolve_ragged_indices(prefix, A, cols)
-        return A.u[cols][resolved...]
+        inner_nd = ndims(A.u[cols])
+        padded = (resolved..., ntuple(_ -> Colon(), max(inner_nd - length(resolved), 0))...)
+        return A.u[cols][padded...]
     else
         col_idxs = _column_indices(A, cols)
+        # Resolve sentinel RaggedEnd/RaggedRange (dim==0) for column selection
+        if col_idxs isa RaggedEnd
+            col_idxs = _resolve_ragged_index(col_idxs, A, 1)
+        elseif col_idxs isa RaggedRange
+            col_idxs = _resolve_ragged_index(col_idxs, A, 1)
+        end
+        # If we're selecting whole inner arrays (all leading indices are Colons),
+        # keep the result as a VectorOfArray to match non-ragged behavior.
+        if all(idx -> idx === Colon(), prefix)
+            if col_idxs isa Int
+                return A.u[col_idxs]
+            else
+                return VectorOfArray(A.u[col_idxs])
+            end
+        end
+        # If col_idxs resolved to a single Int, handle it directly
+        if col_idxs isa Int
+            resolved = _resolve_ragged_indices(prefix, A, col_idxs)
+            inner_nd = ndims(A.u[col_idxs])
+            padded = (resolved..., ntuple(_ -> Colon(), max(inner_nd - length(resolved), 0))...)
+            return A.u[col_idxs][padded...]
+        end
         vals = map(col_idxs) do col
             resolved = _resolve_ragged_indices(prefix, A, col)
-            A.u[col][resolved...]
+            inner_nd = ndims(A.u[col])
+            padded = (resolved..., ntuple(_ -> Colon(), max(inner_nd - length(resolved), 0))...)
+            A.u[col][padded...]
         end
         return stack(vals)
     end

test/basic_indexing.jl

@@ -173,6 +173,8 @@ ragged = VectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0]])
 @test ragged[1:end, 1] == [1.0, 2.0]
 @test ragged[1:end, 2] == [3.0, 4.0, 5.0]
 @test ragged[1:end, 3] == [6.0, 7.0, 8.0, 9.0]
+@test ragged[:, end] == [6.0, 7.0, 8.0, 9.0]
+@test ragged[:, 2:end] == VectorOfArray(ragged.u[2:end])
 
 ragged2 = VectorOfArray([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0], [7.0, 8.0, 9.0]])
 @test ragged2[end, 1] == 4.0
@@ -188,6 +190,8 @@ ragged2 = VectorOfArray([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0], [7.0, 8.0, 9.0]])
 @test ragged2[2:end, 1] == [2.0, 3.0, 4.0]
 @test ragged2[2:end, 2] == [6.0]
 @test ragged2[2:end, 3] == [8.0, 9.0]
+@test ragged2[:, end] == [7.0, 8.0, 9.0]
+@test ragged2[:, 2:end] == VectorOfArray(ragged2.u[2:end])
 @test ragged2[1:(end - 1), 1] == [1.0, 2.0, 3.0]
 @test ragged2[1:(end - 1), 2] == [5.0]
 @test ragged2[1:(end - 1), 3] == [7.0, 8.0]
@@ -209,6 +213,10 @@ u[1, :, 2] .= [1.0, 2.0, 3.0]
 # partial column selection by indices
 u[1, [1, 3], 2] .= [7.0, 9.0]
 @test u.u[2] == [7.0 2.0 9.0]
+# test scalar indexing with end
+@test u[1, 1, end] == u.u[end][1, 1]
+@test u[1, end, end] == u.u[end][1, end]
+@test u[1, 2:end, end] == vec(u.u[end][1, 2:end])
 
 # 3D inner arrays (tensors) with ragged third dimension
 u = VectorOfArray([zeros(2, 1, n) for n in (2, 3)])
@@ -223,6 +231,10 @@ u[1:2, 1, [1, 3], 2] .= [1.0 3.0; 2.0 4.0]
 @test u.u[2][2, 1, 1] == 2.0
 @test u.u[2][1, 1, 3] == 3.0
 @test u.u[2][2, 1, 3] == 4.0
+@test u[:, :, end] == u.u[end]
+@test u[:, :, 2:end] == VectorOfArray(u.u[2:end])
+@test u[1, 1, end] == u.u[end][1, 1, end]
+@test u[end, 1, end] == u.u[end][end, 1, end]
 
 # Test that views can be modified
 f3 = VectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]])