Skip to content

Commit

Permalink
Fix missing levels keyword argument to categorical
Browse files Browse the repository at this point in the history
It was documented, but not actually added. Also ensure the levels vector
is copied. Improve tests.
  • Loading branch information
nalimilan committed Jun 21, 2020
1 parent b8530d4 commit ad9609c
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 53 deletions.
18 changes: 11 additions & 7 deletions src/array.jl
Expand Up @@ -731,7 +731,7 @@ function levels!(A::CategoricalArray{T, N, R}, newlevels::Vector;

# replace the pool and recode refs to reflect new pool
if newlevels != oldlevels
newpool = CategoricalPool{nonmissingtype(T), R}(newlevels, isordered(A.pool))
newpool = CategoricalPool{nonmissingtype(T), R}(copy(newlevels), isordered(A.pool))
update_refs!(A, newlevels)
A.pool = newpool
end
Expand Down Expand Up @@ -832,7 +832,7 @@ function Base.reshape(A::CategoricalArray{T, N}, dims::Dims) where {T, N}
end

"""
categorical(A::AbstractArray; compress=false, levels=nothing, ordered=false)
categorical(A::AbstractArray; levels=nothing, ordered=false, compress=false)
Construct a categorical array with the values from `A`.
Expand All @@ -856,16 +856,20 @@ If `A` is already a `CategoricalArray`, its levels, orderedness and reference ty
are preserved unless explicitly overriden.
"""
@inline function categorical(A::AbstractArray{T, N};
compress::Bool=false, ordered=_isordered(A)) where {T, N}
levels::Union{AbstractVector, Nothing}=nothing,
ordered=_isordered(A),
compress::Bool=false) where {T, N}
# @inline is needed so that return type is inferred when compress is not provided
RefType = compress ? reftype(length(unique(A))) : DefaultRefType
CategoricalArray{fixstringtype(T), N, RefType}(A, ordered=ordered)
CategoricalArray{fixstringtype(T), N, RefType}(A, levels=levels, ordered=ordered)
end
@inline function categorical(A::CategoricalArray{T, N, R};
compress::Bool=false, ordered=_isordered(A)) where {T, N, R}
levels::Union{AbstractVector, Nothing}=nothing,
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(levels(A))) : R
CategoricalArray{fixstringtype(T), N, RefType}(A, ordered=ordered)
RefType = compress ? reftype(length(CategoricalArrays.levels(A))) : R
CategoricalArray{fixstringtype(T), N, RefType}(A, levels=levels, ordered=ordered)
end

function in(x::Any, y::CategoricalArray{T, N, R}) where {T, N, R}
Expand Down
6 changes: 4 additions & 2 deletions test/07_levels.jl
Expand Up @@ -72,8 +72,10 @@ using CategoricalArrays: DefaultRefType, levels!
@test_throws ArgumentError levels!(pool, reverse(levels(pool)))

# Adding levels while preserving existing ones
@test levels!(pool, [2, 1, 3, 4, 0, 10, 11, 12, 13, 15, 14]) === pool
@test levels(pool) == [2, 1, 3, 4, 0, 10, 11, 12, 13, 15, 14]
levs = [2, 1, 3, 4, 0, 10, 11, 12, 13, 15, 14]
@test levels!(pool, levs) === pool
@test levels(pool) == levs
@test levels(pool) !== levs

@test isa(pool.levels, Vector{Int})
@test length(pool) === 11
Expand Down
94 changes: 50 additions & 44 deletions test/13_arraycommon.jl
Expand Up @@ -1071,54 +1071,60 @@ end
@test CategoricalArrays.pool(x).levels !== levs
end

