Improve performance of constructor and levels! with many levels (JuliaData#100)

levels! was very slow due to repeated calls to in, either directly or via issubset. We do not actually need to call levels!(::CategoricalArray, ::Vector) from the constructor since we know the levels will all be preserved (just reordered): we can use levels!(CategoricalPool, ::Vector) directly.

Also inline get!(::CategoricalPool, x) for maximum performance.

Author: nalimilan

benchmark/benchmarks.jl

@@ -30,3 +30,20 @@ end
     @bench "CategoricalArray{String}" sumequals(ca, ca[1])
     @bench "CategoricalArray{Union{String, Missing}}" sumequals(nca, nca[1])
 end
+
+# With many levels, checking whether some levels have been dropped can be very slow (#93)
+@benchgroup "CategoricalArray{String} with many levels" begin
+    a = rand([@sprintf("id%010d", k) for k in 1:1000], 10000)
+    @bench "CategoricalArray(::Vector{String})" CategoricalArray(a)
+
+    a = rand([@sprintf("id%010d", k) for k in 1:1000], 10000)
+    ca = CategoricalArray(a)
+    levs = levels(ca)
+    @bench "levels! with original levels" levels!(ca, levs)
+
+    levs = reverse(levels(ca))
+    @bench "levels! with resorted levels" levels!(ca, levs)
+
+    levs = [levels(ca); [@sprintf("id2%010d", k) for k in 1:1000]]
+    @bench "levels! with many additional levels" levels!(ca, levs)
+end

src/array.jl

@@ -246,7 +246,7 @@ function convert(::Type{CategoricalArray{T, N, R}}, A::AbstractArray{S, N}) wher
     # if order is defined for level type, automatically apply it
     L = leveltype(res)
     if method_exists(isless, Tuple{L, L})
-        levels!(res, sort(levels(res)))
+        levels!(res.pool, sort(levels(res.pool)))
     end
 
     res
@@ -553,9 +553,10 @@ function levels!(A::CategoricalArray{T}, newlevels::Vector; allow_missing=false)
                                    join(unique(filter(x->sum(newlevels.==x)>1, newlevels)), ", "))))
     end
 
-    # first pass to check whether changes can be applied without error
+    # first pass to check whether, if some levels are removed, changes can be applied without error
     # TODO: save original levels and undo changes in case of error to skip this step
-    if !all(l->l in newlevels, index(A.pool))
+    # equivalent to issubset but faster due to JuliaLang/julia#24624
+    if !isempty(setdiff(index(A.pool), newlevels))
         deleted = [!(l in newlevels) for l in index(A.pool)]
         @inbounds for (i, x) in enumerate(A.refs)
             if T >: Missing

src/pool.jl

@@ -88,7 +88,7 @@ Base.getindex(pool::CategoricalPool, i::Integer) = pool.valindex[i]
 Base.get(pool::CategoricalPool, level::Any) = pool.invindex[level]
 Base.get(pool::CategoricalPool, level::Any, default::Any) = get(pool.invindex, level, default)
 
-function Base.get!(pool::CategoricalPool{T, R}, level::Any) where {T, R}
+@inline function Base.get!(pool::CategoricalPool{T, R}, level::Any) where {T, R}
     get!(pool.invindex, level) do
         x = convert(T, level)
         n = length(pool)
@@ -151,7 +151,8 @@ function levels!(pool::CategoricalPool{S, R}, newlevels::Vector) where {S, R}
     end
 
     # No deletions: can preserve position of existing levels
-    if issubset(pool.index, levs)
+    # equivalent to issubset but faster due to JuliaLang/julia#24624
+    if isempty(setdiff(pool.index, levs))
         append!(pool, setdiff(levs, pool.index))
     else
         empty!(pool.invindex)