support map with GPU arrays

Author: aplavin

src/structarray.jl

@@ -495,6 +495,8 @@ function Base.showarg(io::IO, s::StructArray{T}, toplevel) where T
     toplevel && print(io, " with eltype ", T)
 end
 
+Base.map(f, s::StructArray) = f.(s)
+
 # broadcast
 import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, Broadcasted, DefaultArrayStyle, Unknown, ArrayConflict
 using Base.Broadcast: combine_styles

test/runtests.jl

@@ -1377,6 +1377,16 @@ Base.BroadcastStyle(::Broadcast.ArrayStyle{MyArray2}, S::Broadcast.DefaultArrayS
         @test @inferred(bcmul2(sa)) isa StructArray
         @test backend(bcmul2(sa)) === backend(sa)
         @test (sa .+= 1) === sa
+
+        @test_broken collect(sa)
+
+        a2 = map(x -> real(x) + 1, sa)
+        @test a2::JLArray == sa.re .+ 1
+        sa2 = map(x -> x + 1, sa)
+        @test sa2.re::JLArray == sa.re .+ 1
+        sa3 = map(x -> (a=x + 1, b=x.re + x.im), sa)
+        @test JLArrays.GPUArrays.collect_to_cpu(sa3.a::JLArray) == sa .+ 1
+        @test sa3.b::JLArray == sa.re .+ sa.im
     end
 
     @testset "StructSparseArray" begin