Make broadcast work for scalar Strings

JuliaStats · Feb 16, 2017 · 937d396 · 937d396
1 parent ed30192
commit 937d396
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/src/broadcast.jl b/src/broadcast.jl
@@ -116,9 +116,11 @@ Base.map!{F}(f::F, B::Union{DataArray, PooledDataArray}, A0, As...) =
                 As[k] <: DataArray ? quote
                     $(Symbol("state_$(k)_")){d-1} = $(Symbol("state_$(k)_d"));
                     $(Symbol("j_$(k)_d")) = $(Symbol("skip_$(k)_d")) ? 1 : i_d
-                end : quote
+                end : (As[k] <: AbstractArray ? quote
                     $(Symbol("j_$(k)_d")) = size($(Symbol("A_$(k)")), d) == 1 ? 1 : i_d
-                end
+                end : quote
+                    $(Symbol("j_$(k)_d")) = 1
+                end)
             for k = 1:N]...))),
 
             # post
@@ -138,10 +140,16 @@ Base.map!{F}(f::F, B::Union{DataArray, PooledDataArray}, A0, As...) =
                     end : nothing
                 for k = 1:N]...))
 
-                # Extract values for ordinary AbstractArrays
+                # Extract values for other type
                 $(Expr(:block, [
-                    :(@inbounds $(Symbol("v_$(k)")) = @nref $nd $(Symbol("A_$(k)")) d->$(Symbol("j_$(k)_d")))
-                    for k = find(t -> !(t <: DataArray || t <: PooledDataArray), As)]...))
+                    As[k] <: AbstractArray  && !(As[k] <: AbstractDataArray) ? quote
+                        # ordinary AbstractArrays
+                        @inbounds $(Symbol("v_$(k)")) = @nref $nd $(Symbol("A_$(k)")) d->$(Symbol("j_$(k)_d"))
+                    end : quote
+                        # non AbstractArrays (e.g. Strings and Numbers)
+                        @inbounds $(Symbol("v_$(k)")) = $(Symbol("A_$(k)"))
+                    end
+                for k = 1:N]...))
 
                 # Compute and store return value
                 $(gen_na_conds(F, nd, As, B))

diff --git a/test/broadcast.jl b/test/broadcast.jl
@@ -113,6 +113,9 @@ rt = Base.return_types(broadcast, (typeof(+), Array{Float64, 3}, DataArray{Int,
 rt = Base.return_types(broadcast!, (typeof(+), DataArray{Float64, 3}, Array{Float64, 3}, Array{Int, 1}))
 @test length(rt) == 1 && rt[1] == DataArray{Float64, 3}
 
+# Test String broadcast
+@test broadcast(==, @data(["a", "b", "c", "d"]), "a") == @data([true,false,false,false])
+
 # Test broadcasting of functions that do something besides propagate NA
 @test isequal(broadcast(isequal, @data([NA, 1]), @data([NA 1])), @data([true false; false true]))
 @test isequal(broadcast(isequal, @pdata([NA, 1]), @data([NA 1])), @data([true false; false true]))