Add spreadmissings

.gitignore

@@ -1,3 +1,4 @@
 *.jl.cov
 *.jl.*.cov
 *.jl.mem
+Manifest.toml

src/Missings.jl

@@ -2,7 +2,7 @@ module Missings
 
 export allowmissing, disallowmissing, ismissing, missing, missings,
        Missing, MissingException, levels, coalesce, passmissing, nonmissingtype,
-       skipmissings
+       skipmissings, spreadmissings
 
 using Base: ismissing, missing, Missing, MissingException
 using Base: @deprecate
@@ -208,6 +208,290 @@ missing
 """
 passmissing(f) = PassMissing{Core.Typeof(f)}(f)
 
+struct SpreadMissings{F} <: Function
+    f::F
+    spread::Symbol
+    function SpreadMissings(f, spread)
+        if !(spread in (:default, :nonmissing, :none))
+            throw(ArgumentError("spread must be either :default, :nonmissing, or :none"))
+        end
+        new{Core.Typeof(f)}(f, spread)
+    end
+end
+
+function non_spreadabble_check(t::Union{AbstractDict, NamedTuple, Tuple})
+    T = typeof(t)
+    s = "Spreadmissings on $T is reserved. Please wrap in Ref to be " *
+        "treated as a scalar."
+    throw(ArgumentError(s))
+end
+non_spreadabble_check(x) = nothing
+
+"""
+Given an input vector `a` where `nonmissinginds` is guaranteed
+to not include any missing values, return a `SubArray` referencing
+the `nonmissinginds`. The element type of the returned output
+does not include `missing`.
+"""
+function nomissing_subarray(a::AbstractVector, nonmissinginds::AbstractVector)
+    T = nonmissingtype(eltype(a)) # Element type
+    N = 1 # Dimension of view
+    P = typeof(a) # Type of parent array
+    I = Tuple{typeof(nonmissinginds)} # Type of the non-missing indices
+    L = Base.IndexStyle(a) == IndexLinear # If the type supports fast linear indexing (assumed true)
+    SubArray{T, N, P, I, L}(a, (nonmissinginds,), 0, 1)
+end
+
+function new_args_subarray(args::Tuple, nonmissinginds::AbstractVector)
+    newargs = ntuple(length(args)) do i
+        a = args[i]
+        if a isa AbstractVector
+            nomissing_subarray(a, nonmissinginds)
+        else
+            a
+        end
+    end
+end
+
+function maybespread_missing(f, newargs, new_kwargs, vecs, nonmissinginds, nonmissingmask)
+    spread = f.spread
+    res = f.f(newargs...; new_kwargs...)
+
+    if res isa AbstractVector
+        # Default and spread have the same behavior if
+        # output is a vector
+        if spread === :default || spread === :nonmissing
+            if length(res) != length(nonmissinginds)
+                s = "When spreading a vector result, " *
+                    "length of output must match number of jointly non-"
+                    "missing indices in inputs. Currently spread = :$(spread)."
+                throw(DimensionMismatch(s))
+            end
+            out = similar(res, Union{eltype(res), Missing}, length(vecs[1]))
+            fill!(out, missing)
+            out[nonmissingmask] .= res
+        elseif spread === :none
+            out = res
+        else
+            throw(ArgumentError("Should not reach 1"))
+        end
+    else
+        if spread === :nonmissing
+            out = Vector{Union{typeof(res), Missing}}(undef, length(vecs[1]))
+            fill!(out, missing)
+            for ind in nonmissinginds
+                out[ind] = res
+            end
+        elseif spread === :default || spread === :none
+            out = res
+        else
+            throw(ArgumentError("Should not reach 2"))
+        end
+    end
+
+    return out
+end
+
+function maybespread_nomissing(f, args, kwargs, vecs)
+    spread = f.spread
+    res = f.f(args...; kwargs...)
+
+    if res isa AbstractVector
+        # Default and spread have the same behavior if
+        # output is a vector
+        if spread === :default || spread === :nonmissing
+            if length(res) != length(first(vecs))
+                s = "When spreading a vector result, " *
+                    "length of output must match number of jointly non-"
+                    "missing indices in inputs. Currently spread = :$(spread)."
+                throw(DimensionMismatch(s))
+            end
+            out = res
+        elseif spread === :none
+            out = res
+        else
+            throw(ArgumentError("Should not reach 1"))
+        end
+    else
+        if spread === :nonmissing
+            out = Vector{typeof(res)}(undef, length(vecs[1]))
+            fill!(out, res)
+        elseif spread === :default || spread === :none
+            out = res
+        else
+            throw(ArgumentError("Should not reach 2"))
+        end
+    end
+
+    return out
+end
+
+function check_indices_match(vecs...)
+    Base.require_one_based_indexing(vecs...)
+    findex = eachindex(first(vecs))
+    # If vectors don't have the same indices, throw a
+    # nice error for the user.
+    if !(all(x -> eachindex(x) == findex, vecs[2:end]))
+        d = Dict()
+        for i in 1:length(vecs)
+            e = eachindex(vecs[i])
+            if eachindex(e) in keys(d)
+                push!(d[i], i)
+            else
+                d[e] = [i]
+            end
+        end
+        s = "The indices of vector-input arguments are not all " *
+            "the same.\n"
+
+        for k in keys(d)
+            inds = join(d[k], ", ", " and ")
+            ind_msg = "Vector inputs $inds have indices $k\n"
+            s = s * ind_msg
+        end
+
+        throw(DimensionMismatch(s))
+    end
+end
+
+
+function (f::SpreadMissings{F})(args...; kwargs...) where {F}
+    kwargs_vals = values(values(kwargs))
+    xs = tuple(args..., kwargs_vals...)
+    foreach(non_spreadabble_check, xs)
+
+    # Detect vector inputs which contain missing in
+    # either the main arguments or keyword arguments
+    if any(x -> x isa AbstractVector{>:Missing}, xs)
+        # Check that all vector inputs have the
+        # same indices. Collect these vector inputs
+        # into a single object.
+        #
+        # TODO: Allow users to protect vector inputs
+        vecs = Base.filter(x -> x isa AbstractVector, xs)
+        check_indices_match(vecs...)
+        # Determine which indices in our collection of
+        # vector inputs have no missing values in
+        # all our inputs.
+        nonmissingmask = fill(true, length(vecs[1]))
+
+        for v in vecs
+            nonmissingmask .&= .!ismissing.(v)
+        end
+        nonmissinginds = findall(nonmissingmask)
+        # Construct new versions of arguments
+        # with no Vector{Union{T, Missing}}
+        newargs = new_args_subarray(args, nonmissinginds)
+        new_kwargs_vals = new_args_subarray(kwargs_vals, nonmissinginds)
+
+        new_kwargs = NamedTuple{keys(kwargs)}(new_kwargs_vals)
+        maybespread_missing(f, newargs, new_kwargs, vecs, nonmissinginds, nonmissingmask)
+    # There is at least one vector, but none of the vectors can contain missing
+    elseif any(x -> x isa AbstractVector, xs)
+        vecs = Base.filter(x -> x isa AbstractVector, xs)
+        check_indices_match(vecs...)
+        maybespread_nomissing(f, args, kwargs, vecs)
+    else
+        f.f(args...; kwargs...)
+    end
+end
+
+"""
+    spreadmissings(f; spread = :default)
+
+Given a function `f`, function `f` but performs a transformation
+on arguments to remove missing values before executing.
+
+### Initial example
+
+```julia-repl
+julia> using Statistics;
+
+julia> xmiss = [1, 2, 3, missing];
+
+julia> ymiss = [missing, 200, 300, 400];
+
+julia> summeans(x, y) = mean(x) + mean(y);
+
+julia> spreadmissings(summeans)(xmiss, ymiss)
+252.5
+```
+
+### Details
+
+Given the call
+
+```
+spreadmissings(f)(x::AbstractVector, y::Integer, z::AbstractVector)
+```
+
+finds the indices which corresond to `missing` values in *both*
+`x` and `z`. Then apply `f` on the `SubArray`s of `x` and `z` which
+contain non-missing values. In essense:
+
+```
+inds = .!missing.(x) .& .!missing.(z)
+sx = view(x, inds); sy = view(y, inds)
+f(sx, y, sy)
+```
+
+!!! note
+    `spreadmissings` does not use the default `view` behavior. Rather,
+    it constructs a `SubArray` directly such that the eltype of the new
+    inputs do not include `Missing`.
+
+### `spread` keyword argument
+
+Control over how the output from `f` is "spread"
+along with respect to missing values.
+
+* `:default`:
+    * If `output` is a `Vector` with the same length as the number of
+      jointly non-missing elements of the inputs `output` is "spread"
+      to match the non-missing elements of the inputs.
+    * If the `output` is a `Vector` whose length is not the same
+      as the length of number of non-missing elements of the inputs,
+      a `DimensionMismatch` error is thrown.
+    * If the output is not a `Vector`, `output` is simply returned directly
+* `:nonmissing`:
+    * If `output` is a `Vector`, behavior is the same as `:default`
+    * If `output` is not a `Vector`, `output` is spread along non-missing
+    elements of the inputs.
+* `:none`: `output` is returned directly, whether a `Vector` or not.
+
+A summary of the behavior is given in the table below:
+
+| spread \\ output type  | Vector           | Non-Vector       |
+|:---------------------- |:---------------- |:---------------- |
+| :default               | spread and match | return           |
+| :nonmissing            | spread and match | spread and match |
+| :none                  | return           | return           |
+
+If there are `AbstractVector` inputs but none of these inputs
+`AbstractVector{>:Missing}`, behavior of `spread` is the same as
+with inputs which allows for missing values. However the returned
+vectors will not allow for `missing`.
+
+If none of the argumets are `AbstractVector`s, `spreadmissings(f)`
+behaves the same as `f` regardpess of `spread`.
+
+### Limitations
+
+`spreadmissings` currently does not support:
+
+* Different length vector inputs. For
+
+```
+spreadmissings(f)([1, 2], [100, 200, 300])
+```
+
+will error.
+
+* Full spreading of scalar outputs across the *full* length of the
+input vector. That is, there is no `spread = :all` option.
+"""
+spreadmissings(f; spread = :default) = SpreadMissings(f, spread)
+
 """
    skipmissings(args...)
 
