Skip to content

Commit 2ff20c2

Browse files
ChrisRackauckasclaude
ChrisRackauckas
and
claude
committed
Port v3 RaggedEnd for per-column end in ragged arrays
Implements the v3 RaggedEnd/RaggedRange mechanism in the ragged sublibrary so that `r[end, i]` returns the last element of column `i` (not the global max). This matches the original v3 VectorOfArray ragged behavior exactly: r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) r[end, 1] # 2.0 (last of column 1) r[end, 2] # 5.0 (last of column 2) r[1:end, 1] # [1.0, 2.0] r[1:end, 2] # [3.0, 4.0, 5.0] `lastindex(r, d)` returns a `RaggedEnd` sentinel for leading dimensions, which is resolved per-column at indexing time via `_resolve_ragged`. Addresses feedback from @JoshuaLampert. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent f9962b4 commit 2ff20c2

2 files changed

Lines changed: 149 additions & 58 deletions

File tree

‎lib/RecursiveArrayToolsRaggedArrays/src/RecursiveArrayToolsRaggedArrays.jl‎

Lines changed: 117 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -271,19 +271,81 @@ Base.iterate(r::AbstractRaggedVectorOfArray, state) = iterate(r.u, state)
271271
Base.firstindex(r::AbstractRaggedVectorOfArray) = firstindex(r.u)
272272
Base.lastindex(r::AbstractRaggedVectorOfArray) = lastindex(r.u)
273273

274-
# lastindex with dimension — needed for `end` in multi-index expressions like r[end, 2]
275-
# dim N (last) = number of inner arrays; dim 1..N-1 = ragged, use max across inner arrays
276-
function Base.lastindex(r::AbstractRaggedVectorOfArray{T, N}, d::Int) where {T, N}
274+
# ─── RaggedEnd: deferred `end` that resolves per-column ─────────────────────
275+
# `lastindex(r, d)` returns a `RaggedEnd` sentinel instead of an integer.
276+
# When indexing actually happens, `_resolve_ragged` resolves it using the
277+
# actual size of the selected inner array, so `r[end, i]` gives the last
278+
# element of column `i` regardless of other columns' lengths.
279+
280+
struct RaggedEnd
281+
dim::Int
282+
offset::Int
283+
end
284+
RaggedEnd(dim::Int) = RaggedEnd(dim, 0)
285+
286+
Base.:+(re::RaggedEnd, n::Integer) = RaggedEnd(re.dim, re.offset + Int(n))
287+
Base.:-(re::RaggedEnd, n::Integer) = RaggedEnd(re.dim, re.offset - Int(n))
288+
Base.:+(n::Integer, re::RaggedEnd) = re + n
289+
Base.broadcastable(x::RaggedEnd) = Ref(x)
290+
291+
struct RaggedRange
292+
dim::Int
293+
start::Int
294+
step::Int
295+
offset::Int
296+
end
297+
298+
Base.:(:)(stop::RaggedEnd) = RaggedRange(stop.dim, 1, 1, stop.offset)
299+
Base.:(:)(start::Integer, stop::RaggedEnd) = RaggedRange(stop.dim, Int(start), 1, stop.offset)
300+
Base.:(:)(start::Integer, step::Integer, stop::RaggedEnd) = RaggedRange(stop.dim, Int(start), Int(step), stop.offset)
301+
Base.:(:)(start::RaggedEnd, stop::RaggedEnd) = RaggedRange(stop.dim, start.offset, 1, stop.offset)
302+
Base.:(:)(start::RaggedEnd, step::Integer, stop::RaggedEnd) = RaggedRange(stop.dim, start.offset, Int(step), stop.offset)
303+
Base.:(:)(start::RaggedEnd, stop::Integer) = RaggedRange(start.dim, start.offset, 1, Int(stop))
304+
Base.:(:)(start::RaggedEnd, step::Integer, stop::Integer) = RaggedRange(start.dim, start.offset, Int(step), Int(stop))
305+
Base.broadcastable(x::RaggedRange) = Ref(x)
306+
307+
# lastindex returns RaggedEnd for leading dims, plain Int for last (column) dim
308+
@inline function Base.lastindex(r::AbstractRaggedVectorOfArray{T, N}, d::Integer) where {T, N}
277309
if d == N
278-
return length(r.u)
310+
return RaggedEnd(0, Int(lastindex(r.u)))
311+
elseif d < N
312+
isempty(r.u) && return RaggedEnd(0, 0)
313+
return RaggedEnd(Int(d), 0)
279314
else
280-
return isempty(r.u) ? 0 : maximum(size(u, d) for u in r.u)
315+
return RaggedEnd(0, 1)
281316
end
282317
end
283318

