Add Missings.replace method (JuliaData#121)

Without it the function fails when the replacement isn't already a categorical value.

Author: gustafsson

src/missingarray.jl

@@ -1,4 +1,4 @@
-import Base: getindex, setindex!, similar, in
+import Base: getindex, setindex!, similar, in, collect
 
 @inline function getindex(A::CategoricalArray{T}, I...) where {T>:Missing}
     @boundscheck checkbounds(A, I...)
@@ -26,3 +26,13 @@ Base.fill!(A::CategoricalArray{>:Missing}, ::Missing) = (fill!(A.refs, 0); A)
 
 in(x::Missing, y::CategoricalArray) = false
 in(x::Missing, y::CategoricalArray{>:Missing}) = !all(v -> v > 0, y.refs)
+
+function Missings.replace(a::CategoricalArray{S, N, R, V, C}, replacement::V) where {S, N, R, V, C}
+    pool = deepcopy(a.pool)
+    v = C(get!(pool, replacement), pool)
+    Missings.replace(a, v)
+end
+
+function collect(r::Missings.EachReplaceMissing{<:CategoricalArray{S, N, R, C}}) where {S, N, R, C}
+    CategoricalArray{C,N}(R[v.level for v in r], r.replacement.pool)
+end

test/12_missingarray.jl

@@ -1066,4 +1066,67 @@ end
     @test missing in x
 end
 
+@testset "Missings.replace should work on CategoricalArrays" begin
+    x = categorical(["a", "b", missing, "a"])
+    y = ["a", "b", "", "a"]
+    r = Missings.replace(x, "")
+    @test isa(r, Missings.EachReplaceMissing)
+    a = collect(r)
+    @test isa(a, CategoricalVector{String})
+    @test y == a
+    @test levels(x) == ["a", "b"]
+    @test levels(a) == ["a", "b", ""]
+
+    r = Missings.replace(x, "b")
+    y = ["a", "b", "b", "a"]
+    @test isa(r, Missings.EachReplaceMissing)
+    a = collect(r)
+    @test isa(a, CategoricalVector{String})
+    @test y == a
+    @test levels(x) == ["a", "b"]
+    @test levels(a) == ["a", "b"]
+
+    @test_throws MethodError Missings.replace(x, 1)
+end
+
+@testset "Missings.replace should work on CategoricalArrays without missing values" begin
+    x = categorical(["a", "b", "", "a"])
+    y = ["a", "b", "", "a"]
+    r = Missings.replace(x, "dummy")
+    @test isa(r, Missings.EachReplaceMissing)
+    a = collect(r)
+    @test isa(a, CategoricalVector{String})
+    @test y == a
+    @test levels(x) == ["", "a", "b"]
+    @test levels(a) == ["", "a", "b", "dummy"]
+
+    r = Missings.replace(x, "")
+    @test isa(r, Missings.EachReplaceMissing)
+    a = collect(r)
+    @test isa(a, CategoricalVector{String})
+    @test y == a
+    @test levels(x) == ["", "a", "b"]
+    @test levels(a) == ["", "a", "b"]
+end
+
+@testset "Missings.replace should work on CategoricalArrays with empty pools" begin
+    x = categorical(Union{String,Missing}[missing])
+    y = [""]
+    r = Missings.replace(x, "")
+    @test isa(r, Missings.EachReplaceMissing)
+    a = collect(r)
+    @test isa(a, CategoricalVector{String})
+    @test y == a
+end
+
+@testset "Missings.replace should work on empty CategoricalArrays" begin
+    x = categorical(Union{String,Missing}[])
+    y = String[]
+    r = Missings.replace(x, "")
+    @test isa(r, Missings.EachReplaceMissing)
+    a = collect(r)
+    @test isa(a, CategoricalVector{String})
+    @test y == a
+end
+
 end