use permute!! or copyto! instead of fastpermute! (#236)

Author: piever

src/IndexedTables.jl

@@ -7,7 +7,7 @@ using OnlineStatsBase: OnlineStat, fit!
 using StructArrays: StructVector, StructArray, foreachfield, fieldarrays,
     collect_structarray, staticschema, ArrayInitializer, refine_perm!, collect_structarray,
     collect_empty_structarray, grow_to_structarray!, collect_to_structarray!, replace_storage,
-    GroupPerm, GroupJoinPerm, roweq, rowcmp, fastpermute!
+    GroupPerm, GroupJoinPerm, roweq, rowcmp
 
 import Tables, TableTraits, IteratorInterfaceExtensions, TableTraitsUtils
 

src/columns.jl

@@ -698,4 +698,5 @@ function init_inputs(f::Tup, input, isvec)
 end
 
 # utils
+refs(v::Columns) = Columns(map(refs, fieldarrays(v)))
 compact_mem(v::Columns) = replace_storage(compact_mem, v)

src/indexedtable.jl

@@ -98,7 +98,7 @@ function table(::Val{:serial}, cols::Tup;
             if copy
                 cs = cs[perm]
             else
-                cs = permute!(cs, perm)
+                Base.permute!!(refs(cs), perm)
             end
         elseif copy
             cs = copyto!(similar(cs), cs)
@@ -345,7 +345,7 @@ Sort rows of `t` by `by` in place. All of `Base.sort` keyword arguments can be u
 """
 function sort!(t::IndexedTable, by...; kwargs...)
     isempty(t.pkey) || error("Tables with primary keys can't be sorted in place")
-    permute!(rows(t), sortperm(rows(t, by...); kwargs...))
+    Base.permute!!(refs(rows(t)), sortperm(rows(t, by...); kwargs...))
     t
 end
 

src/join.jl

@@ -574,8 +574,8 @@ function _broadcast(f::Function, B::NDSparse, C::NDSparse; dimmap=nothing)
         idx, iperm, vals = _bcast_loop(f, B, C, B_common_cols, B_perm)
         A = NDSparse(idx, vals, copy=false, presorted=true)
         if !issorted(A.index)
-            fastpermute!(A.index, iperm)
-            fastpermute!(A.data, iperm)
+            copyto!(A.data, A.data[iperm])
+            Base.permute!!(refs(A.index), iperm)
         end
     else
         # TODO

src/utils.jl

@@ -267,9 +267,17 @@ getsubfields(t::Tuple, fields) = t[fields]
 product(itr) = itr
 product(itrs...) = Base.Generator(reverse, Iterators.product(reverse(itrs)...))
 
+# refs
+
+refs(v::PooledArray) = v.refs
+refs(v::AbstractArray) = v
+function refs(v::StringArray{T}) where {T}
+    S = Union{WeakRefString{UInt8}, typeintersect(T, Missing)}
+    convert(StringArray{S}, v)
+end
+
 # pool non isbits types
 
 compact_mem(v::PooledArray) = v
 compact_mem(v::AbstractArray{T}) where {T} = isbitstype(T) ? v : PooledArray(v)
-compact_mem(v::StringArray{String}) =
-    map(String, PooledArray(convert(StringArray{WeakRefString{UInt8}}, v)))
+compact_mem(v::StringArray{T}) where {T} = map(t -> convert(T, t), PooledArray(refs(v)))

test/test_utils.jl

@@ -37,18 +37,40 @@ end
 @testset "compact_mem" begin
     v = Columns(a=rand(10), b = fill("string", 10))
     v_pooled = IndexedTables.replace_storage(v) do c
-        isbitstype(eltype(c)) ? c : convert(PooledArrays.PooledArray, c)
+        isbitstype(eltype(c)) ? c : convert(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)
+    @test !isa(v_pooled.a, PooledArray)
+    @test isa(v_pooled.b, PooledArray)
     @test v_pooled == IndexedTables.compact_mem(v)
     s = WeakRefStrings.StringArray(["a", "b", "c"])
-    @test IndexedTables.compact_mem(s) isa PooledArrays.PooledArray{String}
+    @test IndexedTables.compact_mem(s) isa PooledArray{String}
     @test IndexedTables.compact_mem(s)[1] == "a"
     @test IndexedTables.compact_mem(s)[2] == "b"
     @test IndexedTables.compact_mem(s)[3] == "c"
     @test IndexedTables.compact_mem(IndexedTables.compact_mem(s)) == IndexedTables.compact_mem(s)
+    s = WeakRefStrings.StringArray(["a", missing, "c"])
+    @test IndexedTables.compact_mem(s) isa PooledArray{Union{String, Missing}}
+    @test all(isequal.(s, IndexedTables.compact_mem(s)))
+end
+
+@testset "refs" begin
+    a = WeakRefStrings.StringVector(["a", "b", "c"])
+    b = PooledArrays.PooledArray(["1", "2", "3"])
+    c = [:a, :b, :c]
+    s = Columns(a=a, b=b, c=c)
+    ref = IndexedTables.refs(s)
+    @test ref[1].a isa WeakRefStrings.WeakRefString{UInt8}
+    @test all(isequal.(ref.a, a))
+    @test ref[1].b isa UInt8
+    @test isequal(ref.b, UInt8.([1, 2, 3]))
+    Base.permute!!(ref, sortperm(s))
+    @test issorted(s)
+
+    a = WeakRefStrings.StringVector(["a", missing, "c"])
+    refa = IndexedTables.refs(a)
+    @test refa isa WeakRefStrings.StringArray{Union{Missing, WeakRefStrings.WeakRefString{UInt8}}}
+    @test all(isequal.(a, refa))
 end