284-
# axes with dimension — needed for `end` translation and range indexing
285-
function Base.axes(r::AbstractRaggedVectorOfArray{T, N}, d::Int) where {T, N}
286-
return Base.OneTo(lastindex(r, d))
319+
# Resolve RaggedEnd/RaggedRange to concrete indices for a specific column
320+
@inline _resolve_ragged(idx, ::AbstractRaggedVectorOfArray, ::Any) = idx
321+
@inline function _resolve_ragged(idx::RaggedEnd, r::AbstractRaggedVectorOfArray, col)
322+
if idx.dim == 0
323+
return idx.offset
324+
else
325+
return lastindex(r.u[col], idx.dim) + idx.offset
326+
end
327+
end
328+
@inline function _resolve_ragged(idx::RaggedRange, r::AbstractRaggedVectorOfArray, col)
329+
stop_val = if idx.dim == 0
330+
idx.offset
331+
else
332+
lastindex(r.u[col], idx.dim) + idx.offset
333+
end
334+
return Base.range(idx.start; step = idx.step, stop = stop_val)
335+
end
336+
337+
# Column index resolution
338+
@inline _column_index(::AbstractRaggedVectorOfArray, idx) = idx
339+
@inline _column_index(::AbstractRaggedVectorOfArray, idx::Colon) = nothing # handled by Colon methods
340+
@inline function _column_index(::AbstractRaggedVectorOfArray, idx::RaggedEnd)
341+
return idx.dim == 0 ? idx.offset : idx
342+
end
343+
@inline function _column_index(::AbstractRaggedVectorOfArray, idx::RaggedRange)
344+
if idx.dim == 0
345+
return Base.range(idx.start; step = idx.step, stop = idx.offset)
346+
else
347+
return idx
348+
end
287349
end
288350

