Support changing storage type (#52)

Author: piever

src/utils.jl

@@ -78,3 +78,38 @@ iscompatible(::Type{S}, ::Type{<:AbstractArray{T}}) where {S, T} = S<:T
 function iscompatible(::Type{S}, ::Type{StructArray{T, N, C}}) where {S, T, N, C}
     all_types(iscompatible, staticschema(S), C)
 end
+
+function replace_storage(f, v::AbstractArray{T, N})::AbstractArray{T, N} where {T, N}
+    f(v)
+end
+
+"""
+`replace_storage(f, s::StructArray)`
+
+Change storage type for fieldarrays: each array `v` is replaced by `f(v)`. `f(v) is expected to have the same
+`eltype` and `size` as `v`.
+
+## Examples
+
+If PooledArrays is loaded, we can pool all columns of non `isbitstype`:
+
+```jldoctest
+julia> using PooledArrays
+
+julia> s = StructArray(a=1:3, b = fill("string", 3));
+
+julia> s_pooled = StructArrays.replace_storage(s) do v
+           isbitstype(eltype(v)) ? v : convert(PooledArray, v)
+       end
+10-element StructArray{NamedTuple{(:a, :b),Tuple{Int64,String}},1,NamedTuple{(:a, :b),Tuple{UnitRange{Int64},PooledArray{String,UInt8,1,Array{UInt8,1}}}}}:
+ (a = 1, b = "string") 
+ (a = 2, b = "string") 
+ (a = 3, b = "string") 
+```
+"""
+function replace_storage(f, s::StructArray{T}) where T
+    cols = fieldarrays(s)
+    newcols = map(t -> replace_storage(f, t), cols)
+    StructArray{T}(newcols)
+end
+

test/runtests.jl

@@ -19,6 +19,18 @@ end
     @test StructArrays.propertynames(StructArrays.fieldarrays(t)) == (:a, :b)
 end
 
+@testset "replace_storage" begin
+    v = StructArray(a=rand(10), b = fill("string", 10))
+    v_pooled = StructArrays.replace_storage(v) do c
+        isbitstype(eltype(c)) ? c : convert(PooledArrays.PooledArray, c)
+    end
+    @test eltype(v) == eltype(v_pooled)
+    @test all(v.a .== v_pooled.a)
+    @test all(v.b .== v_pooled.b)
+    @test !isa(v_pooled.a, PooledArrays.PooledArray)
+    @test isa(v_pooled.b, PooledArrays.PooledArray)
+end
+
 @testset "permute" begin
     a = WeakRefStrings.StringVector(["a", "b", "c"])
     b = PooledArrays.PooledArray([1, 2, 3])