Fix performance problem pooling strings (JuliaArrays#63)

* optimize pooling of stringarray

* minor change

* test pooling pooled arrays

Author: piever

src/StructArrays.jl

@@ -18,10 +18,16 @@ function __init__()
     Requires.@require Tables="bd369af6-aec1-5ad0-b16a-f7cc5008161c" include("tables.jl")
     Requires.@require WeakRefStrings="ea10d353-3f73-51f8-a26c-33c1cb351aa5" begin
         fastpermute!(v::WeakRefStrings.StringArray, p::AbstractVector) = permute!(v, p)
+        function to_weakrefs(a::WeakRefStrings.StringArray{String})
+            convert(WeakRefStrings.StringArray{WeakRefStrings.WeakRefString{UInt8}}, a)
+        end
         @inline function roweq(a::WeakRefStrings.StringArray{String}, i, j)
-            weaksa = convert(WeakRefStrings.StringArray{WeakRefStrings.WeakRefString{UInt8}}, a)
+            weaksa = to_weakrefs(a)
             @inbounds isequal(weaksa[i], weaksa[j])
         end
+        function pool(v::WeakRefStrings.StringArray{String}, condition = !isbitstype∘eltype)
+            condition(v) ? map(String, PooledArray(to_weakrefs(v))) : v
+        end
     end
 end
 

src/sort.jl

@@ -9,7 +9,8 @@ function Base.permute!(c::StructArray, p::AbstractVector)
     return c
 end
 
-pool(v::AbstractArray, condition = !isbitstype∘eltype) = condition(v) ? convert(PooledArray, v) : v
+pool(v::PooledArray, condition = !isbitstype∘eltype) = v
+pool(v::AbstractArray, condition = !isbitstype∘eltype) = condition(v) ? PooledArray(v) : v
 pool(v::StructArray, condition = !isbitstype∘eltype) = replace_storage(t -> pool(t, condition), v)
 
 struct GroupPerm{V<:AbstractVector, P<:AbstractVector{<:Integer}, U<:AbstractUnitRange}

test/runtests.jl

@@ -28,6 +28,12 @@ end
     @test !isa(v_pooled.a, PooledArrays.PooledArray)
     @test isa(v_pooled.b, PooledArrays.PooledArray)
     @test v_pooled == StructArrays.pool(v)
+    s = WeakRefStrings.StringArray(["a", "b", "c"])
+    @test StructArrays.pool(s) isa PooledArrays.PooledArray{String}
+    @test StructArrays.pool(s)[1] == "a"
+    @test StructArrays.pool(s)[2] == "b"
+    @test StructArrays.pool(s)[3] == "c"
+    @test StructArrays.pool(StructArrays.pool(s)) == StructArrays.pool(s)
 end
 
 @testset "roweq" begin