fix code and add more tests

Author: OkonSamuel

src/arrays.jl

@@ -196,18 +196,18 @@ function Base.Broadcast.broadcasted(
     _classes = classes(u)
     _classes_pool = CategoricalArrays.pool(_classes)
     T = eltype(v) >: Missing ? Missing : Union{}
-    v_loc_flat = Vector{Union{Int, T}}(undef, length(v))
+    v_loc_flat = Vector{Tuple{Union{R, T}, Int}}(undef, length(v))
 
 
     for (i, x) in enumerate(v)
-        v_loc_flat_i = ismissing(x) ? missing : get(_classes_pool, x, zero(R))
-        isequal(v_loc_flat_i, 0) && throw(err_missing_class(x))
-        v_loc_flat[i] = v_loc_flat_i
+        cv_ref = ismissing(x) ? missing : get(_classes_pool, x, zero(R))
+        isequal(cv_ref, 0) && throw(err_missing_class(x))
+        v_loc_flat[i] = (cv_ref, i)
     end
 
-    getter(cv_loc) =
-        _getindex(get(u.prob_given_ref, _classes[cv_loc], zero(P)), cv_loc)
-    getter(::Missing) = missing
+    getter((cv_ref, i)) =
+        _getindex(get(u.prob_given_ref, cv_ref, zero(P)), i)
+    getter(::Tuple{Missing,Any}) = missing
     ret_flat = getter.(v_loc_flat)
     return reshape(ret_flat, size(u))
 end

test/arrays.jl

@@ -149,6 +149,9 @@ end
     u2 = UnivariateFinite(v[1:2], probs, augment=true)
     @test pdf.(u2, v[3]) == zeros(3)
     @test isequal(logpdf.(u2, v[3]), log.(zeros(3)))
+
+    ## Check that the appropriate errors are thrown
+    @test_throws DomainError pdf.(u,"strange_level")
 end
 
 _skip(v) = collect(skipmissing(v))
@@ -168,23 +171,23 @@ _skip(v) = collect(skipmissing(v))
     end
 
     ## Check that the appropriate errors are thrown
-    v1 = [v0[1:end-1];"strange_level"]
+    v1 = categorical([v0[1:end-1];"strange_level"])
     v2 = [v0...;rand(rng, v0)] #length(u) !== length(v2)
     @test_throws DimensionMismatch broadcast(pdf, u, v2)
     @test_throws DomainError broadcast(pdf, u, v1)
 
 end
 
-@testset "broadcasting: check indexing in `getter((cv, i), dtype)` see PR#375" begin
+@testset "broadcasting: check indexing in `getter((cv_ref, i))` see PR#375 from MLJBase" begin
     c  = categorical([0,1,1])
     d = UnivariateFinite(c[1:1], [1 1 1]')
     v = categorical([0,1,1,1])
     @test broadcast(pdf, d, v[2:end]) == [0,0,0]
 end
 
 @testset "_getindex" begin
-   @test CategoricalDistributions._getindex(collect(1:4), 2, Int64) == 2
-   @test CategoricalDistributions._getindex(nothing, 2, Int64) == zero(Int64)
+   @test CategoricalDistributions._getindex(collect(1:4), 2) == 2
+   @test CategoricalDistributions._getindex(0, 2) === 0
 end
 
 @testset "broadcasting mode" begin