@@ -258,7 +542,7 @@ struct SkipMissings{V, T}
     others::T
 end
 
-Base.@propagate_inbounds function _anymissingindex(others::Tuple{Vararg{AbstractArray}}, i)    
+Base.@propagate_inbounds function _anymissingindex(others::Tuple{Vararg{AbstractArray}}, i)
    for oth in others
         oth[i] === missing && return true
     end
@@ -267,7 +551,7 @@ Base.@propagate_inbounds function _anymissingindex(others::Tuple{Vararg{Abstract
 end
 
 @inline function _anymissingiterate(others::Tuple, state)
-    for oth in others 
+    for oth in others
         y = iterate(oth, state)
         y !== nothing && first(y) === missing && return true
     end
@@ -278,7 +562,7 @@ end
 const SkipMissingsofArrays = SkipMissings{V, T} where
     {V <: AbstractArray, T <: Tuple{Vararg{AbstractArray}}}
 
-function Base.show(io::IO, mime::MIME"text/plain", itr::SkipMissings{V}) where V 
+function Base.show(io::IO, mime::MIME"text/plain", itr::SkipMissings{V}) where V
     print(io, SkipMissings, '{', V, '}', '(', itr.x, ')', " comprised of " *
           "$(length(itr.others) + 1) iterators")
 end
@@ -335,7 +619,7 @@ end
 @inline function Base.getindex(itr::SkipMissingsofArrays, i)
     @boundscheck checkbounds(itr.x, i)
     @inbounds xi = itr.x[i]
-    if xi === missing || @inbounds _anymissingindex(itr.others, i) 
+    if xi === missing || @inbounds _anymissingindex(itr.others, i)
         throw(MissingException("the value at index $i is missing for some element"))
     end
     return xi
@@ -380,9 +664,9 @@ Base.mapreduce_impl(f, op, A::SkipMissingsofArrays, ifirst::Integer, ilast::Inte
     A = itr.x
     if ifirst == ilast
         @inbounds a1 = A[ifirst]
-        if a1 === missing 
+        if a1 === missing
             return nothing
-        elseif _anymissingindex(itr.others, ifirst) 
+        elseif _anymissingindex(itr.others, ifirst)
             return nothing
         else
             return Some(Base.mapreduce_first(f, op, a1))
@@ -436,7 +720,7 @@ end
 Return a vector similar to the array wrapped by the given `SkipMissings` iterator
 but skipping all elements with a `missing` value in one of the iterators passed
 to `skipmissing` and elements for which `f` returns `false`. This method
-only applies when all iterators passed to `skipmissings` are arrays. 
+only applies when all iterators passed to `skipmissings` are arrays.
 
 # Examples
 ```