Make levels return a CategoricalArray

Having `levels` preserve the eltype of the input is sometimes useful
to write generic code. This is only slightly breaking as the result
still compares equal to the previous behvior returning unwrapped values.

Author: nalimilan

benchmark/benchmarks.jl

@@ -55,11 +55,11 @@ SUITE["many levels"]["CategoricalArray(::Vector{String})"] =
 a = rand([@sprintf("id%010d", k) for k in 1:1000], 10000)
 ca = CategoricalArray(a)
 
-levs = levels(ca)
+levs = unwrap.(levels(ca))
 SUITE["many levels"]["levels! with original levels"] =
     @benchmarkable levels!(ca, levs)
 
-levs = reverse(levels(ca))
+levs = reverse(unwrap.(levels(ca)))
 SUITE["many levels"]["levels! with resorted levels"] =
     @benchmarkable levels!(ca, levs)
 

docs/src/using.md

@@ -20,7 +20,7 @@ By default, the levels are lexically sorted, which is clearly not correct in our
 
 ```jldoctest using
 julia> levels(x)
-3-element Vector{String}:
+3-element CategoricalArray{String,1,UInt32}:
  "Middle"
  "Old"
  "Young"
@@ -68,7 +68,7 @@ To get rid of the `"Old"` group, just call the [`droplevels!`](@ref) function:
 
 ```jldoctest using
 julia> levels(x)
-3-element Vector{String}:
+3-element CategoricalArray{String,1,UInt32}:
  "Young"
  "Middle"
  "Old"
@@ -81,7 +81,7 @@ julia> droplevels!(x)
  "Young"
 
 julia> levels(x)
-2-element Vector{String}:
+2-element CategoricalArray{String,1,UInt32}:
  "Young"
  "Middle"
 
@@ -139,7 +139,7 @@ Levels still need to be reordered manually:
 
 ```jldoctest using
 julia> levels(y)
-3-element Vector{String}:
+3-element CategoricalArray{String,1,UInt32}:
  "Middle"
  "Old"
  "Young"
@@ -263,15 +263,15 @@ true
 Likewise, assigning a `CategoricalValue` from `y` to an entry in `x` expands the levels of `x` with all levels from `y`, *respecting the ordering of levels of both vectors if possible*:
 ```jldoctest using
 julia> levels(x)
-2-element Vector{String}:
+2-element CategoricalArray{String,1,UInt32}:
  "Middle"
  "Old"
 
 julia> x[1] = y[1]
 CategoricalValue{String, UInt32} "Young" (1/2)
 
 julia> levels(x)
-3-element Vector{String}:
+3-element CategoricalArray{String,1,UInt32}:
  "Young"
  "Middle"
  "Old"
@@ -296,7 +296,7 @@ julia> ab = vcat(a, b)
  "c"
 
 julia> levels(ab)
-3-element Vector{String}:
+3-element CategoricalArray{String,1,UInt32}:
  "a"
  "b"
  "c"
@@ -320,7 +320,7 @@ julia> ab2 = vcat(a, b)
  "c"
 
 julia> levels(ab2)
-3-element Vector{String}:
+3-element CategoricalArray{String,1,UInt32}:
  "a"
  "b"
  "c"

ext/CategoricalArraysArrowExt.jl

@@ -7,6 +7,8 @@ import Arrow: ArrowTypes
 const CATARRAY_ARROWNAME = Symbol("JuliaLang.CategoricalArrays.CategoricalArray")
 ArrowTypes.arrowname(::Type{<:CategoricalValue}) = CATARRAY_ARROWNAME
 ArrowTypes.arrowmetadata(::Type{CategoricalValue{T, R}}) where {T, R} = string(R)
+ArrowTypes.ArrowType(::Type{<:CategoricalValue{T}}) where {T} = T
+ArrowTypes.toarrow(x::CategoricalValue) = unwrap(x)
 
 ArrowTypes.arrowname(::Type{Union{<:CategoricalValue, Missing}}) = CATARRAY_ARROWNAME
 ArrowTypes.arrowmetadata(::Type{Union{CategoricalValue{T, R}, Missing}}) where {T, R} =

