Skip to content

Commit 105d574

Browse files
oschulzoschulz
committed
Specialize broadcasting for getindex and findall
1 parent 3570351 commit 105d574

File tree

4 files changed

+203
-0
lines changed

4 files changed

+203
-0
lines changed

‎src/ArraysOfArrays.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ include("util.jl")
1919
include("functions.jl")
2020
include("array_of_similar_arrays.jl")
2121
include("vector_of_arrays.jl")
22+
include("broadcasting.jl")
2223

2324
@static if !isdefined(Base, :get_extension)
2425
include("../ext/ArraysOfArraysAdaptExt.jl")

‎src/broadcasting.jl

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
2+
# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT).
3+
4+
const _RefLike{T} = Union{Tuple{T}, Ref{T}}
5+
6+
7+
_idx_type(::VectorOfSimilarArrays) = Int
8+
_idx_type(A::VectorOfArrays) = eltype(A.elem_ptr)
9+
10+
_similar_idx_vector(A::VectorOfSimilarArrays, ::Type{T}, n::Integer) where T = similar(A.data, T, n)
11+
_similar_idx_vector(A::VectorOfArrays, ::Type{T}, n::Integer) where T = similar(A.elem_ptr, T, n)
12+
13+
function _new_vector_of_arrays_with_lengths(
14+
A::Union{VectorOfSimilarArrays, VectorOfArrays}, ::Type{T},
15+
new_kernel_size::AbstractArray{<:Tuple{Vararg{Integer,M}}},
16+
new_lengths::AbstractVector{<:Integer}
17+
) where {T,M}
18+
new_data = similar(A.data, T, sum(new_lengths))
19+
20+
new_elem_ptr = _similar_idx_vector(A, _idx_type(A), length(new_lengths) + 1)
21+
fill!(new_elem_ptr, 0)
22+
cumsum!(view(new_elem_ptr, firstindex(new_elem_ptr) + 1:lastindex(new_elem_ptr)), new_lengths)
23+
new_elem_ptr .+= firstindex(new_data)
24+
25+
newA = VectorOfArrays(new_data, new_elem_ptr, new_kernel_size, no_consistency_checks)
26+
return newA
27+
end
28+
29+
30+
# A view on the result of view, using an array of indices, would allocate due
31+
# to reindexing, so call SubArray directly:
32+
_noreindex_view(A, idxs...) = SubArray(A, idxs)
33+
_noreindex_view(A::AbstractArray{T,N}, ::Vararg{Colon,N}) where {T,N} = A
34+
35+
_generic_size(A) = size(A)
36+
_generic_size(tpl::Tuple) = (length(tpl),)
37+
38+
Base.Base.@propagate_inbounds function _to_indices(A, idxs)
39+
new_idxs = Base.to_indices(A, idxs)
40+
@boundscheck Base.checkbounds(A, new_idxs...)
41+
return new_idxs
42+
end
43+
44+
# Limited to vectors of vectors for now.
45+
# ToDo: Extend to vectors of arrays.
46+
function Base.Broadcast.broadcasted(
47+
::typeof(getindex),
48+
A::Union{VectorOfSimilarVectors, VectorOfVectors},
49+
Idxs::Union{AbstractVector{<:AbstractVector{<:Integer}},AbstractVector{Colon},_RefLike{<:Union{AbstractVector{<:Integer},Colon}}}...
50+
)
51+
# Checks size compatibility:
52+
bcsz = Base.Broadcast.broadcast_shape(size(A), map(_generic_size, Idxs)...)
53+
54+
new_sizes = _similar_idx_vector(A, NTuple{length(Idxs),_idx_type(A)}, prod(bcsz))
55+
broadcast!(new_sizes, A, Idxs...) do a, idxs...
56+
map(length, _to_indices(a, idxs))
57+
end
58+
59+
new_lengths = prod.(new_sizes)
60+
new_kernel_size = Base.tail.(new_sizes)
61+
62+
T = eltype(A.data)
63+
newA = _new_vector_of_arrays_with_lengths(A, T, new_kernel_size, new_lengths)
64+
newA .= _noreindex_view.(A, Idxs...)
65+
return newA
66+
end
67+
68+
69+
# Limited to vectors of vectors for now.
70+
# ToDo: Extend to vectors of arrays.
71+
function Base.Broadcast.broadcasted(
72+
::typeof(getindex),
73+
A::VectorOfSimilarVectors,
74+
Idxs::Union{VectorOfSimilarVectors{<:Integer},AbstractVector{Colon},_RefLike{<:Union{AbstractVector{<:Integer},Colon}}}...
75+
)
76+
# Checks size compatibility:
77+
Base.Broadcast.broadcast_shape(size(A), map(size, Idxs)...)
78+
79+
sz_inner = map(only ∘ innersize, Idxs)
80+
sz_outer = size(A)
81+
new_data = similar(A.data, (sz_inner..., sz_outer...))
82+
newA = VectorOfSimilarVectors(new_data)
83+
84+
newA .= _noreindex_view.(A, Idxs...)
85+
return newA
86+
end
87+
88+
89+
# Limited to vectors of vectors for now.
90+
# ToDo: Extend to vectors of arrays.
91+
function Base.Broadcast.broadcasted(
92+
::typeof(getindex),
93+
A::VectorOfSimilarVectors,
94+
Idxs::_RefLike{<:Union{AbstractVector{<:Integer},Colon}}...
95+
)
96+
data = A.data
97+
inner_idxs = map(only, Idxs)
98+
outer_idxs = (:,)
99+
new_data = data[inner_idxs..., outer_idxs...]
100+
return VectorOfSimilarVectors(new_data)
101+
end
102+
103+
104+
105+
Base.@propagate_inbounds function _findall!(B::AbstractVector, A::AbstractVector{Bool})
106+
@boundscheck let n::Int = 0
107+
@inbounds for i in eachindex(A)
108+
n += A[i] ? 1 : 0
109+
end
110+
n == length(B) || throw(ArgumentError("_findall! requires output array of correct size"))
111+
end
112+
113+
i_B::Int = firstindex(B)
114+
#@inbounds
115+
for i_A in eachindex(A)
116+
if A[i_A]
117+
B[i_B] = i_A
118+
i_B += 1
119+
end
120+
end
121+
return B
122+
end
123+
124+
function Base.Broadcast.broadcasted(
125+
::typeof(findall),
126+
A::Union{VectorOfSimilarVectors{Bool}, VectorOfVectors{Bool}}
127+
)
128+
new_lengths = _similar_idx_vector(A, _idx_type(A), length(A))
129+
new_lengths .= sum.(A)
130+
131+
new_kernel_size = map(_ -> (), new_lengths)
132+
133+
newA = _new_vector_of_arrays_with_lengths(A, Int, new_kernel_size, new_lengths)
134+
_findall!.(newA, A)
135+
return newA
136+
end

