diff --git a/src/array.jl b/src/array.jl index c1a5c7fe..c1b54cf5 100644 --- a/src/array.jl +++ b/src/array.jl @@ -943,16 +943,17 @@ function Base.sort!(v::CategoricalVector; # compute the order in which to read from counts ord = Base.Sort.ord(lt, by, rev, order) - index = eltype(v) >: Missing ? [missing; v.pool.valindex] : v.pool.valindex seen = counts .> 0 anymissing = eltype(v) >: Missing && seen[1] - perm = sortperm(view(index, seen), order=ord) - nzcounts = counts[seen] + levs = eltype(v) >: Missing ? [missing; v.pool.valindex] : v.pool.valindex + sortedlevs = sort!(Vector(view(levs, seen)), order=ord) + levelsmap = something.(indexin(sortedlevs, levs)) j = 0 refs = v.refs - @inbounds for ref in perm - tmpj = j + nzcounts[ref] - refs[(j+1):tmpj] .= ref - anymissing + @inbounds for i in eachindex(levelsmap) + ref = levelsmap[i] + tmpj = j + counts[ref] + refs[(j+1):tmpj] .= ref - (eltype(v) >: Missing) j = tmpj end diff --git a/test/13_arraycommon.jl b/test/13_arraycommon.jl index a834ab0a..3aa859b7 100644 --- a/test/13_arraycommon.jl +++ b/test/13_arraycommon.jl @@ -926,7 +926,11 @@ end cv = categorical(v) levels!(cv, ["b", "a", "c", "d"]) @test sort(cv, rev=rev) ≅ - ["b", "a", "c", "d", missing][sort([5, 1, 2, 3, 4][cv.refs .+ 1], rev=rev)] + ["b", "a", "c", "d", missing][sort([5; 1:4][cv.refs .+ 1], rev=rev)] + + levels!(cv, ["x", "z", "b", "a", "y", "c", "d", "0"]) + @test sort(cv, rev=rev) ≅ + ["x", "z", "b", "a", "y", "c", "d", "0", missing][sort([9; 1:8][cv.refs .+ 1], rev=rev)] cv = categorical(v) @test sort(cv, rev=rev, lt=(x, y) -> isless(y, x)) ≅