From a4b75e75e5226406a70bebc7046825b73990678d Mon Sep 17 00:00:00 2001
From: Alexey Stukalov <astukalov@gmail.com>
Date: Tue, 27 Jul 2021 14:35:33 +0200
Subject: [PATCH] add in-place Base.unique!(::CatVec) -> CatVec

---
 src/array.jl            | 11 ++++++++++-
 test/11_array.jl        | 12 ++++++++++++
 test/12_missingarray.jl |  4 ++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/array.jl b/src/array.jl
index aa2c9003..766afb83 100644
--- 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)
@@ -826,6 +826,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
index dfc78eb4..38ba9f3a 100644
--- a/test/11_array.jl
+++ b/test/11_array.jl
@@ -706,6 +706,10 @@ 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))
@@ -713,6 +717,10 @@ end
     @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)
@@ -722,6 +730,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
index 854a8e02..893a5a4a 100644
--- 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