rename pool back to compact_mem and implement here (#235)

* rename pool back to compact_mem and implement here

* add tests

Author: piever

src/IndexedTables.jl

@@ -6,7 +6,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!, pool,
+    collect_empty_structarray, grow_to_structarray!, collect_to_structarray!, replace_storage,
     GroupPerm, GroupJoinPerm, roweq, rowcmp, fastpermute!
 
 import Tables, TableTraits, IteratorInterfaceExtensions, TableTraitsUtils

src/columns.jl

@@ -696,3 +696,6 @@ function init_inputs(f::Tup, input, isvec)
     # functions and input
     NT((fs...,)), rows(NT((xs...,)))
 end
+
+# utils
+compact_mem(v::Columns) = replace_storage(compact_mem, v)

src/ndsparse.jl

@@ -398,7 +398,7 @@ function aggregate!(f, x::NDSparse)
     idxs, data = x.index, x.data
     n = length(idxs)
     newlen = 0
-    for ii in GroupPerm(pool(idxs), Base.OneTo(n))
+    for ii in GroupPerm(compact_mem(idxs), Base.OneTo(n))
         newlen += 1
         if newlen != last(ii)
             copyrow!(idxs, newlen, last(ii))

src/sortperm.jl

@@ -35,10 +35,10 @@ function sortpermby(t, by; cache=false, return_keys=false)
 
     if matched_cols == length(by)
         # first n index columns
-        return return_keys ? (partial_perm, pool(rows(t, by))) : partial_perm
+        return return_keys ? (partial_perm, compact_mem(rows(t, by))) : partial_perm
     end
 
-    byrows = pool(rows(t, by))
+    byrows = compact_mem(rows(t, by))
     bycols = columns(byrows)
     perm = if matched_cols > 0
         nxtcol = bycols[matched_cols+1]

src/utils.jl

@@ -76,7 +76,7 @@ astuple(n::NamedTuple) = Tuple(n)
 
 # optimized sortperm: pool non isbits types before calling sortperm_fast or sortperm_by
 
-sortperm_fast(x) = sortperm(pool(x))
+sortperm_fast(x) = sortperm(compact_mem(x))
 
 function append_n!(X, val, n)
     l = length(X)
@@ -266,3 +266,10 @@ getsubfields(t::Tuple, fields) = t[fields]
 
 product(itr) = itr
 product(itrs...) = Base.Generator(reverse, Iterators.product(reverse(itrs)...))
+
+# 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)))

test/test_utils.jl

@@ -33,3 +33,22 @@ let x = rand(3), y = rand(3), v = rand(3), w = rand(3)
     @test vcat((Columns((x,y))), Columns((v,w))) == Columns((vcat(x,v), vcat(y,w)))
     @test vcat(Columns(x=x,y=y), Columns(x=v,y=w)) == Columns(x=vcat(x,v), y=vcat(y,w))
 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)
+    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 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)[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)
+end