add in-place Base.unique!(::CatVec) -> CatVec

JuliaData · Apr 8, 2024 · b9fd8fc · b9fd8fc
1 parent 5f14ef9
commit b9fd8fc
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 1 deletion.
diff --git a/src/array.jl b/src/array.jl
@@ -1,7 +1,7 @@
 ## Code for CategoricalArray
 
 import Base: Array, convert, collect, copy, getindex, setindex!, similar, size,
-             unique, vcat, in, summary, float, complex, copyto!
+             unique, unique!, vcat, in, summary, float, complex, copyto!
 
 # Used for keyword argument default value
 _isordered(x::AbstractCategoricalArray) = isordered(x)
@@ -890,6 +890,15 @@ end
 unique(A::CatArrOrSub{T}) where T =
     CategoricalVector{T}(_uniquerefs(A), copy(pool(A)))
 
+function unique!(A::CategoricalVector)
+    _urefs = _uniquerefs(A)
+    if length(_urefs) != length(A)
+        resize!(A.refs, length(_urefs))
+        copyto!(A.refs, _urefs)
+    end
+    return A
+end
+
 """
     droplevels!(A::CategoricalArray)
 

diff --git a/test/11_array.jl b/test/11_array.jl
@@ -741,13 +741,21 @@ end
     @test levels(x) == ["Unused1", "Young", "Middle", "Old", "Unused2"]
     @test unique(x) == ["Old", "Young", "Middle"]
 
+    y = copy(x)
+    @test unique!(y) === y
+    @test y == unique(x)
+
     x = CategoricalArray(String[])
     @test isa(levels(x), Vector{String}) && isempty(levels(x))
     @test isa(unique(x), typeof(x)) && isempty(unique(x))
     @test levels!(x, ["Young", "Middle", "Old"]) === x
     @test levels(x) == ["Young", "Middle", "Old"]
     @test isa(unique(x), typeof(x)) && isempty(unique(x))
 
+    y = copy(x)
+    @test unique!(y) === y
+    @test y == unique(x)
+
     # To test short-circuiting
     x = CategoricalArray(repeat(1:10, inner=10))
     @test levels(x) == collect(1:10)
@@ -757,6 +765,10 @@ end
     @test levels(x) == [19:-1:1; 20]
     @test unique(x) == collect(1:10)
     @test unique(x) isa typeof(x)
+
+    y = copy(x)
+    @test unique!(y) === y
+    @test y == 1:10
 end
 
 end
diff --git a/test/12_missingarray.jl b/test/12_missingarray.jl
@@ -23,6 +23,10 @@ const ≅ = isequal
                 @test size(x) === (3,)
                 @test length(x) === 3
 
+                y = copy(x)
+                @test y === unique!(y)
+                @test y == unique(x)
+
                 @test convert(CategoricalArray, x) === x
                 @test convert(CategoricalArray{Union{String, Missing}}, x) === x
                 @test convert(CategoricalArray{Union{String, Missing}, 1}, x) === x