289351
Base.keys(r::AbstractRaggedVectorOfArray) = keys(r.u)
@@ -401,6 +463,53 @@ function Base.setindex!(
401463
return v
402464
end
403465

466+
# ─── RaggedEnd/RaggedRange indexing resolution ───────────────────────────────
467+
# Catch-all: any getindex with RaggedEnd or RaggedRange indices gets resolved
468+
# per-column before dispatching to the concrete methods above.
469+
470+
function Base.getindex(r::AbstractRaggedVectorOfArray, I::RaggedEnd, col::Int)
471+
return r.u[col][_resolve_ragged(I, r, col)]
472+
end
473+
function Base.getindex(r::AbstractRaggedVectorOfArray, I::RaggedRange, col::Int)
474+
return r.u[col][_resolve_ragged(I, r, col)]
475+
end
476+
function Base.getindex(r::AbstractRaggedVectorOfArray, I, col::RaggedEnd)
477+
c = _resolve_ragged(col, r, nothing)
478+
return r[I, c]
479+
end
480+
function Base.getindex(r::AbstractRaggedVectorOfArray, I::RaggedEnd, col::RaggedEnd)
481+
c = _resolve_ragged(col, r, nothing)
482+
return r.u[c][_resolve_ragged(I, r, c)]
483+
end
484+
function Base.getindex(r::AbstractRaggedVectorOfArray, ::Colon, col::RaggedEnd)
485+
c = _resolve_ragged(col, r, nothing)
486+
return r.u[c]
487+
end
488+
function Base.getindex(r::AbstractRaggedVectorOfArray, I::RaggedRange, col::RaggedEnd)
489+
c = _resolve_ragged(col, r, nothing)
490+
return r.u[c][_resolve_ragged(I, r, c)]
491+
end
492+
# A[:, start:end] with RaggedRange in column position
493+
function Base.getindex(r::RaggedVectorOfArray, ::Colon, I::RaggedRange)
494+
cols = _resolve_ragged(I, r, nothing)
495+
return RaggedVectorOfArray(r.u[cols])
496+
end
497+
function Base.getindex(r::RaggedDiffEqArray, ::Colon, I::RaggedRange)
498+
cols = _resolve_ragged(I, r, nothing)
499+
return RaggedDiffEqArray(
500+
r.u[cols], r.t[cols], r.p, r.sys; discretes = r.discretes,
501+
interp = r.interp, dense = r.dense
502+
)
503+
end
504+
# Resolve column RaggedRange for _column_index (dim=0 means already resolved)
505+
function _resolve_ragged(idx::RaggedRange, ::AbstractRaggedVectorOfArray, ::Nothing)
506+
return Base.range(idx.start; step = idx.step,
507+
stop = idx.dim == 0 ? idx.offset : error("Cannot resolve inner-dim RaggedRange without column"))
508+
end
509+
function _resolve_ragged(idx::RaggedEnd, ::AbstractRaggedVectorOfArray, ::Nothing)
510+
return idx.dim == 0 ? idx.offset : error("Cannot resolve inner-dim RaggedEnd without column")
511+
end
512+
404513
# A[:, i] returns the i-th inner array directly (no zero-padding)
405514
Base.getindex(r::RaggedVectorOfArray, ::Colon, i::Int) = r.u[i]
406515
Base.setindex!(r::RaggedVectorOfArray, v, ::Colon, i::Int) = (r.u[i] = v)

‎lib/RecursiveArrayToolsRaggedArrays/test/runtests.jl‎

Lines changed: 32 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -404,69 +404,51 @@ using Test
404404
end
405405

406406
@testset "v3 end indexing with ragged arrays (ported)" begin
407+
# `end` in the first dimension is per-column via RaggedEnd (matching v3 behavior)
407408
ragged = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0]])
408-
# A[j, i] accesses the j-th element of i-th inner array
409-
@test ragged[2, 1] == 2.0
410-
@test ragged[3, 2] == 5.0
411-
@test ragged[4, 3] == 9.0
412-
@test ragged[1, 1] == 1.0
413-
@test ragged[2, 2] == 4.0
414-
@test ragged[3, 3] == 8.0
415-
416-
# `end` in the last (column) dimension
417-
@test ragged[:, end] == [6.0, 7.0, 8.0, 9.0]
418-
@test ragged[:, (end - 1):end] isa RaggedVectorOfArray
419-
@test length(ragged[:, (end - 1):end]) == 2
420-
421-
# `end` in first dim = max inner array length (4 here)
422-
# So ragged[end, 1] = ragged[4, 1] which is out of bounds for inner array 1
423-
@test_throws BoundsError ragged[end, 1] # inner[1] has only 2 elements
424-
@test_throws BoundsError ragged[end, 2] # inner[2] has only 3 elements
425-
@test ragged[end, 3] == 9.0 # inner[3] has 4 elements, end=4 works
426-
427-
# Range indexing into inner arrays
428-
@test ragged[1:2, 1] == [1.0, 2.0]
429-
@test ragged[1:3, 2] == [3.0, 4.0, 5.0]
430-
@test ragged[1:4, 3] == [6.0, 7.0, 8.0, 9.0]
431-
# 1:end uses max inner length (4), so only works for longest column
432-
@test_throws BoundsError ragged[1:end, 1] # 1:4 but inner[1] has 2
433-
@test_throws BoundsError ragged[1:end, 2] # 1:4 but inner[2] has 3
434-
@test ragged[1:end, 3] == [6.0, 7.0, 8.0, 9.0] # 1:4 matches inner[3]
409+
@test ragged[end, 1] == 2.0 # end = lastindex(inner[1]) = 2
410+
@test ragged[end, 2] == 5.0 # end = lastindex(inner[2]) = 3
411+
@test ragged[end, 3] == 9.0 # end = lastindex(inner[3]) = 4
412+
@test ragged[end - 1, 1] == 1.0
413+
@test ragged[end - 1, 2] == 4.0
414+
@test ragged[end - 1, 3] == 8.0
415+
416+
# 1:end also resolves per-column
417+
@test ragged[1:end, 1] == [1.0, 2.0]
418+
@test ragged[1:end, 2] == [3.0, 4.0, 5.0]
419+
@test ragged[1:end, 3] == [6.0, 7.0, 8.0, 9.0]
435420

