Skip to content

Commit

Permalink
keep component array types when slicing (#243)
Browse files Browse the repository at this point in the history
Unify getindex and view to use the same approach. This helps keep component types better than the Base slicing fallback.
  • Loading branch information
aplavin authored Sep 16, 2022
1 parent d3901c6 commit 1e77673
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 5 deletions.
11 changes: 6 additions & 5 deletions src/structarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -350,16 +350,17 @@ map(c -> c[I...], Tuple(cols))
end
@inline get_ith(::Tuple{}, I...) = ()

Base.@propagate_inbounds function Base.getindex(x::StructArray{T, <:Any, <:Any, CartesianIndex{N}}, I::Vararg{Int, N}) where {T, N}
Base.@propagate_inbounds Base.getindex(x::StructArray, I...) = _getindex(x, to_indices(x, I)...)

Base.@propagate_inbounds function _getindex(x::StructArray{T}, I::Vararg{Int}) where {T}
cols = components(x)
@boundscheck checkbounds(x, I...)
return createinstance(T, get_ith(cols, I...)...)
end

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

@inline function Base.view(s::StructArray{T, N, C}, I...) where {T, N, C}
Expand Down
3 changes: 3 additions & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,9 @@ Base.convert(::Type{Millimeters}, x::Meters) = Millimeters(x.x*1000)
# Test that explicit `setindex!` returns the entire array
# (Julia's parser ensures that chained assignment returns the value)
@test setindex!(x, 22, 3) === x

s = StructArray(a=1:5)
@test s[2:3].a === 2:3
end

@testset "eltype conversion" begin
Expand Down

0 comments on commit 1e77673

Please sign in to comment.