‎test/broadcasting.jl

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT).
2+
3+
using ArraysOfArrays
4+
using Test
5+
6+
@testset "broadcasting" begin
7+
ref_flatview(A::AbstractVector{<:AbstractArray}) = vcat(map(vec, Array(A))...)
8+
9+
ref_VoA1(T::Type, n::Integer) = n == 0 ? [Array{T}(undef, 5)][1:0] : [rand(T, rand(1:9)) for i in 1:n]
10+
ref_VoA2(T::Type, n::Integer) = n == 0 ? [Array{T}(undef, 4, 2)][1:0] : [rand(T, rand(1:4), rand(1:4)) for i in 1:n]
11+
ref_VoA3(T::Type, n::Integer) = n == 0 ? [Array{T}(undef, 3, 2, 4)][1:0] : [rand(T, rand(1:3), rand(1:3), rand(1:3)) for i in 1:n]
12+
13+
ref_AosA1(T::Type, n::Integer) = [rand(T, 7) for i in 1:n]
14+
15+
@testset "getindex broadcast specializations" begin
16+
let A = VectorOfArrays(ref_VoA1(Float32, 100))
17+
refA = Array(A)
18+
19+
for Idxs in [
20+
([rand(eachindex(a), rand(1:length(a))) for a in A],),
21+
(VectorOfVectors([rand(eachindex(a), rand(1:length(a))) for a in A]),),
22+
(tuple(1:1),), (tuple([1, 1, 1]),), (tuple(:),),
23+
(Ref(1:1),), (Ref([1, 1, 1]),), (Ref(:),),
24+
]
25+
@test @inferred(broadcast(getindex, A, Idxs...)) isa VectorOfArrays{eltype(eltype(A))}
26+
@test getindex.(A, Idxs...) == getindex.(refA, Idxs...)
27+
end
28+
end
29+
30+
let A = ArrayOfSimilarArrays(ref_AosA1(Float32, 100))
31+
refA = Array(A)
32+
33+
for Idxs in [
34+
([rand(eachindex(a), rand(1:length(a))) for a in A],),
35+
(VectorOfVectors([rand(eachindex(a), rand(1:length(a))) for a in A]),),
36+
]
37+
@test @inferred(broadcast(getindex, A, Idxs...)) isa VectorOfArrays{eltype(eltype(A))}
38+
@test getindex.(A, Idxs...) == getindex.(refA, Idxs...)
39+
end
40+
41+
for Idxs in [
42+
(VectorOfSimilarVectors([rand(eachindex(a), 5) for a in A]),),
43+
(tuple(3:5),), (tuple([2, 5, 6]),), (tuple(:),),
44+
(Ref(3:5),), (Ref([2, 5, 6]),), (Ref(:),),
45+
]
46+
refA = Array(A)
47+
48+
@test @inferred(broadcast(getindex, A, Idxs...)) isa ArrayOfSimilarArrays{eltype(eltype(A))}
49+
@test getindex.(A, Idxs...) == getindex.(refA, Idxs...)
50+
end
51+
end
52+
end
53+
54+
@testset "findall" begin
55+
for A in [
56+
VectorOfArrays(ref_VoA1(Bool, 100)),
57+
ArrayOfSimilarArrays(ref_AosA1(Bool, 100))
58+
]
59+
refA = Array(A)
60+
61+
@test @inferred(broadcast(findall, A)) isa VectorOfArrays{Int}
62+
@test findall.(A) == findall.(refA)
63+
end
64+
end
65+
end

‎test/runtests.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,5 +7,6 @@ Test.@testset "Package ArraysOfArrays" begin
77
include("functions.jl")
88
include("array_of_similar_arrays.jl")
99
include("vector_of_arrays.jl")
10+
include("broadcasting.jl")
1011
include("test_docs.jl")
1112
end # testset

0 commit comments

Comments
 (0)