436421
# Colon returns actual arrays
437422
@test ragged[:, 1] == [1.0, 2.0]
438423
@test ragged[:, 2] == [3.0, 4.0, 5.0]
439424
@test ragged[:, 3] == [6.0, 7.0, 8.0, 9.0]
440425

441-
# Subset selection via range
426+
# `end` in last (column) dimension
427+
@test ragged[:, end] == [6.0, 7.0, 8.0, 9.0]
442428
@test ragged[:, 2:end] isa RaggedVectorOfArray
443-
@test ragged[:, 2:end][:, 1] == [3.0, 4.0, 5.0]
429+
@test ragged[:, (end - 1):end] isa RaggedVectorOfArray
430+
@test length(ragged[:, (end - 1):end]) == 2
444431

445432
ragged2 = RaggedVectorOfArray([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0], [7.0, 8.0, 9.0]])
446-
# end in first dim = max length = 4
447-
@test ragged2[end, 1] == 4.0 # inner[1] has 4 elements
448-
@test_throws BoundsError ragged2[end, 2] # inner[2] has only 2
449-
@test_throws BoundsError ragged2[end, 3] # inner[3] has only 3
450-
@test ragged2[end - 1, 1] == 3.0 # end-1 = 3
451-
@test_throws BoundsError ragged2[end - 1, 2] # 3 > length([5,6])
452-
@test ragged2[end - 1, 3] == 9.0 # 3 <= length([7,8,9])
453-
@test ragged2[:, 1] == [1.0, 2.0, 3.0, 4.0]
454-
@test ragged2[:, 2] == [5.0, 6.0]
455-
@test ragged2[:, 3] == [7.0, 8.0, 9.0]
456-
# Explicit range indexing (not using end)
457-
@test ragged2[1:4, 1] == [1.0, 2.0, 3.0, 4.0]
458-
@test ragged2[1:2, 2] == [5.0, 6.0]
459-
@test ragged2[1:3, 3] == [7.0, 8.0, 9.0]
460-
@test ragged2[2:4, 1] == [2.0, 3.0, 4.0]
461-
@test ragged2[2:2, 2] == [6.0]
462-
@test ragged2[2:3, 3] == [8.0, 9.0]
463-
# end in last dim works fine
433+
@test ragged2[end, 1] == 4.0 # end = 4
434+
@test ragged2[end, 2] == 6.0 # end = 2
435+
@test ragged2[end, 3] == 9.0 # end = 3
436+
@test ragged2[end - 1, 1] == 3.0
437+
@test ragged2[end - 1, 2] == 5.0
438+
@test ragged2[end - 1, 3] == 8.0
439+
@test ragged2[end - 2, 1] == 2.0
440+
@test ragged2[1:end, 1] == [1.0, 2.0, 3.0, 4.0]
441+
@test ragged2[1:end, 2] == [5.0, 6.0]
442+
@test ragged2[1:end, 3] == [7.0, 8.0, 9.0]
443+
@test ragged2[2:end, 1] == [2.0, 3.0, 4.0]
444+
@test ragged2[2:end, 2] == [6.0]
445+
@test ragged2[2:end, 3] == [8.0, 9.0]
464446
@test ragged2[:, end] == [7.0, 8.0, 9.0]
465447
@test ragged2[:, 2:end] isa RaggedVectorOfArray
466-
# end in first dim = 4 (max), 1:(end-1) = 1:3
448+
@test ragged2[:, (end - 1):end] isa RaggedVectorOfArray
467449
@test ragged2[1:(end - 1), 1] == [1.0, 2.0, 3.0]
468-
@test_throws BoundsError ragged2[1:(end - 1), 2] # 1:3 but inner[2] has 2
469-
@test ragged2[1:(end - 1), 3] == [7.0, 8.0, 9.0] # 1:3 matches inner[3]
450+
@test ragged2[1:(end - 1), 2] == [5.0]
451+
@test ragged2[1:(end - 1), 3] == [7.0, 8.0]
470452
end
471453

472454
@testset "v3 push! making array ragged (ported)" begin

0 commit comments

Comments
 (0)