v = T["b", "c", "a"]
if levs === nothing || unique(v) levs
for x in (CategoricalArray(v, levels=levs, ordered=ord),
CategoricalArray{T}(v, levels=levs, ordered=ord),
CategoricalArray{T, 1}(v, levels=levs, ordered=ord),
CategoricalArray{T, 1, UInt32}(v, levels=levs, ordered=ord),
CategoricalVector(v, levels=levs, ordered=ord),
CategoricalVector{T}(v, levels=levs, ordered=ord),
CategoricalVector{T, UInt32}(v, levels=levs, ordered=ord),
CategoricalArray(v, levels=levs, ordered=ord))
@test x isa CategoricalVector{T, UInt32}
@test x == v
@test levels(x) == something(levs, sort!(unique(x)))
@test isordered(x) === ord
@test CategoricalArrays.pool(x).levels !== levs
for v in (T["b", "c", "a"], categorical(T["b", "c", "a"]))
if levs === nothing || unique(v) levs
for x in (categorical(v, levels=levs, ordered=ord),
CategoricalArray(v, levels=levs, ordered=ord),
CategoricalArray{T}(v, levels=levs, ordered=ord),
CategoricalArray{T, 1}(v, levels=levs, ordered=ord),
CategoricalArray{T, 1, UInt32}(v, levels=levs, ordered=ord),
CategoricalVector(v, levels=levs, ordered=ord),
CategoricalVector{T}(v, levels=levs, ordered=ord),
CategoricalVector{T, UInt32}(v, levels=levs, ordered=ord),
CategoricalArray(v, levels=levs, ordered=ord))
@test x isa CategoricalVector{T, UInt32}
@test x == v
@test levels(x) == something(levs, sort!(unique(x)))
@test isordered(x) === ord
@test CategoricalArrays.pool(x).levels !== levs
end
else
@test_throws ArgumentError categorical(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T}(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T, 1}(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T, 1, UInt32}(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalVector(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalVector{T}(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalVector{T, UInt32}(v, levels=levs, ordered=ord)
end
else
@test_throws ArgumentError CategoricalArray(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T}(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T, 1}(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T, 1, UInt32}(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalVector(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalVector{T}(v, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalVector{T, UInt32}(v, levels=levs, ordered=ord)
end

m = T["c" "b"; "a" "b"]
if levs === nothing || unique(m) levs
for x in (CategoricalArray{T}(m, levels=levs, ordered=ord),
CategoricalArray{T, 2}(m, levels=levs, ordered=ord),
CategoricalArray{T, 2, UInt32}(m, levels=levs, ordered=ord),
CategoricalMatrix(m, levels=levs, ordered=ord),
CategoricalMatrix{T}(m, levels=levs, ordered=ord),
CategoricalMatrix{T, UInt32}(m, levels=levs, ordered=ord))
@test x isa CategoricalMatrix{T, UInt32}
@test x == m
@test levels(x) == something(levs, sort!(unique(x)))
@test isordered(x) === ord
@test CategoricalArrays.pool(x).levels !== levs
for m in (T["c" "b"; "a" "b"], categorical(T["c" "b"; "a" "b"]))
if levs === nothing || unique(m) levs
for x in (categorical(m, levels=levs, ordered=ord),
CategoricalArray{T}(m, levels=levs, ordered=ord),
CategoricalArray{T, 2}(m, levels=levs, ordered=ord),
CategoricalArray{T, 2, UInt32}(m, levels=levs, ordered=ord),
CategoricalMatrix(m, levels=levs, ordered=ord),
CategoricalMatrix{T}(m, levels=levs, ordered=ord),
CategoricalMatrix{T, UInt32}(m, levels=levs, ordered=ord))
@test x isa CategoricalMatrix{T, UInt32}
@test x == m
@test levels(x) == something(levs, sort!(unique(x)))
@test isordered(x) === ord
@test CategoricalArrays.pool(x).levels !== levs
end
else
@test_throws ArgumentError categorical(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T}(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T, 2}(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T, 2, UInt32}(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalMatrix(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalMatrix{T}(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalMatrix{T, UInt32}(m, levels=levs, ordered=ord)
end
else
@test_throws ArgumentError CategoricalArray(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T}(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T, 2}(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalArray{T, 2, UInt32}(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalMatrix(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalMatrix{T}(m, levels=levs, ordered=ord)
@test_throws ArgumentError CategoricalMatrix{T, UInt32}(m, levels=levs, ordered=ord)
end
end
end
Expand Down

0 comments on commit ad9609c

Please sign in to comment.