ext/CategoricalArraysRecipesBaseExt.jl

@@ -9,7 +9,7 @@ else
 end
 
 RecipesBase.@recipe function f(::Type{T}, v::T) where T <: CategoricalValue
-    level_strings = [map(string, levels(v)); missing]
+    level_strings = [map(string, CategoricalArrays._levels(v)); missing]
     ticks --> eachindex(level_strings)
     v -> ismissing(v) ? length(level_strings) : Int(CategoricalArrays.refcode(v)),
     i -> level_strings[Int(i)]

src/array.jl

@@ -239,7 +239,7 @@ function CategoricalArray{T, N, R}(A::CategoricalArray{S, N, Q};
         catch err
             err isa LevelsException || rethrow(err)
             throw(ArgumentError("encountered value(s) not in specified `levels`: " *
-                                "$(setdiff(CategoricalArrays.levels(res), levels))"))
+                                "$(setdiff(_levels(res), levels))"))
         end
     end
     return res
@@ -358,18 +358,18 @@ function _convert(::Type{CategoricalArray{T, N, R}}, A::AbstractArray{S, N};
     copyto!(res, A)
 
     if levels !== nothing
-        CategoricalArrays.levels(res) == levels ||
+        _levels(res) == levels ||
             throw(ArgumentError("encountered value(s) not in specified `levels`: " *
-                                "$(setdiff(CategoricalArrays.levels(res), levels))"))
+                                "$(setdiff(_levels(res), levels))"))
     else
         # if order is defined for level type, automatically apply it
         L = leveltype(res)
         if Base.OrderStyle(L) isa Base.Ordered
-            levels!(res, sort(CategoricalArrays.levels(res)))
+            levels!(res, sort(_levels(res)))
         elseif hasmethod(isless, (L, L))
             # isless may throw an error, e.g. for AbstractArray{T} of unordered T
             try
-                levels!(res, sort(CategoricalArrays.levels(res)))
+                levels!(res, sort(_levels(res)))
             catch e
                  e isa MethodError || rethrow(e)
             end
@@ -382,7 +382,7 @@ end
 # From CategoricalArray (preserve levels, ordering and R)
 function convert(::Type{CategoricalArray{T, N, R}}, A::CategoricalArray{S, N}) where {S, T, N, R}
     if length(A.pool) > typemax(R)
-        throw(LevelsException{T, R}(levels(A)[typemax(R)+1:end]))
+        throw(LevelsException{T, R}(_levels(A)[typemax(R)+1:end]))
     end
 
     if !(T >: Missing) && S >: Missing && any(iszero, A.refs)
@@ -460,7 +460,7 @@ size(A::CategoricalArray) = size(A.refs)
 Base.IndexStyle(::Type{<:CategoricalArray}) = IndexLinear()
 
 function update_refs!(A::CategoricalArray, newlevels::AbstractVector)
-    oldlevels = levels(A)
+    oldlevels = _levels(A)
     levelsmap = similar(A.refs, length(oldlevels)+1)
     # 0 maps to a missing value
     levelsmap[1] = 0
@@ -478,7 +478,7 @@ function merge_pools!(A::CatArrOrSub,
                       updaterefs::Bool=true,
                       updatepool::Bool=true)
     newlevels, ordered = merge_pools(pool(A), pool(B))
-    oldlevels = levels(A)
+    oldlevels = _levels(A)
     pA = A isa SubArray ? parent(A) : A
     ordered!(pA, ordered)
     # If A's levels are an ordered superset of new (merged) pool, no need to recompute refs
