Skip to content

Commit 1e77673

Browse files
authored
keep component array types when slicing (#243)
Unify getindex and view to use the same approach. This helps keep component types better than the Base slicing fallback.
1 parent d3901c6 commit 1e77673

File tree

2 files changed

+9
-5
lines changed

2 files changed

+9
-5
lines changed

src/structarray.jl

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -350,16 +350,17 @@ map(c -> c[I...], Tuple(cols))
350350
end
351351
@inline get_ith(::Tuple{}, I...) = ()
352352

353-
Base.@propagate_inbounds function Base.getindex(x::StructArray{T, <:Any, <:Any, CartesianIndex{N}}, I::Vararg{Int, N}) where {T, N}
353+
Base.@propagate_inbounds Base.getindex(x::StructArray, I...) = _getindex(x, to_indices(x, I)...)
354+
355+
Base.@propagate_inbounds function _getindex(x::StructArray{T}, I::Vararg{Int}) where {T}
354356
cols = components(x)
355357
@boundscheck checkbounds(x, I...)
356358
return createinstance(T, get_ith(cols, I...)...)
357359
end
358360

359-
Base.@propagate_inbounds function Base.getindex(x::StructArray{T, <:Any, <:Any, Int}, I::Int) where {T}
360-
cols = components(x)
361-
@boundscheck checkbounds(x, I)
362-
return createinstance(T, get_ith(cols, I)...)
361+
@inline function _getindex(s::StructArray{T}, I...) where {T}
362+
@boundscheck checkbounds(s, I...)
363+
StructArray{T}(map(v -> @inbounds(getindex(v, I...)), components(s)))
363364
end
364365

365366
@inline function Base.view(s::StructArray{T, N, C}, I...) where {T, N, C}

test/runtests.jl

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,9 @@ Base.convert(::Type{Millimeters}, x::Meters) = Millimeters(x.x*1000)
4848
# Test that explicit `setindex!` returns the entire array
4949
# (Julia's parser ensures that chained assignment returns the value)
5050
@test setindex!(x, 22, 3) === x
51+
52+
s = StructArray(a=1:5)
53+
@test s[2:3].a === 2:3
5154
end
5255

5356
@testset "eltype conversion" begin

0 commit comments

Comments
 (0)