support map with GPU arrays

JuliaArrays · Mar 10, 2024 · f8483b5 · f8483b5
1 parent 0b99716
commit f8483b5
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/src/structarray.jl b/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

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1377,6 +1377,15 @@ 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 sa3.b::JLArray == sa.re .+ sa.im
     end
 
     @testset "StructSparseArray" begin