@@ -537,8 +537,8 @@ function copyto!(dest::CatArrOrSub{T, N, R}, dstart::Integer,
 
     # try converting src to dest type to avoid partial copy corruption of dest
     # in the event that the src cannot be copied into dest
-    slevs = convert(Vector{T}, levels(src))
-    dlevs = levels(dest)
+    slevs = convert(Vector{T}, _levels(src))
+    dlevs = _levels(dest)
     if eltype(src) >: Missing && !(eltype(dest) >: Missing) && !all(x -> x > 0, srefs)
         throw(MissingException("cannot copy array with missing values to an array with element type $T"))
     end
@@ -591,7 +591,7 @@ function copyto!(dest::CatArrOrSub{T1, N, R}, dstart::Integer,
         return invoke(copyto!, Tuple{AbstractArray, Integer, AbstractArray, Integer, Integer},
                       dest, dstart, src, sstart, n)
     end
-    newdestlevs = destlevs = copy(levels(dest)) # copy since we need original levels below
+    newdestlevs = destlevs = copy(_levels(dest)) # copy since we need original levels below
     srclevsnm = T2 >: Missing ? setdiff(srclevs, [missing]) : srclevs
     if !(srclevsnm ⊆ destlevs)
         # if order is defined for level type, automatically apply it
@@ -701,7 +701,7 @@ While this will reduce memory use, this function is type-unstable, which can aff
 performance inside the function where the call is made. Therefore, use it with caution.
 """
 function compress(A::CategoricalArray{T, N}) where {T, N}
-    R = reftype(length(levels(A.pool)))
+    R = reftype(length(_levels(A.pool)))
     convert(CategoricalArray{T, N, R}, A)
 end
 
@@ -719,11 +719,11 @@ decompress(A::CategoricalArray{T, N}) where {T, N} =
     convert(CategoricalArray{T, N, DefaultRefType}, A)
 
 function vcat(A::CategoricalArray...)
-    ordered = any(isordered, A) && all(a->isordered(a) || isempty(levels(a)), A)
-    newlevels, ordered = mergelevels(ordered, map(levels, A)...)
+    ordered = any(isordered, A) && all(a->isordered(a) || isempty(_levels(a)), A)
+    newlevels, ordered = mergelevels(ordered, map(_levels, A)...)
 
     refsvec = map(A) do a
-        ii = convert(Vector{Int}, indexin(levels(a.pool), newlevels))
+        ii = convert(Vector{Int}, indexin(_levels(a.pool), newlevels))
         [x==0 ? 0 : ii[x] for x in a.refs]::Array{Int,ndims(a)}
     end
 
@@ -761,23 +761,25 @@ This may include levels which do not actually appear in the data
 `missing` will be included only if it appears in the data and
 `skipmissing=false` is passed.
 
-The returned vector is an internal field of `x` which must not be mutated
+The returned vector is owned by `x` and must not be mutated
 as doing so would corrupt it.
 """
-@inline function DataAPI.levels(A::CatArrOrSub{T}; skipmissing::Bool=true) where T
+@inline function DataAPI.levels(A::CatArrOrSub; skipmissing::Bool=true)
     if eltype(A) >: Missing && !skipmissing
         if any(==(0), refs(A))
-            T[levels(pool(A)); missing]
+            eltype(A)[levels(pool(A)); missing]
         else
-            convert(Vector{T}, levels(pool(A)))
+            levels_missing(pool(A))
         end
     else
         levels(pool(A))
     end
 end
 
+_levels(A::CatArrOrSub) = _levels(pool(A))
+
 """
-    levels!(A::CategoricalArray, newlevels::Vector; allowmissing::Bool=false)
+    levels!(A::CategoricalArray, newlevels::AbstractVector; allowmissing::Bool=false)
 
 Set the levels categorical array `A`. The order of appearance of levels will be respected
 by [`levels`](@ref DataAPI.levels), which may affect display of results in some operations; if `A` is
@@ -791,7 +793,7 @@ Else, `newlevels` must include all levels which appear in the data.
 """
 function levels!(A::CategoricalArray{T, N, R}, newlevels::AbstractVector;
                  allowmissing::Bool=false) where {T, N, R}
-    (levels(A) == newlevels) && return A # nothing to do
+    (_levels(A) == newlevels) && return A # nothing to do
 
     # map each new level to its ref code
     newlv2ref = Dict{eltype(newlevels), Int}()
@@ -806,7 +808,7 @@ function levels!(A::CategoricalArray{T, N, R}, newlevels::AbstractVector;
     end
 
     # map each old ref code to new ref code (or 0 if no such level)
-    oldlevels = levels(pool(A))
+    oldlevels = _levels(pool(A))
     oldref2newref = fill(0, length(oldlevels) + 1)
     for (i, lv) in enumerate(oldlevels)
         oldref2newref[i + 1] = get(newlv2ref, lv, 0)
@@ -867,7 +869,7 @@ end
 function _uniquerefs(A::CatArrOrSub{T}) where T
     arefs = refs(A)
     res = similar(arefs, 0)
-    nlevels = length(levels(A))
+    nlevels = length(_levels(A))
     maxunique = nlevels + (T >: Missing ? 1 : 0)
     seen = fill(false, nlevels + 1) # always +1 for 0 (missing ref)
     @inbounds for ref in arefs
@@ -900,7 +902,7 @@ returned by [`levels`](@ref DataAPI.levels)).
 """
 function droplevels!(A::CategoricalArray)
     arefs = refs(A)
-    nlevels = length(levels(A)) + 1 # +1 for missing
+    nlevels = length(_levels(A)) + 1 # +1 for missing
     seen = fill(false, nlevels)
     seen[1] = true # assume that missing is always observed to simplify checks
     nseen = 1
@@ -913,7 +915,7 @@ function droplevels!(A::CategoricalArray)
     end
 
     # replace the pool
-    A.pool = typeof(pool(A))(@inbounds(levels(A)[view(seen, 2:nlevels)]), isordered(A))
+    A.pool = typeof(pool(A))(@inbounds(_levels(A)[view(seen, 2:nlevels)]), isordered(A))
     # recode refs to keep only the seen ones (optimized version of update_refs!())
     seen[1] = false # to start levelsmap from 0
     levelsmap = cumsum(seen)
@@ -1030,7 +1032,7 @@ end
                              ordered=_isordered(A),
                              compress::Bool=false) where {T, N, R}
     # @inline is needed so that return type is inferred when compress is not provided
-    RefType = compress ? reftype(length(CategoricalArrays.levels(A))) : R
+    RefType = compress ? reftype(length(_levels(A))) : R
     CategoricalArray{T, N, RefType}(A, levels=levels, ordered=ordered)
 end
 
@@ -1043,7 +1045,7 @@ function in(x::CategoricalValue, y::CategoricalArray{T, N, R}) where {T, N, R}
     if x.pool === y.pool
         return refcode(x) in y.refs
     else
-        ref = get(y.pool, levels(x.pool)[refcode(x)], zero(R))
+        ref = get(y.pool, _levels(x.pool)[refcode(x)], zero(R))
         return ref != 0 ? ref in y.refs : false
     end
 end

src/pool.jl

@@ -21,8 +21,8 @@ Base.convert(::Type{CategoricalPool{S}}, pool::CategoricalPool{T, R}) where {S,
     convert(CategoricalPool{S, R}, pool)
 
 function Base.convert(::Type{CategoricalPool{T, R}}, pool::CategoricalPool) where {T, R <: Integer}
-    if length(levels(pool)) > typemax(R)
-        throw(LevelsException{T, R}(levels(pool)[typemax(R)+1:end]))
+    if length(pool.levels) > typemax(R)
+        throw(LevelsException{T, R}(pool.levels[typemax(R)+1:end]))
     end
 
     levelsT = convert(Vector{T}, pool.levels)
@@ -37,10 +37,10 @@ Base.copy(pool::CategoricalPool{T, R}) where {T, R} =
 function Base.show(io::IO, pool::CategoricalPool{T, R}) where {T, R}
     @static if VERSION >= v"1.6.0"
         @printf(io, "%s{%s, %s}([%s])", CategoricalPool, T, R,
-                join(map(repr, levels(pool)), ", "))
+                join(map(repr, pool.levels), ", "))
     else
         @printf(io, "%s{%s,%s}([%s])", CategoricalPool, T, R,
-                join(map(repr, levels(pool)), ", "))
+                join(map(repr, pool.levels), ", "))
     end
 
     pool.ordered && print(io, " with ordered levels")
@@ -65,6 +65,7 @@ it doesn't do this itself to avoid doing a dict lookup twice
 
     i = R(n + 1)
     push!(pool.levels, x)
+    push!(pool.levelsinds, i)
     pool_hash = pool.hash
     if pool_hash !== nothing
         pool.hash = hash(x, pool_hash)
@@ -185,10 +186,10 @@ function merge_pools(a::CategoricalPool{T}, b::CategoricalPool) where {T}
         newlevs = T[]
         ordered = isordered(a)
     elseif length(a) == 0
-        newlevs = Vector{T}(levels(b))
+        newlevs = Vector{T}(b.levels)
         ordered = isordered(b)
     elseif length(b) == 0
-        newlevs = copy(levels(a))
+        newlevs = copy(a.levels)
         ordered = isordered(a)
     else
         ordered = isordered(a) && (isordered(b) || b ⊆ a)
@@ -200,7 +201,7 @@ end
 
 @inline function Base.hash(pool::CategoricalPool, h::UInt)
     if pool.hash === nothing
-        pool.hash = hashlevels(levels(pool))
+        pool.hash = hashlevels(pool.levels)
     end
     hash(pool.hash, h)
 end
@@ -246,9 +247,9 @@ end
 
 # Contrary to the CategoricalArray one, this method only allows adding new levels at the end
 # so that existing CategoricalValue objects still point to the same value
-function levels!(pool::CategoricalPool{S, R}, newlevels::Vector;
+function levels!(pool::CategoricalPool{S, R}, newlevels::AbstractVector;
                  checkunique::Bool=true) where {S, R}
-    levs = convert(Vector{S}, newlevels)
+    levs = newlevels isa CategoricalVector{S} ? newlevels : convert(Vector{S}, newlevels)
     if checkunique && !allunique(levs)
         throw(ArgumentError(string("duplicated levels found in levs: ",
                                    join(unique(filter(x->sum(levs.==x)>1, levs)), ", "))))
@@ -259,24 +260,30 @@ function levels!(pool::CategoricalPool{S, R}, newlevels::Vector;
     n = length(levs)
 
     if n > typemax(R)
-        throw(LevelsException{S, R}(setdiff(levs, levels(pool))[typemax(R)-length(levels(pool))+1:end]))
+        throw(LevelsException{S, R}(setdiff(levs, pool.levels)[typemax(R)-length(pool.levels)+1:end]))
     end
 
     empty!(pool.invindex)
     resize!(pool.levels, n)
+    resize!(pool.levelsinds, n)
     pool.hash = nothing
     pool.equalto = C_NULL
     pool.subsetof = C_NULL
     for i in 1:n
         v = levs[i]
         pool.levels[i] = v
+        pool.levelsinds[i] = i
         pool.invindex[v] = i
     end
 
     return pool
 end
 
-DataAPI.levels(pool::CategoricalPool) = pool.levels
+DataAPI.levels(pool::CategoricalPool{T}) where {T} =
+    CategoricalVector{T}(pool.levelsinds, pool)
+levels_missing(pool::CategoricalPool{T}) where {T} =
+    CategoricalVector{Union{T, Missing}}(pool.levelsinds, pool)
+_levels(pool::CategoricalPool) = pool.levels
 
 isordered(pool::CategoricalPool) = pool.ordered
 ordered!(pool::CategoricalPool, ordered) = (pool.ordered = ordered; pool)