add pure kwarg to map

Project.toml

@@ -1,6 +1,6 @@
 name = "PooledArrays"
 uuid = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
-version = "1.2.1"
+version = "1.3.0"
 
 [deps]
 DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"

src/PooledArrays.jl

@@ -304,7 +304,73 @@ Base.findall(pdv::PooledVector{Bool}) = findall(convert(Vector{Bool}, pdv))
 ##
 ##############################################################################
 
-function Base.map(f, x::PooledArray{T,R}) where {T,R<:Integer}
+"""
+    map(f, x::PooledArray; pure::Bool=false)
+
+Transform `PooledArray` `x` by applying `f` to each element.
+
+If `pure=true` then `f` is applied to each element of pool of `x`
+exactly once (even if some elements in pool are not present it `x`).
+This will typically be much faster when the proportion of unique values
+in `x` is small.
+"""
+function Base.map(f, x::PooledArray; pure::Bool=false)
+    pure && return _map_pure(f, x)
+    length(x) == 0 && return PooledArray(collect(Base.Generator(f, x)))
+    v1 = f(x[1])
+    T = typeof(v1)
+    invpool = Dict(v1 => one(eltype(x.refs)))
+    pool = [v1]
+    labels = similar(x.refs)
+    labels[1] = 1
+    nlabels = 1
+    return _map_notpure(f, x, 2, invpool, pool, labels, nlabels)
+end
+
+function _map_notpure(f, xs::PooledArray,
+                      start,
+                      invpool::Dict{T,I},
+                      pool::Vector{T},
+                      labels::Array{I},
+                      nlabels::Int) where {T, I<:Integer}
+    for i in start:length(xs)
+        vi = f(xs[i])
+        Ti = typeof(vi)
+        lbl = get(invpool, x, zero(I))
+        if lbl !== zero(I)
+            labels[i] = lbl
+        else
+            if nlabels == typemax(I) || !(Ti isa T)
+                if nlabels == typemax(I)
+                    I2 = _widen(I)
+                else
+                    I2 = I
+                end
+                if Ti isa T
+                    T2 = T
+                else
+                    T2 = typejoin(T, Ti)
+                end
+                nlabels += 1
+                invpool2 = convert(Dict{T2, I2}, invpool)
+                invpool2[vi] = nlabels
+                pool2 = convert(Vector{T2}, pool)
+                push!(pool2, vi)
+                labels2 = convert(Vector{I2}, labels)
+                labels2[i] = nlabels
+                return _map_notpure(f, xs, invpool2, pool2,
+                                    labels2, i+1, nlabels)
+            end
+            nlabels += 1
+            labels[i] = nlabels
+            invpool[vi] = nlabels
+            push!(pool, vi)
+        end
+    end
+    return PooledArray(RefArray(labels), invpool, pool)
+end
+
+function _map_pure(f, x::PooledArray)
     ks = collect(keys(x.invpool))
     vs = collect(values(x.invpool))
     ks1 = map(f, ks)

test/runtests.jl

@@ -500,3 +500,29 @@ end
     pa2 = repeat(pa1, inner = (2, 1))
     @test pa2 == [1 2; 1 2; 3 4; 3 4]
 end
+
+@testset "map pure tests" begin
+    x = PooledArray([1, 2, 3])
+    x[3] = 1
+    y = map(-, x, pure=true)
+    @test refpool(y) = [-1, -2, -3]
+    @test y == [-1, -2, -1]
+
+    y = map(-, x)
+    @test refpool(y) = [-1, -2]
+    @test y == [-1, -2, -1]
+
+    function f()
+        i = Ref(0)
+        return x -> (i[] -= 1; i[])
+    end
+
+    # the order is strange as we iterate invpool which is a Dict
+    y = map(f(), x, pure=true)
+    @test refpool(y) = [-3, -1, -2]
+    @test y == [-3, -1, -3]
+
+    y = map(f(), x)
+    @test refpool(y) = [-1, -2, -3]
+    @test y == [-1, -2, -3]
+end