optimize pdf broadcasting

Author: OkonSamuel

src/arrays.jl

@@ -151,10 +151,10 @@ end
 
 # dummy function
 # returns `x[i]` for `Array` inputs `x`
-# For non-Array inputs returns `zero(dtype)`
+# For non-Array inputs returns the input
 #This avoids using an if statement
-_getindex(x::Array, i, dtype)=x[i]
-_getindex(::Nothing, i, dtype) = zero(dtype)
+_getindex(x::Array, i) = x[i]
+_getindex(x, i) = x
 
 # pdf.(u, cv)
 function Base.Broadcast.broadcasted(
@@ -164,15 +164,16 @@ function Base.Broadcast.broadcasted(
 
     # we assume that we compare categorical values by their unwrapped value
     # and pick the index of this value from classes(u)
-    cv_loc = findfirst(==(cv), classes(u))
-    cv_loc == 0 && throw(err_missing_class(cv))
+    _classes = classes(u)
+    cv_loc = get(CategoricalArrays.pool(_classes), cv, zero(R))
+    isequal(cv_loc, 0) && throw(err_missing_class(cv))
 
     f() = zeros(P, size(u)) #default caller function
 
     return Base.Broadcast.Broadcasted(
         identity,
         (get(f, u.prob_given_ref, cv_loc),)
-        )
+    )
 end
 
 Base.Broadcast.broadcasted(
@@ -191,15 +192,23 @@ function Base.Broadcast.broadcasted(
     length(u) == length(v) ||throw(DimensionMismatch(
         "Arrays could not be broadcast to a common size; "*
         "got a dimension with lengths $(length(u)) and $(length(v))"))
+    
+    _classes = classes(u)
+    _classes_pool = CategoricalArrays.pool(_classes)
+    T = eltype(v) >: Missing ? Missing : Union{}
+    v_loc_flat = Vector{Union{Int, T}}(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
+    end
 
-    v_loc_flat = [(ismissing(x) ? missing : findfirst(==(x), classes(u)), i)
-                  for (i, x) in enumerate(v)]
-    any(isequal(0), v_loc_flat) && throw(err_missing_class(cv))
-
-    getter((cv_loc, i), dtype) =
-        _getindex(get(u.prob_given_ref, cv_loc, nothing), i, dtype)
-    getter(::Tuple{Missing,Any}, dtype) = missing
-    ret_flat = getter.(v_loc_flat, P)
+    getter(cv_loc) =
+        _getindex(get(u.prob_given_ref, _classes[cv_loc], zero(P)), cv_loc)
+    getter(::Missing) = missing
+    ret_flat = getter.(v_loc_flat)
     return reshape(ret_flat, size(u))
 end
 

test/arrays.jl

@@ -166,6 +166,13 @@ _skip(v) = collect(skipmissing(v))
         @test _skip(broadcast(logpdf, u, unwrap.(v))) ==
                       _skip([logpdf(u[i], v[i]) for i in 1:length(u)])
     end
+
+    ## Check that the appropriate errors are thrown
+    v1 = [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